fix: Improve validation for system shared memory register (#8273) (#8310)

Co-authored-by: Sai Kiran Polisetty <spolisetty@nvidia.com>

Author: mc-nv

qa/L0_shared_memory/shared_memory_test.py

@@ -509,6 +509,24 @@ def test_python_client_leak(self):
             "client memory usage is increasing",
         )
 
+    def test_register_reserved_names(self):
+        """
+        Test that registration fails if attempting to use a reserved
+        prefix for the shm key.
+        """
+        # This matches kTritonSharedMemoryRegionPrefix in the server code.
+        reserved_prefix = "triton_python_backend_shm_region_"
+
+        # The shared memory key cannot start with the reserved prefix.
+        shm_name = "my_test_shm_name"
+        shm_key = f"{reserved_prefix}_my_test_shm_key"
+
+        with self.assertRaisesRegex(
+            utils.InferenceServerException,
+            f"cannot register shared memory region '{shm_name}' with key '{shm_key}' as the key contains the reserved prefix '{reserved_prefix}'",
+        ) as e:
+            self.triton_client.register_system_shared_memory(shm_name, shm_key, 10000)
+
 
 def callback(user_data, result, error):
     if error:

qa/L0_shared_memory/test.sh

@@ -55,6 +55,7 @@ for i in \
         test_infer_byte_size_out_of_bound \
         test_infer_integer_overflow \
         test_register_out_of_bound \
+        test_register_reserved_names \
         test_python_client_leak; do
     for client_type in http grpc; do
         SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=1 ${SERVER_ARGS_EXTRA}"

src/common.h

@@ -45,6 +45,10 @@ constexpr char kContentEncodingHTTPHeader[] = "Content-Encoding";
 constexpr char kContentTypeHeader[] = "Content-Type";
 constexpr char kContentLengthHeader[] = "Content-Length";
 
+// This prefix is reserved for shm regions created internally by Triton
+constexpr char kTritonSharedMemoryRegionPrefix[] =
+    "triton_python_backend_shm_region_";
+
 constexpr int MAX_GRPC_MESSAGE_SIZE = INT32_MAX;
 
 /// The value for a dimension in a shape that indicates that that

src/shared_memory_manager.cc

@@ -354,6 +354,17 @@ SharedMemoryManager::RegisterSystemSharedMemory(
     const std::string& name, const std::string& shm_key, const size_t offset,
     const size_t byte_size)
 {
+  // Check if the shared memory key starts with the reserved prefix
+  if (shm_key.rfind(kTritonSharedMemoryRegionPrefix, 0) == 0) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        std::string(
+            "cannot register shared memory region '" + name + "' with key '" +
+            shm_key + "' as the key contains the reserved prefix '" +
+            kTritonSharedMemoryRegionPrefix + "'")
+            .c_str());
+  }
+
   std::lock_guard<std::mutex> lock(mu_);
 
   if (shared_memory_map_.find(name) != shared_memory_map_.end()) {
@@ -670,7 +681,6 @@ SharedMemoryManager::UnregisterAll(TRITONSERVER_MemoryType memory_type)
       ++next_it;
       if (it->second->kind_ == TRITONSERVER_MEMORY_GPU) {
         TRITONSERVER_Error* err = UnregisterHelper(it->first, memory_type);
-        ;
         if (err != nullptr) {
           unregister_fails.push_back(it->first);
           LOG_VERBOSE(1) << TRITONSERVER_ErrorMessage(err);