Improve ONNX Runtime input validation

operators/tokenizer/bert_tokenizer.cc

@@ -374,11 +374,18 @@ void KernelBertTokenizer::Compute(const ortc::Tensor<std::string>& input,
   auto* p_out2 = output2.Allocate(output_dim);
   std::copy(attention_mask.begin(), attention_mask.end(), p_out2);
 
-  std::vector<int64_t> offset_dim{static_cast<int64_t>(input_ids.size()), 2};  // tuple of offsets for each input id
-
   if (offset_mapping.has_value()) {
+    // Size the output from the same source that drives the write loop below. The
+    // number of offset pairs produced by tokenization does not necessarily equal
+    // input_ids.size() (e.g. truncation shrinks input_ids but not offset_map), so
+    // sizing from input_ids would under-allocate and overflow the buffer.
+    size_t offset_rows = 0;
+    for (auto& res : offset_map) {
+      offset_rows += res.size();
+    }
+    std::vector<int64_t> offset_dim{static_cast<int64_t>(offset_rows), 2};  // tuple of offsets for each input id
     auto* offset = (*offset_mapping)->Allocate(offset_dim);
-    int idx2 = 0;
+    size_t idx2 = 0;
     for (auto& res : offset_map) {
       for (auto& mapping : res) {
         offset[idx2] = mapping.first;
@@ -436,10 +443,17 @@ void KernelHfBertTokenizer::Compute(const ortc::Tensor<std::string>& input,
     std::copy(token_type_ids.begin(), token_type_ids.end(), p_out2);
   }
 
-  std::vector<int64_t> offset_dim{static_cast<int64_t>(input_ids.size()), 2};  // tuple of offsets for each input id
   if (compute_offset_mapping) {
+    // Size the output from the same source that drives the write loop below,
+    // not from input_ids.size(); the two counts diverge and sizing from
+    // input_ids would under-allocate and overflow the buffer.
+    size_t offset_rows = 0;
+    for (auto& res : offset_map) {
+      offset_rows += res.size();
+    }
+    std::vector<int64_t> offset_dim{static_cast<int64_t>(offset_rows), 2};  // tuple of offsets for each input id
     auto* offset = (*offset_mapping)->Allocate(offset_dim);
-    int idx2 = 0;
+    size_t idx2 = 0;
     for (auto& res : offset_map) {
       for (auto& mapping : res) {
         offset[idx2] = mapping.first;

test/test_bert_tokenizer.py

@@ -34,6 +34,27 @@ def _run_combined_case(input, vocab_path):
     np.testing.assert_array_equal(result[1], expect_result["token_type_ids"])
     np.testing.assert_array_equal(result[2], expect_result["attention_mask"])
 
+def _run_truncated_offset_bounds_check(input, vocab_path):
+    # Truncation shortens input_ids after offset_map has been generated. The
+    # output tensor must be allocated from offset_map so its write loop cannot
+    # run past a buffer sized from the truncated IDs.
+    t2stc = PyOrtFunction.from_customop(
+        BertTokenizer,
+        vocab_file=vocab_path,
+        do_lower_case=0,
+        strip_accents=1,
+        max_length=5,
+    )
+    result = t2stc([input])
+    offset_mapping = np.asarray(result[3])
+    input_ids = np.asarray(result[0])
+    assert offset_mapping.shape[1] == 2, offset_mapping.shape
+    assert offset_mapping.shape[0] > input_ids.shape[0], (
+        offset_mapping.shape,
+        input_ids.shape,
+    )
+
+
 def _run_basic_with_offset_check(input, vocab_path):
     t2stc = PyOrtFunction.from_customop(
         BertTokenizer, vocab_file=vocab_path, do_lower_case=0, strip_accents=1
@@ -80,6 +101,10 @@ def test_text_to_case1(self):
             ["网 易 云 音 乐", "cat isnot playing toyssss"],
             vocab_path=util.get_test_data_file("data", "bert_basic_cased_vocab.txt"),
         )
+        _run_truncated_offset_bounds_check(
+            "cat isnot playing toyssss",
+            vocab_path=util.get_test_data_file("data", "bert_basic_cased_vocab.txt"),
+        )
         print("\n****** Input ids, token type ids, and attention mask tests complete. ******\n\n\n")
         print("*** Starting offset mapping tests. ***\n")