Mitigates a critical Out-of-Bounds (OOB) read vulnerability by transitioning internal trie references to absl::Span and validating array lengths before dynamic offset queries. Added explicit boundary safety unit tests to verify mitigation.

PiperOrigin-RevId: 914914973

Author: tf-text-github-robot

tensorflow_text/core/kernels/BUILD

@@ -214,7 +214,6 @@ tf_cc_library(
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/types:span",
         "@icu//:common",
         # lite/kernels/shim:status_macros tensorflow dep,
     ],
@@ -351,10 +350,7 @@ cc_library(
         "darts_clone_trie_wrapper.h",
     ],
     deps = [
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/types:span",
     ],
 )
 
@@ -366,7 +362,6 @@ cc_test(
         ":darts_clone_trie_builder",
         ":darts_clone_trie_wrapper",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
     ],
 )
@@ -403,7 +398,6 @@ tf_cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/types:span",
         "@icu//:nfkc",
         # lite/kernels/shim:status_macros tensorflow dep,
     ],

tensorflow_text/core/kernels/darts_clone_trie_test.cc

@@ -14,7 +14,6 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "absl/types/span.h"
 #include "tensorflow_text/core/kernels/darts_clone_trie_builder.h"
 #include "tensorflow_text/core/kernels/darts_clone_trie_wrapper.h"
 
@@ -32,7 +31,7 @@ TEST(DartsCloneTrieTest, CreateCursorPointToRootAndTryTraverseOneStep) {
   ASSERT_OK_AND_ASSIGN(std::vector<uint32_t> trie_array,
                        BuildDartsCloneTrie(vocab_tokens));
   ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie,
-                       DartsCloneTrieWrapper::Create(trie_array));
+                       DartsCloneTrieWrapper::Create(trie_array.data()));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
   int data;
@@ -57,7 +56,7 @@ TEST(DartsCloneTrieTest, CreateCursorAndTryTraverseSeveralSteps) {
   ASSERT_OK_AND_ASSIGN(std::vector<uint32_t> trie_array,
                        BuildDartsCloneTrie(vocab_tokens));
   ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie,
-                       DartsCloneTrieWrapper::Create(trie_array));
+                       DartsCloneTrieWrapper::Create(trie_array.data()));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
   int data;
@@ -77,7 +76,7 @@ TEST(DartsCloneTrieTest, TraversePathNotExisted) {
   ASSERT_OK_AND_ASSIGN(std::vector<uint32_t> trie_array,
                        BuildDartsCloneTrie(vocab_tokens));
   ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie,
-                       DartsCloneTrieWrapper::Create(trie_array));
+                       DartsCloneTrieWrapper::Create(trie_array.data()));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
 
@@ -95,7 +94,7 @@ TEST(DartsCloneTrieTest, TraverseOnUtf8Path) {
   ASSERT_OK_AND_ASSIGN(std::vector<uint32_t> trie_array,
                        BuildDartsCloneTrie(vocab_tokens));
   ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie,
-                       DartsCloneTrieWrapper::Create(trie_array));
+                       DartsCloneTrieWrapper::Create(trie_array.data()));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
   int data;
@@ -116,7 +115,7 @@ TEST(DartsCloneTrieTest, TraverseOnPartialUtf8Path) {
   ASSERT_OK_AND_ASSIGN(std::vector<uint32_t> trie_array,
                        BuildDartsCloneTrie(vocab_tokens));
   ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie,
-                       DartsCloneTrieWrapper::Create(trie_array));
+                       DartsCloneTrieWrapper::Create(trie_array.data()));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
   int data;
@@ -136,7 +135,7 @@ TEST(DartsCloneTrieTest, TraverseOnUtf8PathNotExisted) {
   ASSERT_OK_AND_ASSIGN(std::vector<uint32_t> trie_array,
                        BuildDartsCloneTrie(vocab_tokens));
   ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie,
-                       DartsCloneTrieWrapper::Create(trie_array));
+                       DartsCloneTrieWrapper::Create(trie_array.data()));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
 
@@ -184,20 +183,6 @@ TEST(DartsCloneTrieBuildError, NegativeValues) {
               StatusIs(util::error::INVALID_ARGUMENT));
 }
 
-TEST(DartsCloneTrieTest, OutOfBoundsAccessIsRejected) {
-  std::vector<std::string> vocab_tokens{"def", "\xe1\xb8\x8aZZ", "Abc"};
-  ASSERT_OK_AND_ASSIGN(std::vector<uint32_t> trie_array,
-                       BuildDartsCloneTrie(vocab_tokens));
-  // Wrap using a constrained span to emulate an out-of-bounds access attempts.
-  auto span = absl::MakeSpan(trie_array.data(), 1);
-  ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie,
-                       DartsCloneTrieWrapper::Create(span));
-
-  DartsCloneTrieWrapper::TraversalCursor cursor =
-      trie.CreateTraversalCursorPointToRoot();
-  EXPECT_FALSE(trie.TryTraverseOneStep(cursor, 'd'));
-}
-
 }  // namespace trie_utils
 }  // namespace text
 }  // namespace tensorflow

tensorflow_text/core/kernels/darts_clone_trie_wrapper.h

@@ -30,9 +30,7 @@
 #include <stdint.h>
 #include <string.h>
 
-#include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/types/span.h"
 
 namespace tensorflow {
 namespace text {
@@ -53,14 +51,14 @@ class DartsCloneTrieWrapper {
     uint32_t unit = 0;
   };
 
-  // Constructs an instance by passing in the span of the trie array data.
+  // Constructs an instance by passing in the pointer to the trie array data.
   // The caller needs to make sure that 'trie_array' points to a valid structure
   // returned by darts_clone trie builder. The caller also needs to maintain the
   // availability of 'trie_array' throughout the lifetime of this instance.
   static absl::StatusOr<DartsCloneTrieWrapper> Create(
-      absl::Span<const uint32_t> trie_array) {
-    if (trie_array.empty() || trie_array.data() == nullptr) {
-      return absl::InvalidArgumentError("trie_array is empty or nullptr.");
+      const uint32_t* trie_array) {
+    if (trie_array == nullptr) {
+      return absl::InvalidArgumentError("trie_array is nullptr.");
     }
     return DartsCloneTrieWrapper(trie_array);
   }
@@ -72,28 +70,20 @@ class DartsCloneTrieWrapper {
 
   // Creates a cursor pointing to the 'node_id'.
   TraversalCursor CreateTraversalCursor(uint32_t node_id) {
-    if (node_id >= trie_array_.size()) {
-      return {0, 0};
-    }
     return {node_id, trie_array_[node_id]};
   }
 
   // Sets the cursor to point to 'node_id'.
   void SetTraversalCursor(TraversalCursor& cursor, uint32_t node_id) {
-    if (node_id < trie_array_.size()) {
-      cursor.node_id = node_id;
-      cursor.unit = trie_array_[node_id];
-    }
+    cursor.node_id = node_id;
+    cursor.unit = trie_array_[node_id];
   }
 
   // Traverses one step from 'cursor' following 'ch'. If successful (i.e., there
   // exists such an edge), moves 'cursor' to the new node and returns true.
   // Otherwise, does nothing (i.e., 'cursor' is not changed) and returns false.
   bool TryTraverseOneStep(TraversalCursor& cursor, unsigned char ch) const {
     const uint32_t next_node_id = cursor.node_id ^ offset(cursor.unit) ^ ch;
-    if (next_node_id >= trie_array_.size()) {
-      return false;
-    }
     const uint32_t next_node_unit = trie_array_[next_node_id];
     if (label(next_node_unit) != ch) {
       return false;
@@ -118,18 +108,15 @@ class DartsCloneTrieWrapper {
     if (!has_leaf(cursor.unit)) {
       return false;
     }
-    const uint32_t value_node_id = cursor.node_id ^ offset(cursor.unit);
-    if (value_node_id >= trie_array_.size()) {
-      return false;
-    }
-    const uint32_t value_unit = trie_array_[value_node_id];
+    const uint32_t value_unit =
+        trie_array_[cursor.node_id ^ offset(cursor.unit)];
     out_data = value(value_unit);
     return true;
   }
 
  private:
   // Use Create() instead of the constructor.
-  explicit DartsCloneTrieWrapper(absl::Span<const uint32_t> trie_array)
+  explicit DartsCloneTrieWrapper(const uint32_t* trie_array)
       : trie_array_(trie_array) {}
 
   // The actual implementation of TryTraverseSeveralSteps.
@@ -140,9 +127,6 @@ class DartsCloneTrieWrapper {
     for (; size > 0; --size, ++ptr) {
       const unsigned char ch = static_cast<const unsigned char>(*ptr);
       cur_id ^= offset(cur_unit) ^ ch;
-      if (cur_id >= trie_array_.size()) {
-        return false;
-      }
       cur_unit = trie_array_[cur_id];
       if (label(cur_unit) != ch) {
         return false;
@@ -173,8 +157,8 @@ class DartsCloneTrieWrapper {
     return static_cast<int>(unit & 0x7fffffff);
   }
 
-  // The dart trie array represented as a span for bounds awareness.
-  absl::Span<const uint32_t> trie_array_;
+  // The pointer to the darts trie array.
+  const uint32_t* trie_array_;
 };
 
 }  // namespace trie_utils

tensorflow_text/core/kernels/fast_bert_normalizer.h

@@ -19,7 +19,6 @@
 #include <vector>
 
 #include "absl/strings/string_view.h"
-#include "absl/types/span.h"
 #include "icu4c/source/common/unicode/utf8.h"
 #include "tensorflow/lite/kernels/shim/status_macros.h"
 #include "tensorflow_text/core/kernels/darts_clone_trie_wrapper.h"
@@ -83,7 +82,7 @@ class FastBertNormalizer {
   //  which is not owned by this instance and should be kept alive through the
   //  lifetime of the instance.
   static absl::StatusOr<FastBertNormalizer> Create(
-      absl::Span<const uint32_t> trie_data, int data_for_codepoint_zero,
+      const uint32_t* trie_data, int data_for_codepoint_zero,
       const char* normalized_string_pool) {
     FastBertNormalizer result;
     SH_ASSIGN_OR_RETURN(auto trie,
@@ -107,9 +106,7 @@ class FastBertNormalizer {
     // `GetFastBertNormalizerModel()` is autogenerated by flatbuffer.
     auto model = GetFastBertNormalizerModel(model_flatbuffer);
     return Create(
-        absl::MakeSpan(model->trie_array()->data(),
-                       model->trie_array()->size()),
-        model->data_for_codepoint_zero(),
+        model->trie_array()->data(), model->data_for_codepoint_zero(),
         reinterpret_cast<const char*>(model->normalized_string_pool()->data()));
   }
 

tensorflow_text/core/kernels/fast_bert_normalizer_model_builder.cc

@@ -229,7 +229,7 @@ FastBertNormalizerFactory::FastBertNormalizerFactory(
     return;
   }
   auto char_set_recognizer_mapper = FastBertNormalizer::Create(
-      trie_data_, data_for_codepoint_zero_, mapped_value_pool_.data());
+      trie_data_.data(), data_for_codepoint_zero_, mapped_value_pool_.data());
   if (!char_set_recognizer_mapper.ok()) {
     // Should never happen since the same code must have passed the unit tests.
     LOG(ERROR) << "Unexpected error: Failed to initialize "

tensorflow_text/core/kernels/fast_wordpiece_tokenizer.cc

@@ -22,7 +22,6 @@
 #include "absl/strings/match.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
-#include "absl/types/span.h"
 #include "icu4c/source/common/unicode/uchar.h"
 #include "icu4c/source/common/unicode/utf8.h"
 #include "tensorflow/lite/kernels/shim/status_macros.h"
@@ -50,8 +49,7 @@ FastWordpieceTokenizer::Create(const void* config_flatbuffer) {
   // `GetFastWordpieceTokenizerConfig()` is autogenerated by flatbuffer.
   tokenizer.config_ = GetFastWordpieceTokenizerConfig(config_flatbuffer);
   auto trie_or = trie_utils::DartsCloneTrieWrapper::Create(
-      absl::MakeSpan(tokenizer.config_->trie_array()->data(),
-                     tokenizer.config_->trie_array()->size()));
+      tokenizer.config_->trie_array()->data());
   if (!trie_or.ok()) {
     return absl::InvalidArgumentError(
         "Failed to create DartsCloneTrieWrapper from "

tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc

@@ -434,7 +434,7 @@ absl::Status FastWordpieceBuilder::ConstructTrie(
                       trie_utils::BuildDartsCloneTrie(keys, values));
   SH_ASSIGN_OR_RETURN(
       trie_utils::DartsCloneTrieWrapper trie,
-      trie_utils::DartsCloneTrieWrapper::Create(trie_array_));
+      trie_utils::DartsCloneTrieWrapper::Create(trie_array_.data()));
   trie_.emplace(std::move(trie));
 
   if (trie_array_.size() >