[ML] Fail gracefully when restoring a categorizer with an out-of-range token ID

docs/changelog/3143.yaml

@@ -0,0 +1,5 @@
+area: Machine Learning
+issues: []
+pr: 3143
+summary: Fail gracefully when restoring a categorizer with an out-of-range token ID
+type: bug

lib/model/CTokenListDataCategorizerBase.cc

@@ -378,6 +378,45 @@ bool CTokenListDataCategorizerBase::acceptRestoreTraverser(core::CStateRestoreTr
         }
     } while (traverser.next());
 
+    // Validate the restored categories against an inconsistent or truncated
+    // state document. Such a document can leave a category referencing data that
+    // is later used to index a container unchecked, which is an out-of-bounds
+    // access resulting in a crash. Fail the restore gracefully here instead of
+    // proceeding with an inconsistent state.
+    const std::size_t numTokens{m_TokenIdLookup.size()};
+    const auto tokenIdsInRange = [numTokens](const TSizeSizePrVec& tokenIds) {
+        return std::all_of(tokenIds.begin(), tokenIds.end(),
+                           [numTokens](const TSizeSizePr& tokenId) {
+                               return tokenId.first < numTokens;
+                           });
+    };
+    for (const auto& category : m_Categories) {
+        // Every token ID referenced by a category must exist in the restored
+        // token ID lookup, as these IDs index m_TokenIdLookup unchecked (for
+        // example when building a reverse search).
+        if (tokenIdsInRange(category.baseTokenIds()) == false ||
+            tokenIdsInRange(category.commonUniqueTokenIds()) == false) {
+            LOG_ERROR(<< "Cannot restore categorizer - a category references a token ID "
+                         "outside the restored token ID lookup (size "
+                      << numTokens << "); the state document is inconsistent");
+            return false;
+        }
+
+        // The ordered common token bounds must describe a valid sub-range of the
+        // category's base token IDs, as they are used to index m_BaseTokenIds
+        // unchecked (for example when matching or updating a category).
+        const TSizeSizePr orderedBounds{category.orderedCommonTokenBounds()};
+        if (orderedBounds.first > orderedBounds.second ||
+            orderedBounds.second > category.baseTokenIds().size()) {
+            LOG_ERROR(<< "Cannot restore categorizer - a category has ordered common token "
+                         "bounds ["
+                      << orderedBounds.first << ", " << orderedBounds.second << ") outside its base token list (size "
+                      << category.baseTokenIds().size()
+                      << "); the state document is inconsistent");
+            return false;
+        }
+    }
+
     // Categories are persisted in order of creation, but this list needs to be
     // sorted by descending count instead
     std::stable_sort(m_CategoriesByCount.begin(), m_CategoriesByCount.end(),

lib/model/unittest/CTokenListDataCategorizerTest.cc

@@ -541,6 +541,88 @@ BOOST_FIXTURE_TEST_CASE(testPersist, CTestFixture) {
     checkMemoryUsageInstrumentation(restoredCategorizer);
 }
 
+BOOST_FIXTURE_TEST_CASE(testRestoreWithInconsistentTokenIdFailsGracefully, CTestFixture) {
+    // A corrupt or truncated state document can leave a category referencing a
+    // token ID beyond the end of the restored token ID lookup. Such an ID was
+    // previously used to index the lookup unchecked (for example when building a
+    // reverse search), causing an out-of-bounds access and a crash. The restore
+    // must now fail gracefully instead.
+    // See https://github.com/elastic/ml-cpp/issues/2875
+    //
+    // This state has a single token (so index 0 is the only valid token ID) but
+    // a category that references token ID 5 in both its base and common unique
+    // token lists.
+    const std::string inconsistentState{
+        R"({"topLevel":{)"
+        R"("a":"foo","b":1,)"
+        R"("c":{"a":"foo bar","b":5,"c":1,"d":7,"j":0,"e":1,"f":5,"g":1,"h":1,"i":2,"k":7},)"
+        R"("d":0}})"};
+
+    TTokenListDataCategorizerKeepsFields categorizer{
+        m_Limits, NO_REVERSE_SEARCH_CREATOR, 0.7, "whatever"};
+
+    std::istringstream stateStrm{inconsistentState};
+    ml::core::CJsonStateRestoreTraverser traverser{stateStrm};
+    const bool restored{traverser.traverseSubLevel(
+        [&categorizer](ml::core::CStateRestoreTraverser& traverser_) {
+            return categorizer.acceptRestoreTraverser(traverser_);
+        })};
+
+    BOOST_REQUIRE_EQUAL(false, restored);
+}
+
+BOOST_FIXTURE_TEST_CASE(testRestoreWithInvalidOrderedTokenBoundsFailsGracefully, CTestFixture) {
+    // A corrupt or truncated state document can leave a category whose ordered
+    // common token bounds fall outside its base token list. Those bounds are
+    // used to index the base token list unchecked (for example when matching or
+    // updating a category), so an out-of-range end index is an out-of-bounds
+    // access. The restore must fail gracefully instead.
+    // See https://github.com/elastic/ml-cpp/issues/2875
+    //
+    // The category has a single, in-range base token ID (0) but an ordered
+    // common token end index of 5, well beyond the single base token.
+    const std::string inconsistentState{
+        R"({"topLevel":{)"
+        R"("a":"foo","b":1,)"
+        R"("c":{"a":"foo","b":0,"c":1,"d":3,"j":0,"e":5,"f":0,"g":1,"h":1,"i":2,"k":3},)"
+        R"("d":0}})"};
+
+    TTokenListDataCategorizerKeepsFields categorizer{
+        m_Limits, NO_REVERSE_SEARCH_CREATOR, 0.7, "whatever"};
+
+    std::istringstream stateStrm{inconsistentState};
+    ml::core::CJsonStateRestoreTraverser traverser{stateStrm};
+    const bool restored{traverser.traverseSubLevel(
+        [&categorizer](ml::core::CStateRestoreTraverser& traverser_) {
+            return categorizer.acceptRestoreTraverser(traverser_);
+        })};
+
+    BOOST_REQUIRE_EQUAL(false, restored);
+}
+
+BOOST_FIXTURE_TEST_CASE(testRestoreWithConsistentTokenIdSucceeds, CTestFixture) {
+    // Sanity check that the range validation added for issue #2875 does not
+    // reject a valid, consistent state document. The category references the
+    // only valid token ID (0) and has in-range ordered common token bounds.
+    const std::string consistentState{
+        R"({"topLevel":{)"
+        R"("a":"foo","b":1,)"
+        R"("c":{"a":"foo","b":0,"c":1,"d":3,"j":0,"e":1,"f":0,"g":1,"h":1,"i":2,"k":3},)"
+        R"("d":0}})"};
+
+    TTokenListDataCategorizerKeepsFields categorizer{
+        m_Limits, NO_REVERSE_SEARCH_CREATOR, 0.7, "whatever"};
+
+    std::istringstream stateStrm{consistentState};
+    ml::core::CJsonStateRestoreTraverser traverser{stateStrm};
+    const bool restored{traverser.traverseSubLevel(
+        [&categorizer](ml::core::CStateRestoreTraverser& traverser_) {
+            return categorizer.acceptRestoreTraverser(traverser_);
+        })};
+
+    BOOST_REQUIRE_EQUAL(true, restored);
+}
+
 BOOST_FIXTURE_TEST_CASE(testLongReverseSearch, CTestFixture) {
     TTokenListDataCategorizerKeepsFields::TTokenListReverseSearchCreatorCPtr reverseSearchCreator{
         new ml::model::CTokenListReverseSearchCreator{"_raw"}};