fix: Improve validation for system shared memory register (#8336) (#8340)

Author: pskiran1

qa/L0_shared_memory/shared_memory_test.py

@@ -566,16 +566,45 @@ def test_register_reserved_names(self):
         """
         # This matches kTritonSharedMemoryRegionPrefix in the server code.
         reserved_prefix = "triton_python_backend_shm_region_"
+        shm_name = "my_test_shm_name"
+
+        # The shared memory key cannot start with the reserved prefix,
+        # regardless of leading slashes.
+        shm_keys_to_test = [
+            f"{reserved_prefix}_my_test_shm_key",
+            f"/{reserved_prefix}_my_test_shm_key",
+            f"///{reserved_prefix}_my_test_shm_key",
+        ]
 
-        # The shared memory key cannot start with the reserved prefix.
+        for shm_key in shm_keys_to_test:
+            with self.subTest(shm_key=shm_key):
+                expected_msg = f"cannot register shared memory region '{shm_name}' with key '{shm_key}' as the key contains the reserved prefix '{reserved_prefix}'"
+                with self.assertRaisesRegex(
+                    utils.InferenceServerException, expected_msg
+                ):
+                    self.triton_client.register_system_shared_memory(
+                        shm_name, shm_key, 10000
+                    )
+
+    def test_register_invalid_shm_key(self):
+        """
+        Test that registration fails if attempting to use an invalid name for the shm key.
+        """
         shm_name = "my_test_shm_name"
-        shm_key = f"{reserved_prefix}_my_test_shm_key"
+        shm_keys_to_test = [
+            "/",
+            "///",
+        ]
 
-        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)
+        for shm_key in shm_keys_to_test:
+            with self.subTest(shm_key=shm_key):
+                expected_msg = f"cannot register shared memory region '{shm_name}' - invalid shm key '{shm_key}'"
+                with self.assertRaisesRegex(
+                    utils.InferenceServerException, expected_msg
+                ):
+                    self.triton_client.register_system_shared_memory(
+                        shm_name, shm_key, 10000
+                    )
 
 
 def callback(user_data, result, error):

qa/L0_shared_memory/test.sh

@@ -70,6 +70,7 @@ for i in \
         test_infer_integer_overflow \
         test_register_out_of_bound \
         test_register_reserved_names \
+        test_register_invalid_shm_key \
         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.cc

@@ -163,4 +163,39 @@ DecodeBase64(
   return nullptr;
 }
 
+TRITONSERVER_Error*
+ValidateSharedMemoryKey(const std::string& name, const std::string& shm_key)
+{
+  std::string_view key_view(shm_key);
+
+  // Find the index of the first character that is not a slash
+  const std::size_t first_non_slash = key_view.find_first_not_of('/');
+
+  // If the entire key is slashes
+  if (first_non_slash == std::string_view::npos) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        std::string(
+            "cannot register shared memory region '" + name +
+            "' - invalid shm key '" + shm_key + "'")
+            .c_str());
+  }
+
+  // Check whether the substring starting at first_non_slash starts with the
+  // reserved prefix
+  if (key_view.substr(first_non_slash)
+          .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());
+  }
+
+  // Valid shm key
+  return nullptr;
+}
+
 }}  // namespace triton::server

src/common.h

@@ -196,6 +196,16 @@ TRITONSERVER_Error* DecodeBase64(
     const char* input, size_t input_len, std::vector<char>& decoded_data,
     size_t& decoded_size, const std::string& name);
 
+
+/// Validate shared memory key
+///
+/// \param name The name of the memory block.
+/// \param shm_key The name of the posix shared memory object
+/// \return The error status.
+TRITONSERVER_Error* ValidateSharedMemoryKey(
+    const std::string& name, const std::string& shm_key);
+
+
 /// Joins container of strings into a single string delimited by
 /// 'delim'.
 ///

src/shared_memory_manager.cc

@@ -355,15 +355,7 @@ SharedMemoryManager::RegisterSystemSharedMemory(
     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());
-  }
+  RETURN_IF_ERR(ValidateSharedMemoryKey(name, shm_key));
 
   std::lock_guard<std::mutex> lock(mu_);