From 65fbcfa6874349629e8161904cc0d1a91427dd9c Mon Sep 17 00:00:00 2001 From: "TF.Text Team" Date: Wed, 9 Sep 2026 06:26:52 -0700 Subject: [PATCH] fix an integer-driven out-of-bounds memory read vulnerability in DartsCloneTrieWrapper Refactored the API surface across DartsCloneTrieWrapper, FastBertNormalizer, and OnePassNormalizer to avoid passing explicit integer size arguments (size_t/int), accepting absl::Span instead so bounds are automatically inferred from containers (std::vector, std::array, absl::FixedArray, and FlatBuffers flatbuffers::Vector). PiperOrigin-RevId: 978509649 --- tensorflow_text/core/kernels/BUILD | 9 +- .../core/kernels/darts_clone_trie_test.cc | 139 +++++++++++++++++- .../core/kernels/darts_clone_trie_wrapper.h | 94 ++++++++---- .../core/kernels/fast_bert_normalizer.h | 24 +-- .../fast_bert_normalizer_model_builder.cc | 3 +- .../core/kernels/fast_bert_normalizer_test.cc | 28 ++++ .../core/kernels/fast_wordpiece_tokenizer.cc | 2 +- .../fast_wordpiece_tokenizer_model_builder.cc | 48 +++--- .../kernels/fast_wordpiece_tokenizer_test.cc | 4 + 9 files changed, 274 insertions(+), 77 deletions(-) diff --git a/tensorflow_text/core/kernels/BUILD b/tensorflow_text/core/kernels/BUILD index f667b6174..5dd9d7b76 100644 --- a/tensorflow_text/core/kernels/BUILD +++ b/tensorflow_text/core/kernels/BUILD @@ -208,13 +208,12 @@ tf_cc_library( name = "fast_bert_normalizer", hdrs = ["fast_bert_normalizer.h"], deps = [ - ":darts_clone_trie_builder", ":darts_clone_trie_wrapper", ":fast_bert_normalizer_model", "@com_google_absl//absl/base:core_headers", - "@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, ], @@ -256,6 +255,7 @@ cc_test( deps = [ ":fast_bert_normalizer", ":fast_bert_normalizer_model_builder", + "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", ], ) @@ -351,7 +351,10 @@ cc_library( "darts_clone_trie_wrapper.h", ], deps = [ + "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", ], ) @@ -362,7 +365,7 @@ cc_test( deps = [ ":darts_clone_trie_builder", ":darts_clone_trie_wrapper", - "@com_google_absl//absl/status", + "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", ], ) diff --git a/tensorflow_text/core/kernels/darts_clone_trie_test.cc b/tensorflow_text/core/kernels/darts_clone_trie_test.cc index a80c28353..c54d6009c 100644 --- a/tensorflow_text/core/kernels/darts_clone_trie_test.cc +++ b/tensorflow_text/core/kernels/darts_clone_trie_test.cc @@ -14,6 +14,7 @@ #include #include +#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 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 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 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 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 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 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,132 @@ TEST(DartsCloneTrieBuildError, NegativeValues) { StatusIs(util::error::INVALID_ARGUMENT)); } +TEST(DartsCloneTrieTest, OutOfBoundsTraverseOneStepRejected) { + // A malicious 1-element trie formatted to yield a large internal offset. + // 0x4E2000 is 5120000; right-shifted by 10 yields offset 5000. + // When traversing with 'a' (97), next_node_id = 0 ^ 5000 ^ 97 = 4905, + // which is far beyond the 1-element vector. + std::vector malicious_trie = {0x4E2000}; + ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie, + DartsCloneTrieWrapper::Create(malicious_trie)); + + EXPECT_EQ(trie.size(), 1); + auto cursor = trie.CreateTraversalCursorPointToRoot(); + EXPECT_FALSE(trie.TryTraverseOneStep(cursor, 'a')); + // Cursor should not have changed. + EXPECT_EQ(cursor.node_id, DartsCloneTrieWrapper::kRootNodeId); +} + +TEST(DartsCloneTrieTest, OutOfBoundsTraverseSeveralStepsRejected) { + std::vector malicious_trie = {0x4E2000}; + ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie, + DartsCloneTrieWrapper::Create(malicious_trie)); + + auto cursor = trie.CreateTraversalCursorPointToRoot(); + EXPECT_FALSE(trie.TryTraverseSeveralSteps(cursor, "abc")); + EXPECT_EQ(cursor.node_id, DartsCloneTrieWrapper::kRootNodeId); +} + +TEST(DartsCloneTrieTest, OutOfBoundsGetDataRejected) { + // 0x4E2100: offset 5000 and has_leaf bit 0x100 set. + std::vector malicious_trie = {0x4E2100}; + ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie, + DartsCloneTrieWrapper::Create(malicious_trie)); + + auto cursor = trie.CreateTraversalCursorPointToRoot(); + int data = 0; + EXPECT_FALSE(trie.TryGetData(cursor, data)); +} + +TEST(DartsCloneTrieTest, CreateWithInvalidSizeOrNullFails) { + std::vector empty_trie; + EXPECT_THAT(DartsCloneTrieWrapper::Create(empty_trie), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT( + DartsCloneTrieWrapper::Create(absl::Span(nullptr, 10)), + StatusIs(absl::StatusCode::kInvalidArgument)); + uint32_t dummy = 0; + EXPECT_THAT( + DartsCloneTrieWrapper::Create(absl::Span(&dummy, 0)), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(DartsCloneTrieWrapper::Create(absl::Span()), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + +TEST(DartsCloneTrieTest, CursorOutOfBoundsSafe) { + std::vector malicious_trie = {0x4E2000}; + ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie, + DartsCloneTrieWrapper::Create(malicious_trie)); + + auto cursor = trie.CreateTraversalCursor(100); + EXPECT_EQ(cursor.node_id, 100); + EXPECT_EQ(cursor.unit, 0); + + trie.SetTraversalCursor(cursor, 200); + EXPECT_EQ(cursor.node_id, 200); + EXPECT_EQ(cursor.unit, 0); + + EXPECT_FALSE(trie.TryTraverseOneStep(cursor, 'a')); +} + +TEST(DartsCloneTrieTest, InvalidCursorCannotTraverseEvenIfNextIdLandsInBounds) { + // A 2-element trie where element 0 has label 100. + // If cursor has node_id = 100 (out of bounds for size 2) and unit = 0: + // 100 ^ offset(0) ^ 100 = 0 (which is in bounds < 2, and label matches 100). + // Traversal must be rejected because the source cursor is out of bounds. + std::vector trie_data = {100, 0}; + ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie, + DartsCloneTrieWrapper::Create(trie_data)); + + auto cursor = trie.CreateTraversalCursor(100); + EXPECT_EQ(cursor.node_id, 100); + EXPECT_FALSE(trie.TryTraverseOneStep(cursor, 100)); + // Cursor must not have been modified. + EXPECT_EQ(cursor.node_id, 100); + + // Several steps with empty path on invalid cursor must also be rejected. + EXPECT_FALSE(trie.TryTraverseSeveralSteps(cursor, "")); + EXPECT_EQ(cursor.node_id, 100); + + // Several steps with non-empty path on invalid cursor must be rejected. + EXPECT_FALSE(trie.TryTraverseSeveralSteps(cursor, "d")); + EXPECT_EQ(cursor.node_id, 100); +} + +TEST(DartsCloneTrieTest, InvalidCursorGetDataRejectedEvenIfLeafOffsetInBounds) { + // A 2-element trie. Element 0 is a leaf with value 42 (0x80000000 | 42). + std::vector trie_data = {0x8000002A, 0}; + ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie, + DartsCloneTrieWrapper::Create(trie_data)); + + // Craft a cursor with out-of-bounds node_id = 100, but has_leaf bit (0x100) + // and offset = 100 ((100 >> 0) << 10 = 100 << 10 = 0x19000). + // value_node_id = 100 ^ offset(unit) = 100 ^ 100 = 0 (< 2). + DartsCloneTrieWrapper::TraversalCursor cursor; + cursor.node_id = 100; + cursor.unit = 0x100 | (100 << 10); + int data = 0; + EXPECT_FALSE(trie.TryGetData(cursor, data)); +} + +TEST(DartsCloneTrieTest, + SeveralStepsPartialMatchFailureLeavesCursorUnmodified) { + std::vector vocab_tokens{"abc", "def"}; + ASSERT_OK_AND_ASSIGN(std::vector trie_array, + BuildDartsCloneTrie(vocab_tokens)); + ASSERT_OK_AND_ASSIGN(DartsCloneTrieWrapper trie, + DartsCloneTrieWrapper::Create(trie_array)); + + auto cursor = trie.CreateTraversalCursorPointToRoot(); + EXPECT_TRUE(trie.TryTraverseOneStep(cursor, 'a')); + const uint32_t original_node_id = cursor.node_id; + + // "bz" matches 'b' but fails at 'z'. + EXPECT_FALSE(trie.TryTraverseSeveralSteps(cursor, "bz")); + // Cursor should still point to 'a', not moved to 'b'. + EXPECT_EQ(cursor.node_id, original_node_id); +} + } // namespace trie_utils } // namespace text } // namespace tensorflow diff --git a/tensorflow_text/core/kernels/darts_clone_trie_wrapper.h b/tensorflow_text/core/kernels/darts_clone_trie_wrapper.h index 43067ec1b..5249dd182 100644 --- a/tensorflow_text/core/kernels/darts_clone_trie_wrapper.h +++ b/tensorflow_text/core/kernels/darts_clone_trie_wrapper.h @@ -27,10 +27,14 @@ #ifndef THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_DARTS_CLONE_TRIE_WRAPPER_H_ #define THIRD_PARTY_TENSORFLOW_TEXT_CORE_KERNELS_DARTS_CLONE_TRIE_WRAPPER_H_ +#include #include #include +#include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" namespace tensorflow { namespace text { @@ -51,31 +55,41 @@ class DartsCloneTrieWrapper { uint32_t unit = 0; }; - // Constructs an instance by passing in the pointer to the trie array data. + // Constructs an instance from an absl::Span. // 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 Create( - const uint32_t* trie_array) { - if (trie_array == nullptr) { + absl::Span trie_array) { + if (trie_array.empty()) { + return absl::InvalidArgumentError("trie_array must not be empty."); + } + if (trie_array.data() == nullptr) { return absl::InvalidArgumentError("trie_array is nullptr."); } - return DartsCloneTrieWrapper(trie_array); + return DartsCloneTrieWrapper(trie_array.data(), trie_array.size()); } // Creates a cursor pointing to the root. - TraversalCursor CreateTraversalCursorPointToRoot() { - return {kRootNodeId, trie_array_[kRootNodeId]}; + TraversalCursor CreateTraversalCursorPointToRoot() const { + return CreateTraversalCursor(kRootNodeId); } // Creates a cursor pointing to the 'node_id'. - TraversalCursor CreateTraversalCursor(uint32_t node_id) { + TraversalCursor CreateTraversalCursor(uint32_t node_id) const { + if (node_id >= trie_array_size_) { + return {node_id, 0}; + } return {node_id, trie_array_[node_id]}; } // Sets the cursor to point to 'node_id'. - void SetTraversalCursor(TraversalCursor& cursor, uint32_t node_id) { + void SetTraversalCursor(TraversalCursor& cursor, uint32_t node_id) const { cursor.node_id = node_id; + if (node_id >= trie_array_size_) { + cursor.unit = 0; + return; + } cursor.unit = trie_array_[node_id]; } @@ -83,7 +97,13 @@ class DartsCloneTrieWrapper { // 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 { + if (cursor.node_id >= trie_array_size_) { + return false; + } 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; @@ -99,34 +119,17 @@ class DartsCloneTrieWrapper { // false. bool TryTraverseSeveralSteps(TraversalCursor& cursor, absl::string_view path) const { - return TryTraverseSeveralSteps(cursor, path.data(), path.size()); - } - - // If the node pointed by 'cursor' has data, read into 'out_data' and returns - // true; otherwise, does nothing and returns false. - bool TryGetData(const TraversalCursor& cursor, int& out_data) const { - if (!has_leaf(cursor.unit)) { + if (cursor.node_id >= trie_array_size_) { return false; } - 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(const uint32_t* trie_array) - : trie_array_(trie_array) {} - - // The actual implementation of TryTraverseSeveralSteps. - bool TryTraverseSeveralSteps(TraversalCursor& cursor, const char* ptr, - int size) const { uint32_t cur_id = cursor.node_id; uint32_t cur_unit = cursor.unit; - for (; size > 0; --size, ++ptr) { - const unsigned char ch = static_cast(*ptr); + for (const char c : path) { + const unsigned char ch = static_cast(c); 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; @@ -137,6 +140,33 @@ class DartsCloneTrieWrapper { return true; } + // If the node pointed by 'cursor' has data, read into 'out_data' and returns + // true; otherwise, does nothing and returns false. + bool TryGetData(const TraversalCursor& cursor, int& out_data) const { + if (cursor.node_id >= trie_array_size_) { + return false; + } + 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]; + out_data = value(value_unit); + return true; + } + + // Returns the number of elements in the trie array. + size_t size() const { return trie_array_size_; } + + private: + // Use Create() instead of the constructor. + explicit DartsCloneTrieWrapper(const uint32_t* trie_array, + size_t trie_array_size) + : trie_array_(trie_array), trie_array_size_(trie_array_size) {} + // The helper functions below are based on // https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/custom_ops/kernel/sentencepiece/double_array_trie.h @@ -159,6 +189,8 @@ class DartsCloneTrieWrapper { // The pointer to the darts trie array. const uint32_t* trie_array_; + // The number of elements in the trie array. + size_t trie_array_size_; }; } // namespace trie_utils diff --git a/tensorflow_text/core/kernels/fast_bert_normalizer.h b/tensorflow_text/core/kernels/fast_bert_normalizer.h index c7b8f1849..403a32e93 100644 --- a/tensorflow_text/core/kernels/fast_bert_normalizer.h +++ b/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" @@ -78,23 +79,27 @@ class FastBertNormalizer { // Creates an instance. // // Args: - // * trie_data: the pointer to the trie data, which is not owned by this - // instance and should be kept alive through the lifetime of the instance. + // * trie_data: the trie data, which is not owned by this instance and should + // be kept alive through the lifetime of the instance. // * data_for_codepoint_zero: the mapped data for the codepoint zero. // * normalized_string_pool: the pointer to the normalized string pool data, // which is not owned by this instance and should be kept alive through the // lifetime of the instance. + // * normalized_string_pool_size: the size of the normalized string pool, if + // known. static absl::StatusOr Create( - const uint32_t* trie_data, int data_for_codepoint_zero, + absl::Span trie_data, int data_for_codepoint_zero, const char* normalized_string_pool, size_t normalized_string_pool_size = static_cast(-1)) { - if (trie_data == nullptr || normalized_string_pool == nullptr) { - return absl::InvalidArgumentError( - "trie_data or normalized_string_pool is null"); + if (trie_data.empty() || trie_data.data() == nullptr) { + return absl::InvalidArgumentError("trie_data is null or empty"); + } + if (normalized_string_pool == nullptr) { + return absl::InvalidArgumentError("normalized_string_pool is null"); } FastBertNormalizer result; - SH_ASSIGN_OR_RETURN(auto trie, - trie_utils::DartsCloneTrieWrapper::Create(trie_data)); + SH_ASSIGN_OR_RETURN( + auto trie, trie_utils::DartsCloneTrieWrapper::Create(trie_data)); result.trie_ = std::make_unique(std::move(trie)); result.data_for_codepoint_zero_ = data_for_codepoint_zero; @@ -123,7 +128,7 @@ class FastBertNormalizer { "FastBertNormalizerModel or its required fields are null"); } return Create( - model->trie_array()->data(), model->data_for_codepoint_zero(), + *model->trie_array(), model->data_for_codepoint_zero(), reinterpret_cast(model->normalized_string_pool()->data()), model->normalized_string_pool()->size()); } @@ -212,6 +217,7 @@ class FastBertNormalizer { auto copy_unchanged_input_to_output = [input_text, output_normalized_text, output_normalized_offset_mapping, &last_pos_to_copy_over](int exclusive_copy_end) { + (void)output_normalized_offset_mapping; // Copy from `last_pos_to_copy_over` to `exclusive_copy_end` and // update `last_pos_to_copy_over` accordingly. if (last_pos_to_copy_over < exclusive_copy_end) { diff --git a/tensorflow_text/core/kernels/fast_bert_normalizer_model_builder.cc b/tensorflow_text/core/kernels/fast_bert_normalizer_model_builder.cc index 808d2a09c..5907c330a 100644 --- a/tensorflow_text/core/kernels/fast_bert_normalizer_model_builder.cc +++ b/tensorflow_text/core/kernels/fast_bert_normalizer_model_builder.cc @@ -229,7 +229,8 @@ 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(), + mapped_value_pool_.size()); 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 " diff --git a/tensorflow_text/core/kernels/fast_bert_normalizer_test.cc b/tensorflow_text/core/kernels/fast_bert_normalizer_test.cc index 73b4a4c51..5887aabf4 100644 --- a/tensorflow_text/core/kernels/fast_bert_normalizer_test.cc +++ b/tensorflow_text/core/kernels/fast_bert_normalizer_test.cc @@ -14,10 +14,15 @@ #include "tensorflow_text/core/kernels/fast_bert_normalizer.h" +#include #include +#include +#include +#include #include #include +#include "absl/types/span.h" #include "tensorflow_text/core/kernels/fast_bert_normalizer_model_builder.h" namespace tensorflow { @@ -219,6 +224,29 @@ TEST_P(TestNormalization, TestNoGetOffsets) { INSTANTIATE_TEST_SUITE_P(FastBertNormalizerTest, TestNormalization, testing::ValuesIn(GetTestSpecs())); + +TEST(FastBertNormalizerCreateTest, NullOrEmptySpanFails) { + std::vector empty_vec; + EXPECT_FALSE(FastBertNormalizer::Create(empty_vec, 0, "").ok()); + EXPECT_FALSE( + FastBertNormalizer::Create(absl::Span(), 0, "").ok()); + uint32_t dummy = 0; + EXPECT_FALSE( + FastBertNormalizer::Create(absl::Span(&dummy, 0), 0, "") + .ok()); + EXPECT_FALSE( + FastBertNormalizer::Create(absl::Span(nullptr, 5), 0, "") + .ok()); + EXPECT_FALSE( + FastBertNormalizer::Create(absl::Span(&dummy, 1), 0, + nullptr) + .ok()); +} + +TEST(FastBertNormalizerCreateTest, NullModelFlatbufferFails) { + EXPECT_FALSE(FastBertNormalizer::Create(nullptr).ok()); +} + } // namespace } // namespace text } // namespace tensorflow diff --git a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer.cc b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer.cc index ff10541f4..3abaec784 100644 --- a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer.cc +++ b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer.cc @@ -56,7 +56,7 @@ 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()); + *tokenizer.config_->trie_array()); if (!trie_or.ok()) { return absl::InvalidArgumentError( "Failed to create DartsCloneTrieWrapper from " diff --git a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc index 9467c4d6e..5aed0bb42 100644 --- a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc +++ b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc @@ -136,8 +136,7 @@ class FastWordpieceBuilder { absl::Status BuildModel(const std::vector& vocab, int max_bytes_per_token, absl::string_view suffix_indicator, - absl::string_view unk_token, - bool no_pretokenization, + absl::string_view unk_token, bool no_pretokenization, bool support_detokenization); absl::StatusOr ExportToFlatBuffer() const; @@ -432,9 +431,8 @@ absl::Status FastWordpieceBuilder::ConstructTrie( } SH_ASSIGN_OR_RETURN(trie_array_, trie_utils::BuildDartsCloneTrie(keys, values)); - SH_ASSIGN_OR_RETURN( - trie_utils::DartsCloneTrieWrapper trie, - trie_utils::DartsCloneTrieWrapper::Create(trie_array_.data())); + SH_ASSIGN_OR_RETURN(trie_utils::DartsCloneTrieWrapper trie, + trie_utils::DartsCloneTrieWrapper::Create(trie_array_)); trie_.emplace(std::move(trie)); if (trie_array_.size() > @@ -484,23 +482,22 @@ absl::Status FastWordpieceBuilder::BuildOutgoingEdgeLabelsAlongVocabToken( if (!trie_->TryTraverseOneStep(cur_node, edge_label)) { // Should never happen, since we built trie using all of `vocab_token`. return absl::FailedPreconditionError(absl::StrCat( - "Cannot traverse from parent id ", cur_node.node_id, - " to child following the edge with label value of ", - static_cast(edge_label), - " when processing a vocabulary token with token ID ", - vocab_token.TokenId(), " (0-based). This error happened at ", - "position ", char_pos, " (0-based) of the token. Before that, ", - "the prefix \"", token.substr(0, char_pos), - "\" of the token had been processed. This should never happen. ", - "This probably indicates that there are some unicode ", - "issues (e.g., byte '\\x0' in the middle) for the above ", - "mentioned token in the vocabulary file. All bytes of this ", - "questionable token (ID ", vocab_token.TokenId(), ") are: [", - absl::StrJoin( - iter::imap([](auto ch) { return static_cast(ch); }, - vocab_token.Token()), - ", "), - "].")); + "Cannot traverse from parent id ", cur_node.node_id, + " to child following the edge with label value of ", + static_cast(edge_label), + " when processing a vocabulary token with token ID ", + vocab_token.TokenId(), " (0-based). This error happened at ", + "position ", char_pos, " (0-based) of the token. Before that, ", + "the prefix \"", token.substr(0, char_pos), + "\" of the token had been processed. This should never happen. ", + "This probably indicates that there are some unicode ", + "issues (e.g., byte '\\x0' in the middle) for the above ", + "mentioned token in the vocabulary file. All bytes of this ", + "questionable token (ID ", vocab_token.TokenId(), ") are: [", + absl::StrJoin(iter::imap([](auto ch) { return static_cast(ch); }, + vocab_token.Token()), + ", "), + "].")); } ++char_pos; } @@ -929,10 +926,9 @@ absl::StatusOr BuildModelAndExportToFlatBuffer( absl::string_view suffix_indicator, absl::string_view unk_token, bool no_pretokenization, bool support_detokenization) { FastWordpieceBuilder builder; - SH_RETURN_IF_ERROR(builder.BuildModel(vocab, max_bytes_per_token, - suffix_indicator, unk_token, - no_pretokenization, - support_detokenization)); + SH_RETURN_IF_ERROR(builder.BuildModel( + vocab, max_bytes_per_token, suffix_indicator, unk_token, + no_pretokenization, support_detokenization)); SH_ASSIGN_OR_RETURN(std::string flatbuffer, builder.ExportToFlatBuffer()); return flatbuffer; } diff --git a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_test.cc b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_test.cc index fea96e3ef..b000618bb 100644 --- a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_test.cc +++ b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_test.cc @@ -2549,6 +2549,10 @@ INSTANTIATE_TEST_SUITE_P( FastWordpieceTokenizerDetokenizeParameterizedTest, TestTokenizeDetokenize, testing::ValuesIn(GetTestSpecsForTokenizeDetokenize())); +TEST(FastWordpieceTokenizerCreateTest, NullConfigFlatbufferFails) { + EXPECT_FALSE(FastWordpieceTokenizer::Create(nullptr).ok()); +} + } // namespace } // namespace text } // namespace tensorflow