Fix: Update element count handling (#8182)

Ensures requests with overly large or invalid element counts are rejected with a relevant error message.

Author: mattwittwer

qa/L0_http/http_test.py

@@ -236,6 +236,104 @@ def test_descriptive_status_code(self):
         )
         t.join()
 
+    def test_buffer_size_overflow(self):
+        model = "onnx_zero_1_float32"
+
+        # Test for overflow within GetElementCount()
+        payload1 = {
+            "inputs": [
+                {
+                    "name": "INPUT0",
+                    "shape": [
+                        2**4,
+                        2**60 + 2,
+                    ],  # This evaluates to 2^64 + 32 during GetElementCount()
+                    "datatype": "FP32",
+                    "data": [1.0],
+                }
+            ]
+        }
+
+        # Test for overflow with type_byte_size multiplication
+        payload2 = {
+            "inputs": [
+                {
+                    "name": "INPUT0",
+                    "shape": [
+                        2**2,
+                        2**60 + 2,
+                    ],  # This evaluates to 2^64 + 32 during type_byte_size multiplication since FP32 is 4 bytes
+                    "datatype": "FP32",
+                    "data": [1.0],
+                }
+            ]
+        }
+
+        # Send request and expect a 400 error with specific overflow message
+        headers = {"Content-Type": "application/json"}
+
+        # Test the first payload (GetElementCount overflow)
+        r1 = requests.post(self._get_infer_url(model), json=payload1, headers=headers)
+
+        self.assertEqual(
+            400,
+            r1.status_code,
+            "Expected error code 400 for GetElementCount overflow check; got: {}".format(
+                r1.status_code
+            ),
+        )
+
+        error_message1 = r1.content.decode()
+        self.assertIn(
+            "causes total element count to exceed maximum size of", error_message1
+        )
+
+        # Test the second payload (type_byte_size multiplication overflow)
+        r2 = requests.post(self._get_infer_url(model), json=payload2, headers=headers)
+
+        self.assertEqual(
+            400,
+            r2.status_code,
+            "Expected error code 400 for type_byte_size multiplication overflow check; got: {}".format(
+                r2.status_code
+            ),
+        )
+
+        error_message2 = r2.content.decode()
+        self.assertIn("byte size overflow for input", error_message2)
+
+    def test_negative_dimensions(self):
+        model = "onnx_zero_1_float32"
+
+        payload = {
+            "inputs": [
+                {
+                    "name": "INPUT0",
+                    "shape": [2, -5],  # Negative dimension should be invalid
+                    "datatype": "FP32",
+                    "data": [1.0],
+                }
+            ]
+        }
+
+        # Send request and expect a 500 error
+        headers = {"Content-Type": "application/json"}
+        r = requests.post(self._get_infer_url(model), json=payload, headers=headers)
+
+        self.assertEqual(
+            500,
+            r.status_code,
+            "Expected error code 500 for negative dimension; got: {}".format(
+                r.status_code
+            ),
+        )
+
+        error_message = r.content.decode()
+        self.assertIn(
+            "Unable to parse 'shape': attempt to access JSON non-unsigned-integer as unsigned-integer",
+            error_message,
+        )
+
     def test_loading_large_invalid_model(self):
         # Generate large base64 encoded data
         data_length = 1 << 31

qa/L0_http/test.sh

@@ -630,7 +630,7 @@ fi
 
 TEST_RESULT_FILE='test_results.txt'
 PYTHON_TEST=http_test.py
-EXPECTED_NUM_TESTS=11
+EXPECTED_NUM_TESTS=13
 set +e
 python $PYTHON_TEST >$CLIENT_LOG 2>&1
 if [ $? -ne 0 ]; then

src/common.cc

@@ -1,4 +1,4 @@
-// Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2020-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -80,6 +80,31 @@ GetEnvironmentVariableOrDefault(
   return value ? value : default_value;
 }
 
+std::string
+ShapeToString(const int64_t* dims, const size_t dims_count)
+{
+  bool first = true;
+
+  std::string str("[");
+  for (size_t i = 0; i < dims_count; ++i) {
+    const int64_t dim = dims[i];
+    if (!first) {
+      str += ",";
+    }
+    str += std::to_string(dim);
+    first = false;
+  }
+
+  str += "]";
+  return str;
+}
+
+std::string
+ShapeToString(const std::vector<int64_t>& shape)
+{
+  return ShapeToString(shape.data(), shape.size());
+}
+
 int64_t
 GetElementCount(const std::vector<int64_t>& dims)
 {
@@ -88,12 +113,20 @@ GetElementCount(const std::vector<int64_t>& dims)
   for (auto dim : dims) {
     if (dim == WILDCARD_DIM) {
       return -1;
+    } else if (dim < 0) {  // invalid dim
+      return -2;
+    } else if (dim == 0) {
+      return 0;
     }
 
     if (first) {
       cnt = dim;
       first = false;
     } else {
+      // Check for overflow before multiplication
+      if (cnt > (INT64_MAX / dim)) {
+        return -3;
+      }
       cnt *= dim;
     }
   }

src/common.h

@@ -158,11 +158,18 @@ std::string GetEnvironmentVariableOrDefault(
 /// Get the number of elements in a shape.
 ///
 /// \param dims The shape.
-/// \return The number of elements, or -1 if the number of elements
+/// \return The number of elements, -1 if the number of elements
 /// cannot be determined because the shape contains one or more
-/// wildcard dimensions.
+/// wildcard dimensions, -2 if the shape contains an invalid dim,
+/// or -3 if the number is too large to represent as an int64_t.
 int64_t GetElementCount(const std::vector<int64_t>& dims);
 
+/// Convert shape to string representation.
+///
+/// \param shape The shape as a vector.
+/// \return The string representation of the shape.
+std::string ShapeToString(const std::vector<int64_t>& shape);
+
 /// Returns if 'vec' contains 'str'.
 ///
 /// \param vec The vector of strings to search.

src/http_server.cc

@@ -2755,12 +2755,29 @@ HTTPAPIServer::ParseJsonTritonIO(
       } else {
         const int64_t element_cnt = GetElementCount(shape_vec);
 
-        // FIXME, element count should never be 0 or negative so
-        // shouldn't we just return an error here?
         if (element_cnt == 0) {
           RETURN_IF_ERR(TRITONSERVER_InferenceRequestAppendInputData(
               irequest, input_name, nullptr, 0 /* byte_size */,
               TRITONSERVER_MEMORY_CPU, 0 /* memory_type_id */));
+        } else if (element_cnt == -2) {
+          // -2 indicates invalid dimension
+          return TRITONSERVER_ErrorNew(
+              TRITONSERVER_ERROR_INVALID_ARG,
+              std::string(
+                  "invalid shape for input '" + std::string(input_name) +
+                  "': shape " + ShapeToString(shape_vec) +
+                  " contains one or more invalid dimensions")
+                  .c_str());
+        } else if (element_cnt == -3) {
+          // -3 indicates integer overflow
+          return TRITONSERVER_ErrorNew(
+              TRITONSERVER_ERROR_INVALID_ARG,
+              std::string(
+                  "invalid shape for input '" + std::string(input_name) +
+                  "': shape " + ShapeToString(shape_vec) +
+                  " causes total element count to exceed maximum size of " +
+                  std::to_string(INT64_MAX))
+                  .c_str());
         } else {
           // JSON... presence of "data" already validated but still
           // checking here. Flow in this endpoint needs to be
@@ -2773,7 +2790,22 @@ HTTPAPIServer::ParseJsonTritonIO(
           if (dtype == TRITONSERVER_TYPE_BYTES) {
             RETURN_IF_ERR(JsonBytesArrayByteSize(tensor_data, &byte_size));
           } else {
-            byte_size = element_cnt * TRITONSERVER_DataTypeByteSize(dtype);
+            const uint32_t type_byte_size =
+                TRITONSERVER_DataTypeByteSize(dtype);
+            if ((type_byte_size > 1) &&
+                (element_cnt > (INT64_MAX / type_byte_size))) {
+              return TRITONSERVER_ErrorNew(
+                  TRITONSERVER_ERROR_INVALID_ARG,
+                  std::string(
+                      "byte size overflow for input '" +
+                      std::string(input_name) + "': element count (" +
+                      std::to_string(element_cnt) + ") * data type size (" +
+                      std::to_string(type_byte_size) +
+                      ") exceeds maximum allowed size (" +
+                      std::to_string(INT64_MAX) + ")")
+                      .c_str());
+            }
+            byte_size = element_cnt * type_byte_size;
           }
 
           infer_req->serialized_data_.emplace_back();