Fix KFS multi-dimensional string tensor deserialization

src/kfs_frontend/kfs_utils.cpp

@@ -16,6 +16,7 @@
 #include "kfs_utils.hpp"
 
 #include <algorithm>
+#include <limits>
 #include <map>
 #include <memory>
 #include <sstream>
@@ -292,6 +293,34 @@ int getBinaryInputsSize(const ::KFSRequest::InferInputTensor& tensor) {
     return tensor.contents().bytes_contents_size();
 }
 
+Status buildShapeFromStringTensorRequest(const ::KFSRequest::InferInputTensor& src, size_t numElements, ov::Shape& shape) {
+    if (src.shape_size() > 0) {
+        size_t shapeElements = 1;
+        bool validShape = true;
+        for (int i = 0; i < src.shape_size(); i++) {
+            int64_t dim = src.shape().at(i);
+            if (dim < 0) {
+                validShape = false;
+                break;
+            }
+            size_t dimSize = static_cast<size_t>(dim);
+            if (dimSize != 0 && shapeElements > std::numeric_limits<size_t>::max() / dimSize) {
+                validShape = false;  // multiplication would overflow
+                break;
+            }
+            shape.push_back(dimSize);
+            shapeElements *= dimSize;
+        }
+        if (!validShape || shapeElements != numElements) {
+            SPDLOG_DEBUG("Input string shape mismatch: shape product {} != parsed string count {}; rejecting request", shapeElements, numElements);
+            return StatusCode::INVALID_STRING_INPUT;
+        }
+    } else {
+        shape = ov::Shape{numElements};
+    }
+    return StatusCode::OK;
+}
+
 Status convertBinaryExtensionStringFromBufferToNativeOVTensor(const ::KFSRequest::InferInputTensor& src, ov::Tensor& tensor, const std::string* buffer) {
     std::vector<uint32_t> stringSizes;
     uint32_t totalStringsLength = 0;
@@ -305,7 +334,12 @@ Status convertBinaryExtensionStringFromBufferToNativeOVTensor(const ::KFSRequest
         SPDLOG_DEBUG("Input string format conversion failed");
         return StatusCode::INVALID_STRING_INPUT;
     }
-    tensor = ov::Tensor(ov::element::Type_t::string, ov::Shape{batchSize});
+    ov::Shape shape;
+    auto status = buildShapeFromStringTensorRequest(src, batchSize, shape);
+    if (!status.ok()) {
+        return status;
+    }
+    tensor = ov::Tensor(ov::element::Type_t::string, shape);
     std::string* data = tensor.data<std::string>();
     size_t tensorStringsOffset = 0;
     for (size_t i = 0; i < stringSizes.size(); i++) {

src/kfs_frontend/kfs_utils.hpp

@@ -73,4 +73,5 @@ Status validateTensor(const TensorInfo& tensorInfo,
 Status convertBinaryExtensionStringFromBufferToNativeOVTensor(const ::KFSRequest::InferInputTensor& src, ov::Tensor& tensor, const std::string* buffer);
 const std::string& getBinaryInput(const ::KFSRequest::InferInputTensor& tensor, size_t i);
 int getBinaryInputsSize(const ::KFSRequest::InferInputTensor& tensor);
+Status buildShapeFromStringTensorRequest(const ::KFSRequest::InferInputTensor& src, size_t numElements, ov::Shape& shape);
 }  // namespace ovms

src/kfs_frontend/tensor_conversion.cpp

@@ -18,6 +18,7 @@
 #include "../tensor_conversion_common.hpp"
 
 #include <algorithm>
+#include <limits>
 #include <memory>
 #include <string>
 #include <utility>
@@ -61,7 +62,28 @@ Status convertStringRequestFromBufferToOVTensor2D(const ::KFSRequest::InferInput
     return StatusCode::OK;
 }
 
-template Status convertStringRequestToOVTensor<::KFSRequest::InferInputTensor>(const ::KFSRequest::InferInputTensor& src, ov::Tensor& tensor, const std::string* buffer);
+template <>
+Status convertStringRequestToOVTensor<::KFSRequest::InferInputTensor>(const ::KFSRequest::InferInputTensor& src, ov::Tensor& tensor, const std::string* buffer) {
+    OVMS_PROFILE_FUNCTION();
+    if (buffer != nullptr) {
+        return convertBinaryExtensionStringFromBufferToNativeOVTensor(src, tensor, buffer);
+    }
+    int numElements = getBinaryInputsSize(src);
+    if (numElements < 0) {
+        return StatusCode::INVALID_STRING_INPUT;
+    }
+    ov::Shape shape;
+    auto status = buildShapeFromStringTensorRequest(src, static_cast<size_t>(numElements), shape);
+    if (!status.ok()) {
+        return status;
+    }
+    tensor = ov::Tensor(ov::element::Type_t::string, shape);
+    std::string* data = tensor.data<std::string>();
+    for (int i = 0; i < numElements; i++) {
+        data[i].assign(getBinaryInput(src, i));
+    }
+    return StatusCode::OK;
+}
 template Status convertNativeFileFormatRequestTensorToOVTensor<::KFSRequest::InferInputTensor>(const ::KFSRequest::InferInputTensor& src, ov::Tensor& tensor, const TensorInfo& tensorInfo, const std::string* buffer);
 }  // namespace ovms
 

src/predict_request_validation_utils.hpp

@@ -323,6 +323,9 @@ Status RequestValidator<RequestType, InputTensorType, choice, IteratorType, Shap
                     inputWidth = getStringInputWidth(*proto);
                 }
                 if (processingHint == TensorInfo::ProcessingHint::STRING_NATIVE) {
+                    // Guard against ov::Tensor memory amplification: N small strings in proto
+                    // allocate N*sizeof(std::string) objects even for empty content.
+                    RETURN_IF_ERR(validateAgainstMaxNativeStringElementCount(inputBatchSize));
                     // Pass through to normal validation
                 } else if (processingHint == TensorInfo::ProcessingHint::STRING_2D_U8) {
                     SPDLOG_DEBUG("[servable name: {} version: {}] Validating request containing 2D string input: name: {}",

src/predict_request_validation_utils_impl.cpp

@@ -45,6 +45,20 @@ Status validateAgainstMax2DStringArraySize(int32_t inputBatchSize, size_t inputW
     }
     return StatusCode::OK;
 }
+Status validateAgainstMaxNativeStringElementCount(int32_t elementCount) {
+    if (elementCount < 0) {
+        SPDLOG_DEBUG("Invalid(negative) element count: {}", elementCount);
+        return Status(StatusCode::INVALID_BATCH_SIZE, "Batch size is negative");
+    }
+    if (static_cast<size_t>(elementCount) > MAX_NATIVE_STRING_ELEMENTS) {
+        std::stringstream ss;
+        ss << "; element count " << elementCount << " exceeds max " << MAX_NATIVE_STRING_ELEMENTS;
+        const std::string details = ss.str();
+        SPDLOG_DEBUG(details);
+        return Status(StatusCode::INVALID_STRING_MAX_SIZE_EXCEEDED, details);
+    }
+    return StatusCode::OK;
+}
 Mode getShapeMode(const shapes_info_map_t& shapeInfo, const std::string& name) {
     if (shapeInfo.size() == 0) {
         return Mode::FIXED;

src/predict_request_validation_utils_impl.hpp

@@ -28,9 +28,13 @@
 
 namespace ovms {
 namespace request_validation_utils {
-const size_t MAX_2D_STRING_ARRAY_SIZE = 1024 * 1024 * 1024 * 1;                                             // 1GB
+const size_t MAX_2D_STRING_ARRAY_SIZE = 1024 * 1024 * 1024 * 1;  // 1GB
+// Each ov::Tensor(string, {N}) allocates N std::string objects (~32 B each).
+// Cap element count so string-object heap growth stays within 1 GB.
+const size_t MAX_NATIVE_STRING_ELEMENTS = (1ULL << 30) / sizeof(std::string);                               // ~33.5M on 64-bit
 Status getRawInputContentsBatchSizeAndWidth(const std::string& buffer, int32_t& batchSize, size_t& width);  // this comes from KFS - may need to move there
 Status validateAgainstMax2DStringArraySize(int32_t inputBatchSize, size_t inputWidth);
+Status validateAgainstMaxNativeStringElementCount(int32_t elementCount);
 Mode getShapeMode(const shapes_info_map_t& shapeInfo, const std::string& name);
 }  // namespace request_validation_utils
 }  // namespace ovms

src/tensor_conversion.hpp

@@ -28,8 +28,8 @@
 #include "tensorinfo.hpp"
 #include "status.hpp"
 
-namespace tensorflow {
-class TensorProto;
+namespace inference {
+class ModelInferRequest_InferInputTensor;
 }
 namespace ovms {
 class Status;
@@ -156,24 +156,12 @@ class InferenceTensor;
 template <>
 Status convertNativeFileFormatRequestTensorToOVTensor(const InferenceTensor& src, ov::Tensor& tensor, const TensorInfo& tensorInfo, const std::string* buffer);
 
-template <typename TensorType>
-Status convertStringRequestToOVTensor(const TensorType& src, ov::Tensor& tensor, const std::string* buffer) {
-    OVMS_PROFILE_FUNCTION();
-    if (buffer != nullptr) {
-        return convertBinaryExtensionStringFromBufferToNativeOVTensor(src, tensor, buffer);
-    }
-    int batchSize = getBinaryInputsSize(src);
-    tensor = ov::Tensor(ov::element::Type_t::string, ov::Shape{static_cast<size_t>(batchSize)});
-    std::string* data = tensor.data<std::string>();
-    for (int i = 0; i < batchSize; i++) {
-        data[i].assign(getBinaryInput(src, i));
-    }
-    return StatusCode::OK;
-}
-
 template <>
 Status convertStringRequestToOVTensor(const InferenceTensor& src, ov::Tensor& tensor, const std::string* buffer);
 
+template <>
+Status convertStringRequestToOVTensor(const ::inference::ModelInferRequest_InferInputTensor& src, ov::Tensor& tensor, const std::string* buffer);
+
 template <typename TensorType>
 Status convertOVTensor2DToStringResponse(const ov::Tensor& tensor, TensorType& dst) {
     if (tensor.get_shape().size() != 2) {

src/test/predict_validation_test.cpp

@@ -1303,6 +1303,53 @@ TYPED_TEST(PredictValidationStringNativeTest, string_not_allowed_with_demultiple
     EXPECT_EQ(status, ovms::StatusCode::INVALID_NO_OF_SHAPE_DIMENSIONS);
 }
 
+TEST(PredictValidationStringNativeKFSTest, negative_over_element_count_limit) {
+    const char* tensorName = DUMMY_MODEL_INPUT_NAME;
+    ovms::tensor_map_t mockedInputsInfo, mockedOutputsInfo;
+    mockedInputsInfo[tensorName] = std::make_shared<ovms::TensorInfo>(
+        tensorName, ovms::Precision::STRING, ovms::Shape{-1}, ovms::Layout{"N..."});
+
+    const size_t N = ovms::request_validation_utils::MAX_NATIVE_STRING_ELEMENTS + 1;
+    // Build raw buffer: N entries, each with 4-byte zero length (empty string).
+    // Buffer size = N * 4 bytes (~128 MB), much less than N * sizeof(std::string) (~1 GB).
+    std::string rawBuffer(N * sizeof(uint32_t), '\0');
+
+    ::KFSRequest request;
+    auto* input = request.add_inputs();
+    input->set_name(tensorName);
+    input->set_datatype("BYTES");
+    input->add_shape(static_cast<int64_t>(N));
+    *request.add_raw_input_contents() = std::move(rawBuffer);
+
+    auto status = ovms::request_validation_utils::validate(
+        request, mockedInputsInfo, mockedOutputsInfo, "dummy", ovms::model_version_t{1});
+    EXPECT_EQ(status, ovms::StatusCode::INVALID_STRING_MAX_SIZE_EXCEEDED) << status.string();
+}
+
+TEST(PredictValidationStringNativeKFSTest, negative_over_element_count_limit_contents) {
+    const char* tensorName = DUMMY_MODEL_INPUT_NAME;
+    ovms::tensor_map_t mockedInputsInfo, mockedOutputsInfo;
+    mockedInputsInfo[tensorName] = std::make_shared<ovms::TensorInfo>(
+        tensorName, ovms::Precision::STRING, ovms::Shape{-1}, ovms::Layout{"N..."});
+
+    const int32_t N = static_cast<int32_t>(ovms::request_validation_utils::MAX_NATIVE_STRING_ELEMENTS) + 1;
+
+    ::KFSRequest request;
+    auto* input = request.add_inputs();
+    input->set_name(tensorName);
+    input->set_datatype("BYTES");
+    input->add_shape(static_cast<int64_t>(N));
+    auto* contents = input->mutable_contents()->mutable_bytes_contents();
+    contents->Reserve(N);
+    for (int32_t i = 0; i < N; i++) {
+        contents->Add();
+    }
+
+    auto status = ovms::request_validation_utils::validate(
+        request, mockedInputsInfo, mockedOutputsInfo, "dummy", ovms::model_version_t{1});
+    EXPECT_EQ(status, ovms::StatusCode::INVALID_STRING_MAX_SIZE_EXCEEDED) << status.string();
+}
+
 #define VERIFY_COMPUTE_BUFFER_SIZE(SHAPE, ELEMENT_SIZE, WILL_NOT_OVERFLOW, EXPECTED_BYTES)                                                  \
     {                                                                                                                                       \
         size_t elementSize = ELEMENT_SIZE;                                                                                                  \

src/test/tensor_conversion_test.cpp

@@ -711,6 +711,102 @@ TYPED_TEST(StringInputsConversionTest, rawInputContents_native_ov_string) {
     }
 }
 
+TEST(StringInputsConversionKFSTest, native_ov_string_2d_shape) {
+    // Request shape [1, 5] with 5 strings - simulates 2D string input (batch=1, seq_len=5)
+    ::KFSRequest::InferInputTensor requestTensor;
+    requestTensor.set_datatype("BYTES");
+    requestTensor.mutable_shape()->Clear();
+    requestTensor.add_shape(1);
+    requestTensor.add_shape(5);
+    std::vector<std::string> expectedStrings = {"aa", "bbb", "c", "dddd", "e"};
+    for (const auto& s : expectedStrings) {
+        auto bytes_val = requestTensor.mutable_contents()->mutable_bytes_contents()->Add();
+        bytes_val->append(s.data(), s.size());
+    }
+    ov::Tensor tensor;
+    ASSERT_EQ(convertStringRequestToOVTensor(requestTensor, tensor, nullptr), ovms::StatusCode::OK);
+    ASSERT_EQ(tensor.get_element_type(), ov::element::string);
+    ASSERT_EQ(tensor.get_shape().size(), 2);
+    ASSERT_THAT(tensor.get_shape(), ::testing::ElementsAre(1, 5));
+    std::string* data = tensor.data<std::string>();
+    for (size_t i = 0; i < expectedStrings.size(); i++) {
+        ASSERT_EQ(data[i], expectedStrings[i]) << " at index " << i;
+    }
+}
+
+TEST(StringInputsConversionKFSTest, rawInputContents_native_ov_string_2d_shape) {
+    // raw_input_contents path: shape [1, 5] with 5 strings in buffer
+    ::KFSRequest::InferInputTensor requestTensor;
+    requestTensor.set_datatype("BYTES");
+    requestTensor.mutable_shape()->Clear();
+    requestTensor.add_shape(1);
+    requestTensor.add_shape(5);
+    std::vector<std::string> expectedStrings = {"a", "bbbbb", "c", "dd", "eee"};
+    std::string rawInputContents;
+    size_t dataSize = 0;
+    for (const auto& s : expectedStrings) {
+        dataSize += s.size() + sizeof(uint32_t);
+    }
+    rawInputContents.resize(dataSize);
+    size_t offset = 0;
+    for (const auto& s : expectedStrings) {
+        uint32_t inputSize = s.size();
+        std::memcpy(rawInputContents.data() + offset, &inputSize, sizeof(uint32_t));
+        offset += sizeof(uint32_t);
+        std::memcpy(rawInputContents.data() + offset, s.data(), s.size());
+        offset += s.size();
+    }
+    ov::Tensor tensor;
+    ASSERT_EQ(convertStringRequestToOVTensor(requestTensor, tensor, &rawInputContents), ovms::StatusCode::OK);
+    ASSERT_EQ(tensor.get_element_type(), ov::element::string);
+    ASSERT_EQ(tensor.get_shape().size(), 2);
+    ASSERT_THAT(tensor.get_shape(), ::testing::ElementsAre(1, 5));
+    std::string* data = tensor.data<std::string>();
+    for (size_t i = 0; i < expectedStrings.size(); i++) {
+        ASSERT_EQ(data[i], expectedStrings[i]) << " at index " << i;
+    }
+}
+
+TEST(StringInputsConversionKFSTest, native_ov_string_shape_mismatch_invalid) {
+    // Shape product [2,5]=10 does not match 3 provided strings - request must be rejected
+    ::KFSRequest::InferInputTensor requestTensor;
+    requestTensor.set_datatype("BYTES");
+    requestTensor.add_shape(2);
+    requestTensor.add_shape(5);
+    std::vector<std::string> strings = {"aa", "bbb", "c"};
+    for (const auto& s : strings) {
+        auto bytes_val = requestTensor.mutable_contents()->mutable_bytes_contents()->Add();
+        bytes_val->append(s.data(), s.size());
+    }
+    ov::Tensor tensor;
+    ASSERT_EQ(convertStringRequestToOVTensor(requestTensor, tensor, nullptr), ovms::StatusCode::INVALID_STRING_INPUT);
+}
+
+TEST(StringInputsConversionKFSTest, rawInputContents_native_ov_string_shape_mismatch_invalid) {
+    // Shape product [2,5]=10 does not match 3 strings in buffer - request must be rejected
+    ::KFSRequest::InferInputTensor requestTensor;
+    requestTensor.set_datatype("BYTES");
+    requestTensor.add_shape(2);
+    requestTensor.add_shape(5);
+    std::vector<std::string> strings = {"a", "bb", "ccc"};
+    std::string rawInputContents;
+    size_t dataSize = 0;
+    for (const auto& s : strings) {
+        dataSize += s.size() + sizeof(uint32_t);
+    }
+    rawInputContents.resize(dataSize);
+    size_t offset = 0;
+    for (const auto& s : strings) {
+        uint32_t inputSize = s.size();
+        std::memcpy(rawInputContents.data() + offset, &inputSize, sizeof(uint32_t));
+        offset += sizeof(uint32_t);
+        std::memcpy(rawInputContents.data() + offset, s.data(), s.size());
+        offset += s.size();
+    }
+    ov::Tensor tensor;
+    ASSERT_EQ(convertStringRequestToOVTensor(requestTensor, tensor, &rawInputContents), ovms::StatusCode::INVALID_STRING_INPUT);
+}
+
 template <typename TensorType>
 class StringOutputsConversionTest : public ::testing::Test {
 public: