From ff8f4d9c31956738bd65b1046202ffbde377957e Mon Sep 17 00:00:00 2001 From: LinZiyuu Date: Sun, 31 May 2026 06:13:35 +0800 Subject: [PATCH] fix: validate tensor dims_count against the shared-memory region PbTensor::LoadFromSharedMemory() reads `dims_count` from shared memory and uses it to compute `name_offset` and to construct the dims vector (`std::vector(dims_ptr, dims_ptr + dims_count)`) without checking it against the region. A corrupted `dims_count` (e.g. written by a model into the backend shm) makes the parent process perform a large out-of-bounds read, crashing the server; the `sizeof(int64_t) * dims_count` product can also overflow and yield a small, controlled `name_offset`. Validate `dims_count` so that the dims array stays within the shared-memory region before it is used, throwing PythonBackendException otherwise. The check uses division to avoid overflowing the product, and mirrors the MemoryShm::byte_size boundary check. Valid tensors are unaffected. Signed-off-by: LinZiyuu --- src/pb_tensor.cc | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/pb_tensor.cc b/src/pb_tensor.cc index 02c37845..c04ed5e3 100644 --- a/src/pb_tensor.cc +++ b/src/pb_tensor.cc @@ -655,6 +655,22 @@ PbTensor::LoadFromSharedMemory( AllocatedSharedMemory tensor_shm = shm_pool->Load(tensor_handle); TensorShm* tensor_shm_ptr = reinterpret_cast(tensor_shm.data_.get()); + + // Validate 'dims_count' against the shared-memory region before using it to + // compute offsets and read the dims array, to avoid an out-of-bounds read + // when the value is corrupted. The division avoids overflowing the + // 'sizeof(int64_t) * dims_count' product. Mirrors the MemoryShm::byte_size + // boundary check. + char* region_end = reinterpret_cast(shm_pool->GetBaseAddress()) + + shm_pool->GetCurrentCapacity(); + char* dims_begin = tensor_shm.data_.get() + sizeof(TensorShm); + if (dims_begin > region_end || + tensor_shm_ptr->dims_count > + static_cast(region_end - dims_begin) / sizeof(int64_t)) { + throw PythonBackendException( + "Attempted to load a tensor with an out-of-bounds 'dims_count'"); + } + size_t name_offset = sizeof(TensorShm) + sizeof(int64_t) * tensor_shm_ptr->dims_count; std::unique_ptr name_shm = PbString::LoadFromSharedMemory(