diff --git a/src/kfs_frontend/kfs_utils.cpp b/src/kfs_frontend/kfs_utils.cpp index 246239d081..15bc97d1ab 100644 --- a/src/kfs_frontend/kfs_utils.cpp +++ b/src/kfs_frontend/kfs_utils.cpp @@ -16,6 +16,7 @@ #include "kfs_utils.hpp" #include +#include #include #include #include @@ -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(dim); + if (dimSize != 0 && shapeElements > std::numeric_limits::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 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(); size_t tensorStringsOffset = 0; for (size_t i = 0; i < stringSizes.size(); i++) { diff --git a/src/kfs_frontend/kfs_utils.hpp b/src/kfs_frontend/kfs_utils.hpp index 8faa5ce2ff..23fd16e8ca 100644 --- a/src/kfs_frontend/kfs_utils.hpp +++ b/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 diff --git a/src/kfs_frontend/tensor_conversion.cpp b/src/kfs_frontend/tensor_conversion.cpp index 138204e410..ae1726c76d 100644 --- a/src/kfs_frontend/tensor_conversion.cpp +++ b/src/kfs_frontend/tensor_conversion.cpp @@ -18,6 +18,7 @@ #include "../tensor_conversion_common.hpp" #include +#include #include #include #include @@ -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(numElements), shape); + if (!status.ok()) { + return status; + } + tensor = ov::Tensor(ov::element::Type_t::string, shape); + std::string* data = tensor.data(); + 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 diff --git a/src/predict_request_validation_utils.hpp b/src/predict_request_validation_utils.hpp index 99d6bb86d6..31e7eb8fb5 100644 --- a/src/predict_request_validation_utils.hpp +++ b/src/predict_request_validation_utils.hpp @@ -323,6 +323,9 @@ Status RequestValidator(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; diff --git a/src/predict_request_validation_utils_impl.hpp b/src/predict_request_validation_utils_impl.hpp index 14eff89771..2208650400 100644 --- a/src/predict_request_validation_utils_impl.hpp +++ b/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 diff --git a/src/tensor_conversion.hpp b/src/tensor_conversion.hpp index c5eb537062..2b4dc86175 100644 --- a/src/tensor_conversion.hpp +++ b/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 -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(batchSize)}); - std::string* data = tensor.data(); - 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 Status convertOVTensor2DToStringResponse(const ov::Tensor& tensor, TensorType& dst) { if (tensor.get_shape().size() != 2) { diff --git a/src/test/predict_validation_test.cpp b/src/test/predict_validation_test.cpp index 74d6cc873b..a4fc43ed91 100644 --- a/src/test/predict_validation_test.cpp +++ b/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( + 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(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( + tensorName, ovms::Precision::STRING, ovms::Shape{-1}, ovms::Layout{"N..."}); + + const int32_t N = static_cast(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(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; \ diff --git a/src/test/tensor_conversion_test.cpp b/src/test/tensor_conversion_test.cpp index d0581daff9..31006ca0fa 100644 --- a/src/test/tensor_conversion_test.cpp +++ b/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 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(); + 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 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(); + 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 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 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 class StringOutputsConversionTest : public ::testing::Test { public: