diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py index f40531b2e6..6f8ba94337 100755 --- a/qa/L0_shared_memory/shared_memory_test.py +++ b/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): 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}" 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..403aec9a6a 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 bde10b525c..058bd73c1d 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_);