Mitigates a critical Out-of-Bounds (OOB) read vulnerability by transitioning internal trie references to absl::Span and validating arra...

PiperOrigin-RevId: 914999991

Author: tf-text-github-robot

tensorflow_text/core/kernels/BUILD

@@ -215,6 +215,7 @@ 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,7 +352,10 @@ 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",
     ],
 )
 
@@ -363,6 +367,7 @@ 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",
     ],
 )
@@ -399,6 +404,7 @@ 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,6 +14,7 @@
 
 #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"
 
@@ -31,7 +32,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.data()));
+                       DartsCloneTrieWrapper::Create(trie_array));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
   int data;
@@ -56,7 +57,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.data()));
+                       DartsCloneTrieWrapper::Create(trie_array));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
   int data;
@@ -76,7 +77,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.data()));
+                       DartsCloneTrieWrapper::Create(trie_array));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
 
@@ -94,7 +95,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.data()));
+                       DartsCloneTrieWrapper::Create(trie_array));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
   int data;
@@ -115,7 +116,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.data()));
+                       DartsCloneTrieWrapper::Create(trie_array));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
   int data;
@@ -135,7 +136,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.data()));
+                       DartsCloneTrieWrapper::Create(trie_array));
 
   DartsCloneTrieWrapper::TraversalCursor cursor;
 
@@ -183,6 +184,20 @@ 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,7 +30,9 @@
 #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 {
@@ -51,14 +53,14 @@ class DartsCloneTrieWrapper {
     uint32_t unit = 0;
   };
 
-  // Constructs an instance by passing in the pointer to the trie array data.
+  // Constructs an instance by passing in the span of 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(
-      const uint32_t* trie_array) {
-    if (trie_array == nullptr) {
-      return absl::InvalidArgumentError("trie_array is nullptr.");
+      absl::Span<const uint32_t> trie_array) {
+    if (trie_array.empty() || trie_array.data() == nullptr) {
+      return absl::InvalidArgumentError("trie_array is empty or nullptr.");
     }
     return DartsCloneTrieWrapper(trie_array);
   }
@@ -70,20 +72,28 @@ 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) {
-    cursor.node_id = node_id;
-    cursor.unit = trie_array_[node_id];
+    if (node_id < trie_array_.size()) {
+      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;
@@ -108,15 +118,18 @@ class DartsCloneTrieWrapper {
     if (!has_leaf(cursor.unit)) {
       return false;
     }
-    const uint32_t value_unit =
-        trie_array_[cursor.node_id ^ offset(cursor.unit)];
+    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];
     out_data = value(value_unit);
     return true;
   }
 
  private:
   // Use Create() instead of the constructor.
-  explicit DartsCloneTrieWrapper(const uint32_t* trie_array)
+  explicit DartsCloneTrieWrapper(absl::Span<const uint32_t> trie_array)
       : trie_array_(trie_array) {}
 
   // The actual implementation of TryTraverseSeveralSteps.
@@ -127,6 +140,9 @@ 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;
@@ -157,8 +173,8 @@ class DartsCloneTrieWrapper {
     return static_cast<int>(unit & 0x7fffffff);
   }
 
-  // The pointer to the darts trie array.
-  const uint32_t* trie_array_;
+  // The dart trie array represented as a span for bounds awareness.
+  absl::Span<const uint32_t> trie_array_;
 };
 
 }  // namespace trie_utils

tensorflow_text/core/kernels/fast_bert_normalizer.h

@@ -22,6 +22,7 @@
 #include "absl/base/optimization.h"
 #include "absl/status/status.h"
 #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"

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(), data_for_codepoint_zero_, mapped_value_pool_.data());
+      trie_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

@@ -24,6 +24,7 @@
 #include "absl/strings/str_cat.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"
@@ -56,7 +57,8 @@ FastWordpieceTokenizer::Create(const void* config_flatbuffer) {
         "FastWordpieceTokenizerConfig or its trie_array is null.");
   }
   auto trie_or = trie_utils::DartsCloneTrieWrapper::Create(
-      tokenizer.config_->trie_array()->data());
+      absl::MakeSpan(tokenizer.config_->trie_array()->data(),
+                     tokenizer.config_->trie_array()->size()));
   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_.data()));
+      trie_utils::DartsCloneTrieWrapper::Create(trie_array_));
   trie_.emplace(std::move(trie));
 
   if (trie_array_.size() >