From 0f8db027a3f321d6ba49af83732506c2c858a85e Mon Sep 17 00:00:00 2001 From: Sai Kiran Polisetty Date: Fri, 8 Aug 2025 20:28:47 +0530 Subject: [PATCH 1/3] Update --- src/common.cc | 35 +++++++++++++++++++++++++++++++++++ src/common.h | 10 ++++++++++ src/shared_memory_manager.cc | 10 +--------- 3 files changed, 46 insertions(+), 9 deletions(-) diff --git a/src/common.cc b/src/common.cc index 90725041d1..a7591b8324 100644 --- a/src/common.cc +++ b/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 diff --git a/src/common.h b/src/common.h index 7df618e965..d411a99240 100644 --- a/src/common.h +++ b/src/common.h @@ -196,6 +196,16 @@ TRITONSERVER_Error* DecodeBase64( const char* input, size_t input_len, std::vector& 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'. /// diff --git a/src/shared_memory_manager.cc b/src/shared_memory_manager.cc index a2d52f5f48..06cb826dcc 100644 --- a/src/shared_memory_manager.cc +++ b/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 lock(mu_); From 7fed43b8cb0185d4081665a04c4ef7378243d905 Mon Sep 17 00:00:00 2001 From: Sai Kiran Polisetty Date: Fri, 8 Aug 2025 20:57:34 +0530 Subject: [PATCH 2/3] Test cases --- qa/L0_shared_memory/shared_memory_test.py | 43 +++++++++++++++++++---- qa/L0_shared_memory/test.sh | 1 + 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index 63f3700c05..d6d6030b79 100755 --- a/qa/L0_shared_memory/shared_memory_test.py +++ b/qa/L0_shared_memory/shared_memory_test.py @@ -556,16 +556,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): diff --git a/qa/L0_shared_memory/test.sh b/qa/L0_shared_memory/test.sh index d34ce1a4e5..3f5758bf36 100755 --- a/qa/L0_shared_memory/test.sh +++ b/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}" From 76b3743e0e6dc5820346d3da4048b6e4941d92b6 Mon Sep 17 00:00:00 2001 From: Sai Kiran Polisetty Date: Fri, 8 Aug 2025 21:08:15 +0530 Subject: [PATCH 3/3] Update --- src/common.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/common.h b/src/common.h index d411a99240..403aec9a6a 100644 --- a/src/common.h +++ b/src/common.h @@ -202,8 +202,8 @@ TRITONSERVER_Error* DecodeBase64( /// \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); +TRITONSERVER_Error* ValidateSharedMemoryKey( + const std::string& name, const std::string& shm_key); /// Joins container of strings into a single string delimited by