Skip to content

Commit a4d1897

Browse files
authored
fix: Proper handling of system shm register offset (#8330)
1 parent 5a97991 commit a4d1897

File tree

4 files changed

+143
-19
lines changed

4 files changed

+143
-19
lines changed

qa/L0_shared_memory/shared_memory_test.py

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -113,10 +113,6 @@ def _configure_server(
113113
shm_op1_handle,
114114
]
115115
# Implicit assumption that input and output byte_sizes are 64 bytes for now
116-
input0_data = np.arange(start=0, stop=16, dtype=np.int32)
117-
input1_data = np.ones(shape=16, dtype=np.int32)
118-
shm.set_shared_memory_region(shm_ip0_handle, [input0_data])
119-
shm.set_shared_memory_region(shm_ip1_handle, [input1_data])
120116
self.triton_client.register_system_shared_memory(
121117
"input0_data", "/input0_data", register_byte_size, offset=register_offset
122118
)
@@ -129,6 +125,16 @@ def _configure_server(
129125
self.triton_client.register_system_shared_memory(
130126
"output1_data", "/output1_data", register_byte_size, offset=register_offset
131127
)
128+
129+
# Write data to shared memory regions
130+
input0_data = np.arange(start=0, stop=16, dtype=np.int32)
131+
input1_data = np.ones(shape=16, dtype=np.int32)
132+
shm.set_shared_memory_region(
133+
shm_ip0_handle, [input0_data], offset=register_offset
134+
)
135+
shm.set_shared_memory_region(
136+
shm_ip1_handle, [input1_data], offset=register_offset
137+
)
132138
self.shm_names = ["input0_data", "input1_data", "output0_data", "output1_data"]
133139

134140
def _cleanup_shm_handles(self):
@@ -292,6 +298,40 @@ def test_too_big_shm(self):
292298
self._shm_handles.append(shm_ip2_handle)
293299
self._cleanup_shm_handles()
294300

301+
def test_large_shm_register_offset(self):
302+
# Test for out of bounds read vulnerability when registering system shared memory with large offset
303+
for platform in ["python", "onnx", "libtorch", "plan", "openvino"]:
304+
model_name = f"{platform}_int32_int32_int32"
305+
306+
# Test for large offset
307+
error_msg = []
308+
page_size = os.sysconf("SC_PAGE_SIZE")
309+
# Create a large shm size (page_size * 1024 is large enough to reproduce a segfault).
310+
# 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
313+
self._configure_server(
314+
create_byte_size=create_byte_size,
315+
register_offset=register_offset,
316+
)
317+
318+
iu.shm_basic_infer(
319+
self,
320+
self.triton_client,
321+
self._shm_handles[0],
322+
self._shm_handles[1],
323+
self._shm_handles[2],
324+
self._shm_handles[3],
325+
error_msg,
326+
register_offset=register_offset,
327+
protocol=self.protocol,
328+
use_system_shared_memory=True,
329+
override_model_name=model_name,
330+
)
331+
self.triton_client.unregister_system_shared_memory()
332+
if len(error_msg) > 0:
333+
raise Exception(str(error_msg))
334+
295335
def test_mixed_raw_shm(self):
296336
# Mix of shared memory and RAW inputs
297337
error_msg = []

qa/L0_shared_memory/test.sh

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,26 @@
2525
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
2626
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
2727

28+
REPO_VERSION=${NVIDIA_TRITON_SERVER_VERSION}
29+
if [ "$#" -ge 1 ]; then
30+
REPO_VERSION=$1
31+
fi
32+
if [ -z "$REPO_VERSION" ]; then
33+
echo -e "Repository version must be specified"
34+
echo -e "\n***\n*** Test Failed\n***"
35+
exit 1
36+
fi
37+
if [ ! -z "$TEST_REPO_ARCH" ]; then
38+
REPO_VERSION=${REPO_VERSION}_${TEST_REPO_ARCH}
39+
fi
40+
2841
CLIENT_LOG="./client.log"
2942
SHM_TEST=shared_memory_test.py
3043
TEST_RESULT_FILE='test_results.txt'
3144

3245
# Configure to support test on jetson as well
3346
TRITON_DIR=${TRITON_DIR:="/opt/tritonserver"}
47+
DATADIR=/data/inferenceserver/${REPO_VERSION}
3448
SERVER=${TRITON_DIR}/bin/tritonserver
3549
BACKEND_DIR=${TRITON_DIR}/backends
3650
SERVER_ARGS_EXTRA="--backend-directory=${BACKEND_DIR}"
@@ -142,6 +156,60 @@ for test_case in \
142156
done
143157
done
144158

159+
# Test large system shared memory offset
160+
rm -rf models/*
161+
# prepare add_sub model of various backends
162+
BACKENDS="python onnx libtorch plan openvino"
163+
for backend in ${BACKENDS} ; do
164+
model="${backend}_int32_int32_int32"
165+
model_dir="models/${model}"
166+
if [[ $backend == "python" ]]; then
167+
mkdir -p ${model_dir}/1
168+
cp ../python_models/add_sub/model.py ${model_dir}/1/
169+
cp ../python_models/add_sub/config.pbtxt ${model_dir}/
170+
sed -i 's/TYPE_FP32/TYPE_INT32/g' ${model_dir}/config.pbtxt
171+
echo "max_batch_size: 8" >> ${model_dir}/config.pbtxt
172+
else
173+
mkdir -p ${model_dir}
174+
cp -r $DATADIR/qa_model_repository/${model}/1 ${model_dir}/1
175+
cp $DATADIR/qa_model_repository/${model}/config.pbtxt ${model_dir}/
176+
cp $DATADIR/qa_model_repository/${model}/output0_labels.txt ${model_dir}/
177+
if [ $backend == "openvino" ]; then
178+
echo 'parameters { key: "ENABLE_BATCH_PADDING" value { string_value: "YES" } }' >> models/${model}/config.pbtxt
179+
fi
180+
fi
181+
done
182+
183+
test_case="test_large_shm_register_offset"
184+
for client_type in http grpc; do
185+
SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=1 ${SERVER_ARGS_EXTRA}"
186+
SERVER_LOG="./${test_case}.${client_type}.server.log"
187+
run_server
188+
if [ "$SERVER_PID" == "0" ]; then
189+
echo -e "\n***\n*** Failed to start $SERVER\n***"
190+
cat $SERVER_LOG
191+
exit 1
192+
fi
193+
194+
export CLIENT_TYPE=$client_type
195+
CLIENT_LOG="./${test_case}.${client_type}.client.log"
196+
set +e
197+
python3 $SHM_TEST SharedMemoryTest.${test_case} >>"$CLIENT_LOG" 2>&1
198+
if [ $? -ne 0 ]; then
199+
cat $CLIENT_LOG
200+
echo -e "\n***\n*** Test Failed - ${client_type}\n***"
201+
RET=1
202+
fi
203+
204+
kill $SERVER_PID
205+
wait $SERVER_PID
206+
if [ $? -ne 0 ]; then
207+
echo -e "\n***\n*** Test Server shut down non-gracefully\n***"
208+
RET=1
209+
fi
210+
set -e
211+
done
212+
145213
if [ $RET -eq 0 ]; then
146214
echo -e "\n***\n*** Test Passed\n***"
147215
else

qa/common/infer_util.py

Lines changed: 30 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#!/usr/bin/env python3
22

3-
# Copyright 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
# Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
44
#
55
# Redistribution and use in source and binary forms, with or without
66
# modification, are permitted provided that the following conditions
@@ -1367,11 +1367,13 @@ def shm_basic_infer(
13671367
big_shm_name="",
13681368
big_shm_size=64,
13691369
default_shm_byte_size=64,
1370+
register_offset=0,
13701371
shm_output_offset=0,
13711372
shm_output_byte_size=64,
13721373
protocol="http",
13731374
use_system_shared_memory=False,
13741375
use_cuda_shared_memory=False,
1376+
override_model_name=None,
13751377
):
13761378
# Lazy shm imports...
13771379
if use_system_shared_memory:
@@ -1381,20 +1383,34 @@ def shm_basic_infer(
13811383
else:
13821384
raise Exception("No shared memory type specified")
13831385

1386+
if override_model_name is None:
1387+
model_name = "simple"
1388+
else:
1389+
model_name = override_model_name
1390+
1391+
if model_name.startswith("libtorch"):
1392+
output_names = ["OUTPUT__0", "OUTPUT__1"]
1393+
else:
1394+
output_names = ["OUTPUT0", "OUTPUT1"]
1395+
13841396
input0_data = np.arange(start=0, stop=16, dtype=np.int32)
13851397
input1_data = np.ones(shape=16, dtype=np.int32)
13861398
inputs = []
13871399
outputs = []
13881400
if protocol == "http":
13891401
inputs.append(httpclient.InferInput("INPUT0", [1, 16], "INT32"))
13901402
inputs.append(httpclient.InferInput("INPUT1", [1, 16], "INT32"))
1391-
outputs.append(httpclient.InferRequestedOutput("OUTPUT0", binary_data=True))
1392-
outputs.append(httpclient.InferRequestedOutput("OUTPUT1", binary_data=False))
1403+
outputs.append(
1404+
httpclient.InferRequestedOutput(output_names[0], binary_data=True)
1405+
)
1406+
outputs.append(
1407+
httpclient.InferRequestedOutput(output_names[1], binary_data=False)
1408+
)
13931409
else:
13941410
inputs.append(grpcclient.InferInput("INPUT0", [1, 16], "INT32"))
13951411
inputs.append(grpcclient.InferInput("INPUT1", [1, 16], "INT32"))
1396-
outputs.append(grpcclient.InferRequestedOutput("OUTPUT0"))
1397-
outputs.append(grpcclient.InferRequestedOutput("OUTPUT1"))
1412+
outputs.append(grpcclient.InferRequestedOutput(output_names[0]))
1413+
outputs.append(grpcclient.InferRequestedOutput(output_names[1]))
13981414

13991415
inputs[0].set_shared_memory("input0_data", default_shm_byte_size)
14001416

@@ -1414,9 +1430,9 @@ def shm_basic_infer(
14141430

14151431
try:
14161432
results = triton_client.infer(
1417-
"simple", inputs, model_version="", outputs=outputs
1433+
model_name, inputs, model_version="", outputs=outputs
14181434
)
1419-
output = results.get_output("OUTPUT0")
1435+
output = results.get_output(output_names[0])
14201436
if protocol == "http":
14211437
output_datatype = output["datatype"]
14221438
output_shape = output["shape"]
@@ -1427,11 +1443,16 @@ def shm_basic_infer(
14271443

14281444
if use_system_shared_memory:
14291445
output_data = shm.get_contents_as_numpy(
1430-
shm_op0_handle, output_dtype, output_shape
1446+
shm_op0_handle,
1447+
output_dtype,
1448+
output_shape,
1449+
offset=register_offset + shm_output_offset,
14311450
)
14321451
elif use_cuda_shared_memory:
14331452
output_data = cudashm.get_contents_as_numpy(
1434-
shm_op0_handle, output_dtype, output_shape
1453+
shm_op0_handle,
1454+
output_dtype,
1455+
output_shape,
14351456
)
14361457

14371458
tester.assertTrue(

src/shared_memory_manager.cc

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -524,12 +524,7 @@ SharedMemoryManager::GetMemoryInfo(
524524
*shm_info = std::static_pointer_cast<const SharedMemoryInfo>(it->second);
525525
}
526526

527-
if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) {
528-
*shm_mapped_addr = (void*)((uint8_t*)it->second->mapped_addr_ +
529-
it->second->offset_ + offset);
530-
} else {
531-
*shm_mapped_addr = (void*)((uint8_t*)it->second->mapped_addr_ + offset);
532-
}
527+
*shm_mapped_addr = (void*)((uint8_t*)it->second->mapped_addr_ + offset);
533528

534529
*memory_type = it->second->kind_;
535530
*device_id = it->second->device_id_;

0 commit comments

Comments
 (0)