Skip to content

Commit e5aecb3

Browse files
whoisjwhoisj
authored
fix: Avoid overflows when reading json inputs (#8676)
Check input from HTTP+JSON sizes before allocating memory for them. Protect against integer overflows while we're at it.
1 parent e0a6a1e commit e5aecb3

6 files changed

Lines changed: 321 additions & 37 deletions

File tree

‎qa/L0_http/generate_endpoint_test.py‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#!/usr/bin/python3
2-
# Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
#
44
# Redistribution and use in source and binary forms, with or without
55
# modification, are permitted provided that the following conditions
@@ -29,6 +29,7 @@
2929

3030
sys.path.append("../common")
3131

32+
import base64
3233
import json
3334
import threading
3435
import time
@@ -274,6 +275,37 @@ def test_invalid_input_types(self):
274275
self.generate_expect_failure(self._model_name, inputs, error_msg)
275276
self.generate_stream_expect_failure(self._model_name, inputs, error_msg)
276277

278+
def test_json_dtype_size_expansion_exceeds_limit_error(self):
279+
"""
280+
Test that when the client sends a JSON input of byte[], that when it
281+
expands to dtype[], it exceeds the maximum allowed input size and
282+
returns an appropriate error message. The test sends a large base64
283+
encoded string as input, which simulates a byte[] input that would
284+
expand to a much larger dtype[] input on the server side when
285+
`sizeof(dtype) > 1`.
286+
The test checks that the error message indicates that the input size
287+
exceeds the limit.
288+
This is important to prevent clients from sending inputs that could
289+
cause excessive memory usage on the server.
290+
"""
291+
292+
input_data = [1] * (
293+
64 * 1024 * 1024
294+
) # 64MB input, which is large but still reasonable for HTTP request body
295+
input_bytes = bytes(input_data)
296+
input_str = base64.b64encode(input_bytes).decode("utf-8")
297+
inputs = {"PROMPT": input_str, "STREAM": False}
298+
error_msg = " bytes exceeds the maximum allowed input size of "
299+
self.generate_expect_failure(self._model_name, inputs, error_msg)
300+
301+
inputs = {
302+
"INPUT0": input_str[0 : (len(input_str) // 2)],
303+
"INPUT1": input_str[(len(input_str) // 2) :],
304+
"STREAM": False,
305+
}
306+
error_msg = " bytes exceeds the maximum allowed input size of "
307+
self.generate_expect_failure(self._model_name, inputs, error_msg)
308+
277309
def test_duplicate_inputs(self):
278310
dupe_prompt = "input 'PROMPT' already exists in request"
279311
dupe_stream = "input 'STREAM' already exists in request"

‎qa/L0_http/http_input_size_limit_test.py‎

Lines changed: 140 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#!/usr/bin/python
2-
# Copyright 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# Copyright 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
#
44
# Redistribution and use in source and binary forms, with or without
55
# modification, are permitted provided that the following conditions
@@ -29,6 +29,7 @@
2929

3030
sys.path.append("../common")
3131

32+
import base64
3233
import gzip
3334
import io
3435
import json
@@ -41,6 +42,9 @@
4142
# Constants for size calculations
4243
# Each FP32 value is 4 bytes, so we need to divide target byte sizes by 4 to get element counts
4344
BYTES_PER_FP32 = 4
45+
BYTES_PER_INT64 = (
46+
8 # For the type size explosion test, we use int64 which is 8 bytes per element
47+
)
4448
MB = 2**20 # 1 MB = 1,048,576 bytes
4549
GB = 2**30 # 1 GB = 1,073,741,824 bytes
4650
DEFAULT_LIMIT_BYTES = 64 * MB # 64MB default limit
@@ -58,7 +62,120 @@
5862

5963
class InferSizeLimitTest(tu.TestResultCollector):
6064
def _get_infer_url(self, model_name):
61-
return "http://localhost:8000/v2/models/{}/infer".format(model_name)
65+
return f"http://localhost:8000/v2/models/{model_name}/infer"
66+
67+
def test_json_dtype_size_expansion_exceeds_limit_error(self):
68+
"""
69+
Test that when the client sends a JSON input of byte[], that when it
70+
expands to dtype[], it exceeds the maximum allowed input size and
71+
returns an appropriate error message. The test sends a large base64
72+
encoded string as input, which simulates a byte[] input that would
73+
expand to a much larger dtype[] input on the server side when
74+
`sizeof(dtype) > 1`.
75+
The test checks that the error message indicates that the input size
76+
exceeds the limit.
77+
This is important to prevent clients from sending inputs that could
78+
cause excessive memory usage on the server.
79+
"""
80+
model = "onnx_zero_1_float32"
81+
82+
# Provided data is 64MB of int8, but the model expects FP32,
83+
# which would expand to 256MB when interpreted as FP32.
84+
bytes_input = np.ones(DEFAULT_LIMIT_BYTES, dtype=np.int8)
85+
input_bytes = bytes_input.tobytes()
86+
data_str = base64.b64encode(input_bytes).decode("utf-8")
87+
headers = {
88+
"Content-Type": "application/json",
89+
"Inference-Header-Content-Length": f"{len(input_bytes)}",
90+
}
91+
shape_size = (
92+
DEFAULT_LIMIT_ELEMENTS // BYTES_PER_INT64
93+
) # Calculate shape size based on int64 element count to match the byte size
94+
95+
payload = {
96+
"inputs": [
97+
{
98+
"name": "INPUT0",
99+
"datatype": "INT64",
100+
"shape": [1, shape_size],
101+
"data": data_str,
102+
}
103+
]
104+
}
105+
106+
response = requests.post(
107+
f"http://localhost:8000/v2/models/{model}/generate",
108+
headers=headers,
109+
json=payload,
110+
)
111+
112+
self.assertEqual(
113+
400,
114+
response.status_code,
115+
f"Expected error code for type/size mismatch, got: {response.status_code}",
116+
)
117+
error_msg = response.content.decode()
118+
print(
119+
f"Error message: {error_msg}", flush=True
120+
) # Print the error message for debugging
121+
self.assertIn(
122+
"Request JSON size of ",
123+
error_msg,
124+
)
125+
self.assertIn(
126+
" bytes exceeds the maximum allowed input size of ",
127+
error_msg,
128+
)
129+
self.assertIn(
130+
"Use --http-max-input-size to increase the limit.",
131+
error_msg,
132+
)
133+
134+
# Test multiple inputs with one that causes size explosion.
135+
payload = {
136+
"inputs": [
137+
{
138+
"name": "INPUT0",
139+
"datatype": "INT64",
140+
"shape": [1, shape_size // 2],
141+
"data": data_str[: len(data_str) // 2],
142+
},
143+
{
144+
"name": "INPUT1",
145+
"datatype": "INT64",
146+
"shape": [1, shape_size // 2],
147+
"data": data_str[len(data_str) // 2 :],
148+
},
149+
]
150+
}
151+
152+
response = requests.post(
153+
f"http://localhost:8000/v2/models/{model}/generate",
154+
headers=headers,
155+
json=payload,
156+
)
157+
158+
self.assertEqual(
159+
400,
160+
response.status_code,
161+
f"Expected error code for type/size mismatch, got: {response.status_code}",
162+
)
163+
error_msg = response.content.decode()
164+
print(
165+
f"Error message: {error_msg}", flush=True
166+
) # Print the error message for debugging
167+
self.assertIn(
168+
"request JSON size of ",
169+
error_msg,
170+
)
171+
self.assertIn(
172+
" bytes exceeds the maximum allowed input size of ",
173+
error_msg,
174+
)
175+
self.assertIn(
176+
"Use --http-max-input-size to increase the limit.",
177+
error_msg,
178+
)
62179

63180
def test_default_limit_raw_binary(self):
64181
"""Test raw binary inputs with default limit"""
@@ -165,9 +282,16 @@ def test_default_limit_json(self):
165282
# Verify error message contains size limit info
166283
error_msg = response.content.decode()
167284
self.assertIn(
168-
"exceeds the maximum allowed value",
285+
"Request JSON size of ",
286+
error_msg,
287+
)
288+
self.assertIn(
289+
" bytes exceeds the maximum allowed input size of ",
290+
error_msg,
291+
)
292+
self.assertIn(
293+
"Use --http-max-input-size to increase the limit.",
169294
error_msg,
170-
"Expected error message about exceeding max input size",
171295
)
172296

173297
# Test case 2: Input just under the 64MB limit (should succeed)
@@ -320,9 +444,16 @@ def test_large_input_json(self):
320444
# Verify error message contains size limit info
321445
error_msg = response.content.decode()
322446
self.assertIn(
323-
"exceeds the maximum allowed value",
447+
"request JSON size of ",
448+
error_msg,
449+
)
450+
self.assertIn(
451+
" bytes exceeds the maximum allowed input size of ",
452+
error_msg,
453+
)
454+
self.assertIn(
455+
"Use --http-max-input-size to increase the limit.",
324456
error_msg,
325-
"Expected error message about exceeding max input size",
326457
)
327458

328459
# Test case 2: Input just under the 128MB configured limit (should succeed)
@@ -405,15 +536,15 @@ def test_large_string_in_json(self):
405536
# Verify error message
406537
error_msg = response.content.decode()
407538
self.assertIn(
408-
"Request JSON size",
539+
"Request JSON size of ",
409540
error_msg,
410541
)
411542
self.assertIn(
412-
"exceeds the maximum allowed value",
543+
" bytes exceeds the maximum allowed input size of ",
413544
error_msg,
414545
)
415546
self.assertIn(
416-
"Use --http-max-input-size to increase the limit",
547+
"Use --http-max-input-size to increase the limit.",
417548
error_msg,
418549
)
419550

‎qa/L0_http/http_test.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#!/usr/bin/python
2-
# Copyright 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# Copyright 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
#
44
# Redistribution and use in source and binary forms, with or without
55
# modification, are permitted provided that the following conditions
@@ -368,7 +368,7 @@ def test_loading_large_invalid_model(self):
368368
error_message,
369369
)
370370
self.assertIn(
371-
"exceeds the maximum allowed value",
371+
" exceeds the maximum allowed input size. ",
372372
error_message,
373373
)
374374
except ValueError:

0 commit comments

Comments
 (0)