refactor: Prevent overflow in backend APIs

include/triton/backend/backend_common.h

@@ -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
@@ -316,28 +316,66 @@ std::string ShapeToString(const int64_t* dims, const size_t dims_count);
 /// \return The string representation.
 std::string ShapeToString(const std::vector<int64_t>& shape);
 
+/// Deprecated. Use TRITONSERVER_Error* GetElementCount instead.
 /// Return the number of elements of a shape.
 ///
 /// \param dims The shape dimensions.
 /// \param dims_count The number of dimensions.
-/// \return The number of elements.
+/// \return The number of elements,
+/// -1 if unable to determine the number,
+/// -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 int64_t* dims, const size_t dims_count);
 
+/// Deprecated. Use TRITONSERVER_Error* GetElementCount instead.
 /// Return the number of elements of a shape.
 ///
 /// \param shape The shape as a vector of dimensions.
-/// \return The number of elements.
+/// \return The number of elements,
+/// -1 if unable to determine the number,
+/// -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>& shape);
 
+/// Return the number of elements of a shape with error checking.
+///
+/// \param dims The shape dimensions.
+/// \param dims_count The number of dimensions.
+/// \param cnt Returns the number of elements.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_Error* GetElementCount(
+    const int64_t* dims, const size_t dims_count, int64_t* cnt);
+
+/// Return the number of elements of a shape with error checking.
+///
+/// \param shape The shape as a vector of dimensions.
+/// \param cnt Returns the number of elements.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_Error* GetElementCount(
+    const std::vector<int64_t>& shape, int64_t* cnt);
+
+/// Deprecated. Use TRITONSERVER_Error* GetByteSize instead.
 /// Get the size, in bytes, of a tensor based on datatype and
 /// shape.
 /// \param dtype The data-type.
 /// \param dims The shape.
-/// \return The size, in bytes, of the corresponding tensor, or -1 if
-/// unable to determine the size.
+/// \return The size, in bytes, of the corresponding tensor,
+/// -1 if unable to determine the size,
+/// -2 if the shape contains an invalid dim,
+/// or -3 if the size is too large to represent as an int64_t.
 int64_t GetByteSize(
     const TRITONSERVER_DataType& dtype, const std::vector<int64_t>& dims);
 
+/// Get the size, in bytes, of a tensor based on datatype and
+/// shape with error checking.
+/// \param dtype The data-type.
+/// \param dims The shape.
+/// \param size Returns the size, in bytes, of the corresponding tensor.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_Error* GetByteSize(
+    const TRITONSERVER_DataType& dtype, const std::vector<int64_t>& dims,
+    int64_t* size);
+
 /// Get an input tensor's contents into a buffer. This overload expects
 /// both 'buffer' and buffers of the input to be in CPU.
 ///

src/backend_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
@@ -166,12 +166,20 @@ GetElementCount(const int64_t* dims, const size_t dims_count)
   for (size_t i = 0; i < dims_count; i++) {
     if (dims[i] == WILDCARD_DIM) {
       return -1;
+    } else if (dims[i] < 0) {  // invalid dim
+      return -2;
+    } else if (dims[i] == 0) {
+      return 0;
     }
 
     if (first) {
       cnt = dims[i];
       first = false;
     } else {
+      // Check for overflow before multiplication
+      if (cnt > INT64_MAX / dims[i]) {
+        return -3;
+      }
       cnt *= dims[i];
     }
   }
@@ -185,6 +193,42 @@ GetElementCount(const std::vector<int64_t>& shape)
   return GetElementCount(shape.data(), shape.size());
 }
 
+TRITONSERVER_Error*
+GetElementCount(const int64_t* dims, const size_t dims_count, int64_t* cnt)
+{
+  *cnt = GetElementCount(dims, dims_count);
+  if (*cnt == -2) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        (std::string("shape") + ShapeToString(dims, dims_count) +
+         " contains an invalid dim.")
+            .c_str());
+  } else if (*cnt == -3) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        "unexpected integer overflow while calculating element count.");
+  }
+  return nullptr;  // success
+}
+
+TRITONSERVER_Error*
+GetElementCount(const std::vector<int64_t>& shape, int64_t* cnt)
+{
+  *cnt = GetElementCount(shape.data(), shape.size());
+  if (*cnt == -2) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        (std::string("shape") + ShapeToString(shape) +
+         " contains an invalid dim.")
+            .c_str());
+  } else if (*cnt == -3) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        "unexpected integer overflow while calculating element count.");
+  }
+  return nullptr;  // success
+}
+
 int64_t
 GetByteSize(
     const TRITONSERVER_DataType& dtype, const std::vector<int64_t>& dims)
@@ -195,13 +239,36 @@ GetByteSize(
   }
 
   int64_t cnt = GetElementCount(dims);
-  if (cnt == -1) {
-    return -1;
+  if (cnt <= 0) {
+    return cnt;
   }
 
+  if ((cnt > INT64_MAX / dt_size)) {
+    return -3;
+  }
   return cnt * dt_size;
 }
 
+TRITONSERVER_Error*
+GetByteSize(
+    const TRITONSERVER_DataType& dtype, const std::vector<int64_t>& dims,
+    int64_t* size)
+{
+  *size = GetByteSize(dtype, dims);
+  if (*size == -2) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        (std::string("shape") + ShapeToString(dims) +
+         " contains an invalid dim.")
+            .c_str());
+  } else if (*size == -3) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        "unexpected integer overflow while calculating byte size.");
+  }
+  return nullptr;  // success
+}
+
 TRITONSERVER_Error*
 ReadInputTensor(
     TRITONBACKEND_Request* request, const std::string& input_name, char* buffer,

src/backend_input_collector.cc

@@ -1,4 +1,4 @@
-// Copyright 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2019-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
@@ -762,11 +762,12 @@ BackendInputCollector::BatchInputShape(
             requests_[req_idx], source_input.c_str(), &input));
         const int64_t* shape_arr;
         uint32_t dims_count;
+        int64_t element_cnt = 0;
         RETURN_IF_ERROR(TRITONBACKEND_InputPropertiesForHostPolicy(
             input, host_policy_cstr_, nullptr, nullptr, &shape_arr, &dims_count,
             nullptr, nullptr));
-        (*shape)[0] =
-            std::max((*shape)[0], GetElementCount(shape_arr, dims_count));
+        RETURN_IF_ERROR(GetElementCount(shape_arr, dims_count, &element_cnt));
+        (*shape)[0] = std::max((*shape)[0], element_cnt);
       }
       break;
     }
@@ -841,7 +842,9 @@ BackendInputCollector::ProcessBatchInput(
     // Calculate the byte size of the buffer
     std::vector<int64_t> shape;
     RETURN_IF_ERROR(BatchInputShape(batch_input, &shape));
-    *dst_buffer_byte_size = GetByteSize(batch_input.DataType(), shape);
+    RETURN_IF_ERROR(GetByteSize(
+        batch_input.DataType(), shape,
+        reinterpret_cast<int64_t*>(dst_buffer_byte_size)));
     BackendMemory* backend_memory = nullptr;
     for (const auto& allowed_type : allowed_input_types) {
       std::vector<BackendMemory::AllocationType> alloc_types;
@@ -945,11 +948,31 @@ BackendInputCollector::ProcessBatchInput(
       const auto& source_input = batch_input.SourceInputs()[0];
       if (data_type == TRITONSERVER_TYPE_FP32) {
         *reinterpret_cast<float*>(input_buffer) = 0;
+        if (*dst_buffer_byte_size < sizeof(float)) {
+          return TRITONSERVER_ErrorNew(
+              TRITONSERVER_ERROR_INVALID_ARG,
+              (std::string(
+                   "Unexpected total byte size for batch input. Expect >= ") +
+               std::to_string(sizeof(float)) + ", got " +
+               std::to_string(*dst_buffer_byte_size))
+                  .c_str());
+        }
+
         RETURN_IF_ERROR(SetAccumulatedElementCount<float>(
             source_input, input_buffer + sizeof(float),
             *dst_buffer_byte_size - sizeof(float)));
       } else {
         *reinterpret_cast<int32_t*>(input_buffer) = 0;
+        if (*dst_buffer_byte_size < sizeof(int32_t)) {
+          return TRITONSERVER_ErrorNew(
+              TRITONSERVER_ERROR_INVALID_ARG,
+              (std::string(
+                   "Unexpected total byte size for batch input. Expect >= ") +
+               std::to_string(sizeof(int32_t)) + ", got " +
+               std::to_string(*dst_buffer_byte_size))
+                  .c_str());
+        }
+
         RETURN_IF_ERROR(SetAccumulatedElementCount<int32_t>(
             source_input, input_buffer + sizeof(int32_t),
             *dst_buffer_byte_size - sizeof(int32_t)));
@@ -1011,11 +1034,12 @@ BackendInputCollector::SetElementCount(
         requests_[req_idx], source_input.c_str(), &input));
     const int64_t* shape;
     uint32_t dims_count;
+    int64_t element_cnt = 0;
     RETURN_IF_ERROR(TRITONBACKEND_InputPropertiesForHostPolicy(
         input, host_policy_cstr_, nullptr, nullptr, &shape, &dims_count,
         nullptr, nullptr));
-    *(reinterpret_cast<T*>(buffer) + req_idx) =
-        GetElementCount(shape, dims_count);
+    RETURN_IF_ERROR(GetElementCount(shape, dims_count, &element_cnt));
+    *(reinterpret_cast<T*>(buffer) + req_idx) = element_cnt;
     buffer_offset += sizeof(T);
   }
   // Set the rest of the buffer to 0
@@ -1046,10 +1070,12 @@ BackendInputCollector::SetAccumulatedElementCount(
         requests_[req_idx], source_input.c_str(), &input));
     const int64_t* shape;
     uint32_t dims_count;
+    int64_t element_cnt = 0;
     RETURN_IF_ERROR(TRITONBACKEND_InputPropertiesForHostPolicy(
         input, host_policy_cstr_, nullptr, nullptr, &shape, &dims_count,
         nullptr, nullptr));
-    accumulated_element_count += GetElementCount(shape, dims_count);
+    RETURN_IF_ERROR(GetElementCount(shape, dims_count, &element_cnt));
+    accumulated_element_count += element_cnt;
     *(reinterpret_cast<T*>(buffer) + req_idx) = accumulated_element_count;
     buffer_offset += sizeof(T);
   }

src/backend_output_responder.cc

@@ -1,4 +1,4 @@
-// Copyright 2019-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2019-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
@@ -109,7 +109,9 @@ BackendOutputResponder::ProcessTensor(
       batch_size_offset += shape[0];
     }
 
-    const size_t tensor_byte_size = GetByteSize(datatype, batchn_shape);
+    int64_t tensor_byte_size = 0;
+    RESPOND_AND_SET_NULL_IF_ERROR(
+        &response, GetByteSize(datatype, batchn_shape, &tensor_byte_size));
 
     TRITONBACKEND_Output* response_output;
     if (response != nullptr) {
@@ -218,7 +220,9 @@ BackendOutputResponder::ProcessStateTensor(
       batch_size_offset += shape[0];
     }
 
-    const size_t tensor_byte_size = GetByteSize(datatype, batchn_shape);
+    int64_t tensor_byte_size = 0;
+    RESPOND_AND_SET_NULL_IF_ERROR(
+        &response, GetByteSize(datatype, batchn_shape, &tensor_byte_size));
 
     TRITONBACKEND_State* output_state;
     if (response != nullptr) {
@@ -554,8 +558,10 @@ BackendOutputResponder::ProcessBatchOutput(
           }
         }
 
-        const size_t tensor_byte_size =
-            GetByteSize(datatype, output_batchn_shape);
+        int64_t tensor_byte_size = 0;
+        RESPOND_AND_SET_NULL_IF_ERROR(
+            &response,
+            GetByteSize(datatype, output_batchn_shape, &tensor_byte_size));
 
         TRITONBACKEND_Output* response_output;
         if (response != nullptr) {