From 4f51e109a3beb03e161fd07d77e4695dd9321cec Mon Sep 17 00:00:00 2001
From: bhkumar007 <45501967+bhkumar007@users.noreply.github.com>
Date: Mon, 23 Mar 2026 13:35:12 -0700
Subject: [PATCH] Calculate num_strings based on all dimensions for
 OrtStringViewTensorStorage
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary
OrtStringViewTensorStorage computes num_strings using only the first shape dimension (shape[0]) instead of the product of all dimensions. This causes a heap buffer overflow when processing string tensors with rank ≥ 2.

Bug
In custom_op_lite.h, the OrtStringViewTensorStorage constructor calculates the number of strings as:

size_t num_strings = 1;
if ((*shape_).size() > 0) {
    num_strings = static_cast<size_t>((*shape_)[0]);  // ← only first dimension
}

This under-counts the total number of strings for any tensor with rank ≥ 2. The under-sized offsets vector is then passed to GetStringTensorContent, which writes beyond its bounds.

For example, a string tensor with shape [2, 3] contains 6 strings, but this code allocates offsets for only 2.

Note: the sibling class OrtStringTensorStorage computes this correctly by multiplying all dimensions.

Fix
Replace the partial dimension calculation with a simple loop over all shape dimensions:

While String tensors in ONNX are often rank 1 in practice, The ORT API doesn't care about intent. GetStringTensorContent expects offsets_count to equal the actual total number of strings in the tensor. If ORT has a [2, 3] tensor with 6 strings and you pass an offsets buffer of size 2, it writes 6 entries into a 2-slot buffer. There's no graceful "I only want the first dimension's worth" mode — it's an unconditional overflow. Also, the sibling class OrtStringTensorStorage handles all dimensions, and both classes implement the same IStringTensorStorage interface — callers have no way to know which backing class they're hitting
---
 include/custom_op/custom_op_lite.h | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/include/custom_op/custom_op_lite.h b/include/custom_op/custom_op_lite.h
index b667d4cf0..44eed159e 100644
--- a/include/custom_op/custom_op_lite.h
+++ b/include/custom_op/custom_op_lite.h
@@ -217,7 +217,9 @@ class OrtStringViewTensorStorage : public IStringTensorStorage<std::string_view>
 
       size_t num_strings = 1;
       if ((*shape_).size() > 0) {
-        num_strings = static_cast<size_t>((*shape_)[0]);
+        for (auto dim : *shape_) {
+          num_strings *= static_cast<size_t>(dim);
+        }
       }
 
       if (num_strings) {