Fix: Update handling of shared mem integer values (#8170)

Added an additional check to prevent the value of byte_size and offset used in a request from exceeding the bounds of shared memory.

Author: mattwittwer

qa/L0_shared_memory/shared_memory_test.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2019-2024, 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
@@ -385,6 +385,59 @@ def test_infer_byte_size_out_of_bound(self):
         )
         self._cleanup_shm_handles()
 
+    def test_infer_integer_overflow(self):
+        # Test for integer overflow vulnerability in offset + byte_size calculation
+        error_msg = []
+        self._configure_server()
+
+        offset = 32
+        byte_size = 2**64 - 32
+
+        if self.protocol == "http":
+            iu.shm_basic_infer(
+                self,
+                self.triton_client,
+                self._shm_handles[0],
+                self._shm_handles[1],
+                self._shm_handles[2],
+                self._shm_handles[3],
+                error_msg,
+                shm_output_offset=offset,
+                shm_output_byte_size=byte_size,
+                protocol=self.protocol,
+                use_system_shared_memory=True,
+            )
+
+            self.assertEqual(len(error_msg), 1)
+            self.assertTrue(
+                "Integer overflow detected: byte_size " in error_msg[0],
+                f"Unexpected error message: {error_msg[0]}",
+            )
+            self._cleanup_shm_handles()
+        else:
+            # The gRPC client utilizes the int64_param and will throw a separate error for values larger than 2**63-1
+            try:
+                iu.shm_basic_infer(
+                    self,
+                    self.triton_client,
+                    self._shm_handles[0],
+                    self._shm_handles[1],
+                    self._shm_handles[2],
+                    self._shm_handles[3],
+                    error_msg,
+                    shm_output_offset=offset,
+                    shm_output_byte_size=byte_size,
+                    protocol=self.protocol,
+                    use_system_shared_memory=True,
+                )
+                self.assertTrue(
+                    False,
+                    "Expected gRPC client to fail on value larger than int64_param maximum",
+                )
+            except ValueError as ex:
+                self.assertIn("Value out of range:", str(ex))
+            self._cleanup_shm_handles()
+
     def test_register_out_of_bound(self):
         create_byte_size = self.DEFAULT_SHM_BYTE_SIZE
 

qa/L0_shared_memory/test.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2019-2024, 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
@@ -53,6 +53,7 @@ for i in \
         test_unregisterall \
         test_infer_offset_out_of_bound \
         test_infer_byte_size_out_of_bound \
+        test_infer_integer_overflow \
         test_register_out_of_bound \
         test_python_client_leak; do
     for client_type in http grpc; do

src/shared_memory_manager.cc

@@ -1,4 +1,4 @@
-// Copyright 2019-2024, 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
@@ -485,6 +485,19 @@ SharedMemoryManager::GetMemoryInfo(
         std::string("Invalid offset for shared memory region: '" + name + "'")
             .c_str());
   }
+
+  // Check for potential integer overflow before validating bounds
+  if (byte_size > (SIZE_MAX - offset)) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        std::string(
+            "Integer overflow detected: byte_size (" +
+            std::to_string(byte_size) + ") + offset (" +
+            std::to_string(offset) + ") exceeds maximum value (" +
+            std::to_string(SIZE_MAX) + ") for region '" + name + "'")
+            .c_str());
+  }
+
   // validate byte_size + offset is within memory bounds
   size_t total_req_shm = offset + byte_size - 1;
   if (total_req_shm > shm_region_end) {