Skip to content

Commit 50ff906

Browse files
authored
fix: System shared memory boundary check (#8335)
1 parent a4d1897 commit 50ff906

File tree

2 files changed

+21
-14
lines changed

2 files changed

+21
-14
lines changed

qa/L0_shared_memory/shared_memory_test.py

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747

4848
class SystemSharedMemoryTestBase(tu.TestResultCollector):
4949
DEFAULT_SHM_BYTE_SIZE = 64
50+
SYS_PAGE_SIZE = os.sysconf("SC_PAGE_SIZE")
5051

5152
def setUp(self):
5253
self._setup_client()
@@ -305,11 +306,10 @@ def test_large_shm_register_offset(self):
305306

306307
# Test for large offset
307308
error_msg = []
308-
page_size = os.sysconf("SC_PAGE_SIZE")
309309
# Create a large shm size (page_size * 1024 is large enough to reproduce a segfault).
310310
# Register offset at 1 page before the end of the shm region to give enough space for the input/output data.
311-
create_byte_size = page_size * 1024
312-
register_offset = page_size * 1023
311+
create_byte_size = self.SYS_PAGE_SIZE * 1024
312+
register_offset = self.SYS_PAGE_SIZE * 1023
313313
self._configure_server(
314314
create_byte_size=create_byte_size,
315315
register_offset=register_offset,
@@ -372,7 +372,12 @@ def test_unregisterall(self):
372372
def test_infer_offset_out_of_bound(self):
373373
# Shared memory offset outside output region - Throws error
374374
error_msg = []
375-
self._configure_server()
375+
create_byte_size = self.SYS_PAGE_SIZE + self.DEFAULT_SHM_BYTE_SIZE
376+
register_offset = self.SYS_PAGE_SIZE
377+
self._configure_server(
378+
create_byte_size=create_byte_size,
379+
register_offset=register_offset,
380+
)
376381
if self.protocol == "http":
377382
# -32 when placed in an int64 signed type, to get a negative offset
378383
# by overflowing
@@ -402,8 +407,13 @@ def test_infer_offset_out_of_bound(self):
402407
def test_infer_byte_size_out_of_bound(self):
403408
# Shared memory byte_size outside output region - Throws error
404409
error_msg = []
405-
self._configure_server()
406-
offset = 60
410+
create_byte_size = self.SYS_PAGE_SIZE + self.DEFAULT_SHM_BYTE_SIZE
411+
register_offset = self.SYS_PAGE_SIZE
412+
self._configure_server(
413+
create_byte_size=create_byte_size,
414+
register_offset=register_offset,
415+
)
416+
offset = 1
407417
byte_size = self.DEFAULT_SHM_BYTE_SIZE
408418

409419
iu.shm_basic_infer(

src/shared_memory_manager.cc

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -483,14 +483,11 @@ SharedMemoryManager::GetMemoryInfo(
483483
}
484484

485485
// validate offset
486-
size_t shm_region_end = 0;
487-
if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) {
488-
shm_region_end = it->second->offset_;
489-
}
486+
size_t shm_region_size = 0;
490487
if (it->second->byte_size_ > 0) {
491-
shm_region_end += it->second->byte_size_ - 1;
488+
shm_region_size += it->second->byte_size_;
492489
}
493-
if (offset > shm_region_end) {
490+
if (offset >= shm_region_size) {
494491
return TRITONSERVER_ErrorNew(
495492
TRITONSERVER_ERROR_INVALID_ARG,
496493
std::string("Invalid offset for shared memory region: '" + name + "'")
@@ -510,8 +507,8 @@ SharedMemoryManager::GetMemoryInfo(
510507
}
511508

512509
// validate byte_size + offset is within memory bounds
513-
size_t total_req_shm = offset + byte_size - 1;
514-
if (total_req_shm > shm_region_end) {
510+
size_t total_req_shm = offset + byte_size;
511+
if (total_req_shm > shm_region_size) {
515512
return TRITONSERVER_ErrorNew(
516513
TRITONSERVER_ERROR_INVALID_ARG,
517514
std::string(

0 commit comments

Comments
 (0)