diff --git a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc index 25dee2d74..0658daedd 100644 --- a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc +++ b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.cc @@ -89,8 +89,9 @@ SharedMemoryManager::RegisterSystemMemory( Error SharedMemoryManager::GetMemoryInfo( - const std::string& name, size_t offset, void** shm_mapped_addr, - TRITONSERVER_MemoryType* memory_type, int64_t* device_id) + const std::string& name, size_t offset, size_t byte_size, + void** shm_mapped_addr, TRITONSERVER_MemoryType* memory_type, + int64_t* device_id) { // protect shared_memory_map_ from concurrent access std::lock_guard lock(mu_); @@ -100,6 +101,29 @@ SharedMemoryManager::GetMemoryInfo( return Error( std::string("Unable to find shared memory region: '" + name + "'")); } + + // validate offset + size_t shm_region_end = 0; + if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) { + shm_region_end = it->second->offset_; + } + if (it->second->byte_size_ > 0) { + shm_region_end += it->second->byte_size_ - 1; + } + if (offset > shm_region_end) { + return Error( + std::string("Invalid offset for shared memory 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) { + return Error(std::string( + "Invalid offset + byte size for shared memory region: '" + + name + "'") + .c_str()); + } + if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) { *shm_mapped_addr = (void*)((uint8_t*)it->second->mapped_addr_ + it->second->offset_ + offset); diff --git a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h index 1ffeb292c..6b2082c44 100644 --- a/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h +++ b/src/c++/perf_analyzer/client_backend/triton_c_api/shared_memory_manager.h @@ -70,10 +70,12 @@ class SharedMemoryManager { Error RegisterSystemMemory( const std::string& name, void* ptr, const size_t byte_size); - /// Get the access information for the shared memory block with the specified - /// name. Return an Error if named block doesn't exist. + /// Get the access information for the shared memory block + /// with the specified name. Return an Error + /// if named block doesn't exist. /// \param name The name of the shared memory block to get. /// \param offset The offset in the block + /// \param byte_size The byte size to request for the shm region /// \param shm_mapped_addr Returns the pointer to the shared /// memory block with the specified name and offset /// \param memory_type Returns the type of the memory @@ -81,8 +83,9 @@ class SharedMemoryManager { /// memory block /// \return an Error indicating success or failure. Error GetMemoryInfo( - const std::string& name, size_t offset, void** shm_mapped_addr, - TRITONSERVER_MemoryType* memory_type, int64_t* device_id); + const std::string& name, size_t offset, size_t byte_size, + void** shm_mapped_addr, TRITONSERVER_MemoryType* memory_type, + int64_t* device_id); /// Removes the named shared memory block of the specified type from /// the manager. Any future attempt to get the details of this block diff --git a/src/c++/perf_analyzer/client_backend/triton_c_api/triton_loader.cc b/src/c++/perf_analyzer/client_backend/triton_c_api/triton_loader.cc index 12c767c3c..35f7657f3 100644 --- a/src/c++/perf_analyzer/client_backend/triton_c_api/triton_loader.cc +++ b/src/c++/perf_analyzer/client_backend/triton_c_api/triton_loader.cc @@ -945,7 +945,8 @@ TritonLoader::Infer( TRITONSERVER_MemoryType memory_type; int64_t memory_type_id; RETURN_IF_ERROR(shm_manager_->GetMemoryInfo( - shm_name, offset, &buf, &memory_type, &memory_type_id)); + shm_name, offset, shm_byte_size, &buf, &memory_type, + &memory_type_id)); alloc_payload.output_map_.emplace( std::piecewise_construct, std::forward_as_tuple(output->Name()), @@ -1149,7 +1150,8 @@ TritonLoader::AddInputs( TRITONSERVER_MemoryType memory_type; int64_t memory_type_id; RETURN_IF_ERROR(shm_manager_->GetMemoryInfo( - shm_name, offset, &buf, &memory_type, &memory_type_id)); + shm_name, offset, shm_byte_size, &buf, &memory_type, + &memory_type_id)); RETURN_IF_TRITONSERVER_ERROR( inference_request_append_input_data_fn_( irequest, input_name, buf, byte_size,