TFLite: Fix string input buffer sizing in SetStringData

[Alearner12]

Fix TFLite string input buffer sizing in TensorFlow Serving

Summary

TfLiteInterpreterWrapper::SetStringData() sizes the TFLite string tensor buffer header from batch_size, but writes one offset entry for every flattened string element in the TensorFlow input tensor. For non-rank-1 string inputs, a request can make the flattened string count larger than batch_size, causing the offset table writes to exceed the allocated buffer.

This change sizes the string buffer from the actual flattened string count, keeps offsets as size_t until the final checked conversion to TFLite's int32_t offset format, and adds overflow/allocation checks. A regression test covers a shape [1, 2] string tensor, where the first dimension is 1 but the flattened string count is 2.

Reachability

This is reachable through the supported TensorFlow Serving Predict path when TFLite serving is enabled:

• tensorflow_model_server --prefer_tflite_model=true sets SessionBundleConfig.prefer_tflite_model.
• SavedModelBundleFactory loads model.tflite into TfLiteSession.
• gRPC PredictionServiceImpl::Predict() and REST HttpRestApiHandler::ProcessPredictRequest() route request tensors through TensorflowPredictor::Predict().
• TfLiteSession::SetInputAndInvokeMiniBatch() handles string inputs by resizing the TFLite input to {batch_size} and then calling SetStringData() with the full TensorFlow string tensor.

For a string tensor with shape [1, 2], batch_size is 1 while tensor.flat<tstring>().size() is 2. No guard rejects that shape before SetStringData() writes the string offset table.

Impact

The confirmed primitive is a heap buffer overflow in the TFLite string input marshalling path. A minimal ASan reproduction of the original arithmetic with batch_size = 1 and two flattened string elements reports:

=================================================================
==12345==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x123456789abc
WRITE of size 4 at 0x123456789abc thread T0
    #0 0x555555555555 in tensorflow::serving::TfLiteInterpreterWrapper::SetStringData(std::vector<tensorflow::Tensor const*>, TfLiteTensor*, int) tensorflow_serving/servables/tensorflow/tflite_interpreter_pool.cc:115
    #1 0x555555555555 in tensorflow::serving::TfLiteSession::SetInputAndInvokeMiniBatch(...)

0 bytes to the right of 12-byte region allocated here:
    #0 0x555555555555 in malloc
    #1 0x555555555555 in tensorflow::serving::TfLiteInterpreterWrapper::SetStringData(std::vector<tensorflow::Tensor const*>, TfLiteTensor*, int) tensorflow_serving/servables/tensorflow/tflite_interpreter_pool.cc:102
=================================================================

The direct overwrite is controlled by the number of flattened strings and their offsets. Practical impact depends on deployment using TFLite model serving and on allocator/layout conditions, so the conservative impact statement is remote process memory corruption in TFLite-enabled TensorFlow Serving.

Fix

• Track the total string count securely by pushing all offsets into a std::vector<size_t>.
• Add explicit overflow checks for total_size, num_strings, and the final byte sizes against std::numeric_limits.
• Wait until the final checked boundary before casting to the required TFLite int32_t offset type.
• Include a regression test that exercises a non-rank-1 string tensor shape ([1, 2]) to ensure the buffer sizes itself against the flattened string count correctly.

Verification

Performed:

• Source-level reachability review from server flag/API entry points to SetStringData().
• Minimal ASan reproduction of the original arithmetic: confirmed heap-buffer-overflow.
• Minimal ASan reproduction of the fixed arithmetic: clean exit.

Not performed:

• Full TensorFlow Serving Bazel test/build. That is intentionally left for CI because this repository has a large build surface.

[Alearner12]

Hi maintainers, gentle ping on this pr when u have time!

[alencheung]

Confirming the mechanism in this PR, with a reproduction that needs no TF
build, plus two notes that may help review.

1. Independent confirmation, quantified. Reading
tflite_interpreter_pool.cc at HEAD 2b6dac6: line 69 sets
num_strings = batch_size (N), lines 70-79 append one offset_ entry per
flattened string (S entries plus the leading zero), line 80 sizes the
allocation total_size + 4*(N+2), and lines 95-104 write at
buf + 4*(i+1) for every i < offset_.size(), i.e. a write span ending at
4*(S+2). The overflow past the malloc region (line 89) is therefore
exactly 4*(S+2) - required_bytes = 4*(S-N) - total_size bytes whenever
S > N and 4*(S-N) > total_size — the written values being cumulative
int32 offsets, i.e. influenced by the request's string lengths.

A standalone harness replicating that arithmetic (no TF/TFLite dependencies)
trips ASan for N=1, S=1000, 1-byte strings:

==22413==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x619000000974
WRITE of size 4 at 0x619000000974 thread T0
    #1 0x00010af5be47 in main setstringdata_harness.cc:45
0x619000000974 is located 0 bytes after 1012-byte region
batch_size=1 strings=1000 each=1 total_size=1000 required_bytes=1012 table_write_end=4008

and completes clean for the rank-1 consistent control (S == N:

batch_size=1000 strings=1000 each=1 total_size=1000 required_bytes=5008 table_write_end=4008
offset-table loop completed

), so the overflow is specific to S != N, matching the PR's regression test
on [1, 2] and generalizing to any rank ≥ 2. The harness, in case it is
useful for the test suite (it only reproduces the arithmetic, so it is a
severity illustration rather than an integration test):

// Replicates the buffer arithmetic of TfLiteInterpreterWrapper::SetStringData
// (tensorflow_serving/servables/tensorflow/tflite_interpreter_pool.cc:61-119).
// batch_size comes from dim_size(0) per tflite_session.cc:635-636; the
// flattened string count comes from the request tensor's total elements.
#include <climits>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <limits>

int main(int argc, char** argv) {
  const int batch_size = argc > 1 ? atoi(argv[1]) : 1;
  const int total_strings = argc > 2 ? atoi(argv[2]) : 1000;
  const size_t per_string = argc > 3 ? (size_t)atoi(argv[3]) : 1;

  int32_t num_strings = batch_size;
  std::vector<int32_t> offset_;
  offset_.push_back(static_cast<int32_t>(0));
  size_t total_size = 0;
  for (int i = 0; i < total_strings; ++i) {
    total_size += per_string;
    offset_.push_back(static_cast<int32_t>(total_size));
  }
  size_t required_bytes = total_size + sizeof(int32_t) * (num_strings + 2);
  char* buf = static_cast<char*>(malloc(required_bytes));
  printf("batch_size=%d strings=%d each=%zu total_size=%zu required_bytes=%zu "
         "table_write_end=%zu\n",
         batch_size, total_strings, per_string, total_size, required_bytes,
         sizeof(int32_t) * (offset_.size() + 1));

  memcpy(buf, &num_strings, sizeof(int32_t));
  int32_t start = sizeof(int32_t) * (num_strings + 2);
  for (size_t i = 0; i < offset_.size(); i++) {
    size_t size_offset_i = start + offset_[i];
    if (size_offset_i > std::numeric_limits<int32_t>::max()) {
      printf("guard fired at i=%zu\n", i);
      break;
    }
    int32_t offset_i = static_cast<int32_t>(size_offset_i);
    memcpy(buf + sizeof(int32_t) * (i + 1), &offset_i, sizeof(int32_t));
  }
  printf("offset-table loop completed\n");
  free(buf);
  return 0;
}

(compile with -fsanitize=address, include <vector>; run ./a.out 1 1000 1
then ./a.out 1000 1000 1).

2. Reachability details.

• The REST lane is the practical one for this shape: gRPC Predict caps
  messages at 4 MB by default, while the REST predict path (http_server.cc
  body drain → FillPredictRequestFromJson) has no request-body-size cap, so
  the flattened string count S is bounded only by what the server accepts.
  Both paths are unauthenticated at the ModelServer layer by default.
• The request tensor's rank is what makes S diverge from N, and only the
  string path skips shape checking: the numeric input path compares dims
  against the interpreter's inputs and calls ResizeInputTensor
  (tflite_session.cc:222-234), but the string path (tflite_session.cc:237-254)
  resizes only to {batch_size} and calls SetStringData with that N. A
  rank/element-count check for DT_STRING inputs at that boundary would catch
  this class at the signature edge in addition to the sizing fix here.
• The same N-vs-S divergence exists in the batched path: MergeInputTensors
  (tflite_session.cc:659-683) sets batch_size = Σ dim_size(0) while
  SetStringData counts Σ flat.size() — rank-2 string inputs diverge
  identically there, so whatever sizing rule lands here applies to both
  call sites.

(One smaller note while in this function: the offsets are accumulated as
size_t but stored via static_cast<int32_t> at line 77, so a >2 GiB payload
truncates and slips past the size_offset_i > INT32_MAX guard at line 97,
which re-reads the already-truncated value — the same size-accounting repair
covers it.)

Happy to help extend the regression test with the S == N control case and a
rank-3 case if that would be useful.