Improve KFS request validation

src/kfs_frontend/kfs_graph_executor_impl.cpp

@@ -922,13 +922,6 @@ static Status createPacketAndPushIntoGraph(const std::string& name, std::shared_
     }
     SPDLOG_DEBUG("Tensor to deserialize:\"{}\"", name);
     OVMS_RETURN_ON_FAIL(validateRequestCoherencyKFS(*request, request->model_name(), MediapipeGraphDefinition::VERSION));
-    if (!request->raw_input_contents().empty() && (request->raw_input_contents().size() != request->inputs().size())) {
-        std::stringstream ss;
-        ss << "Size of raw_input_contents: " << request->raw_input_contents().size() << " is different than number of inputs: " << request->inputs().size();
-        const std::string details = ss.str();
-        SPDLOG_DEBUG("[servable name: {} version: {}] Invalid message structure - {}", request->model_name(), request->model_version(), details);
-        return Status(StatusCode::INVALID_MESSAGE_STRUCTURE, details);
-    }
     std::unique_ptr<T> inputTensor;
     OVMS_RETURN_ON_FAIL(deserializeTensor(name, *request, inputTensor, pythonBackend));
     SPDLOG_TRACE("Current Timestamp before actual pushing:{}", timestamp.Value());

src/kfs_frontend/kfs_utils.cpp

@@ -24,11 +24,11 @@
 #include <unordered_map>
 #include <vector>
 
-#include "../logging.hpp"
-#include "../profiler.hpp"
-#include "../status.hpp"
-#include "../tensorinfo.hpp"
-#include "../tensor_conversion_common.hpp"
+#include "src/logging.hpp"
+#include "src/profiler.hpp"
+#include "src/status.hpp"
+#include "src/tensorinfo.hpp"
+#include "src/tensor_conversion_common.hpp"
 
 namespace ovms {
 Precision KFSPrecisionToOvmsPrecision(const KFSDataType& datatype) {
@@ -184,6 +184,13 @@ void setStringPrecision(KFSTensorOutputProto& proto) {
 
 Status validateRequestCoherencyKFS(const KFSRequest& request, const std::string servableName, model_version_t servableVersion) {
     if (!request.raw_input_contents().empty()) {
+        if (request.raw_input_contents().size() != request.inputs_size()) {
+            std::stringstream ss;
+            ss << "Size of raw_input_contents: " << request.raw_input_contents().size() << " is different than number of inputs: " << request.inputs_size();
+            const std::string details = ss.str();
+            SPDLOG_DEBUG("[servable name: {} version: {}] Invalid message structure - {}", servableName, servableVersion, details);
+            return Status(StatusCode::INVALID_MESSAGE_STRUCTURE, details);
+        }
         for (auto& input : request.inputs()) {
             if (input.has_contents()) {
                 std::stringstream ss;

src/test/mediapipe_validation_test.cpp

@@ -142,6 +142,12 @@ TEST_F(MediapipeValidationTest, NoDataInRawField) {
     ASSERT_EQ(impl->ModelInfer(nullptr, &request, &response).error_code(), grpc::StatusCode::INVALID_ARGUMENT);
 }
 
+TEST_F(MediapipeValidationTest, RawInputContentsCountMismatch) {
+    prepareDoubleInput();
+    request.mutable_raw_input_contents()->RemoveLast();
+    ASSERT_EQ(impl->ModelInfer(nullptr, &request, &response).error_code(), grpc::StatusCode::INVALID_ARGUMENT);
+}
+
 TEST_F(MediapipeValidationTest, NegativeShape) {
     prepareSingleInput();
     request.mutable_inputs(0)->mutable_shape()->Set(0, -1);

src/test/predict_validation_test.cpp

@@ -571,6 +571,25 @@ TEST_F(KFSPredictValidationInputTensorContent, RequestInputTensorContentAndRawIn
     EXPECT_EQ(status, ovms::StatusCode::INVALID_MESSAGE_STRUCTURE) << status.string();
 }
 
+TEST_F(KFSPredictValidationInputTensorContent, RawInputContentsCountMismatch) {
+    const std::string inputName1 = "input1";
+    const std::string inputName2 = "input2";
+    servableInputs = ovms::tensor_map_t({
+        {inputName1, std::make_shared<ovms::TensorInfo>(inputName1, ovms::Precision::FP32, ovms::shape_t{1, 4}, ovms::Layout{"NC"})},
+        {inputName2, std::make_shared<ovms::TensorInfo>(inputName2, ovms::Precision::FP32, ovms::shape_t{1, 4}, ovms::Layout{"NC"})},
+    });
+    ON_CALL(*instance, getInputsInfo()).WillByDefault(ReturnRef(servableInputs));
+    ON_CALL(*instance, getBatchSize()).WillByDefault(Return(1));
+    ON_CALL(*instance, getModelConfig()).WillByDefault(ReturnRef(modelConfig));
+    preparePredictRequest(request,
+        {{inputName1, std::tuple<ovms::signed_shape_t, ovms::Precision>{{1, 4}, ovms::Precision::FP32}},
+            {inputName2, std::tuple<ovms::signed_shape_t, ovms::Precision>{{1, 4}, ovms::Precision::FP32}}});
+    // Remove one buffer, leaving 1 raw_input_contents for 2 inputs
+    request.mutable_raw_input_contents()->RemoveLast();
+    auto status = instance->mockValidate(&request);
+    EXPECT_EQ(status, ovms::StatusCode::INVALID_MESSAGE_STRUCTURE) << status.string();
+}
+
 TEST_P(KFSPredictValidationInputTensorContent, RequestCorrectContentSizeInputTensorContent) {
     ovms::Precision testedPrecision = GetParam();
     const std::string inputName = "someName";