Skip to content

Commit 411ac70

Browse files
yinggehyinggeh
committed
Remove HTTP_MAX_CHUNKS
1 parent 982dd3b commit 411ac70

5 files changed

Lines changed: 70 additions & 105 deletions

File tree

‎qa/L0_http/http_request_many_chunks.py‎

Lines changed: 11 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,8 @@ def setUp(self):
3434
self._model_name = "simple"
3535
self._local_host = "localhost"
3636
self._http_port = 8000
37-
self._max_chunk_count = (
38-
1000000 # defined in server/src/common.h HTTP_MAX_CHUNKS
37+
self._malicious_chunk_count = (
38+
1000000 # large enough to cause a stack overflow if using alloca()
3939
)
4040
self._parse_error = (
4141
"failed to parse the request JSON buffer: Invalid value. at 0"
@@ -80,13 +80,6 @@ def send_chunked_request(
8080
finally:
8181
s.close()
8282

83-
def chunk_exceed_error(self, request_kind: str, chunk_count: int):
84-
# defined in server/src/http_server.cc
85-
if request_kind == "input":
86-
return f"Chunks in the {request_kind} buffer exceed the limit {self._max_chunk_count}, got {chunk_count} chunks"
87-
else:
88-
return f"Chunks in the {request_kind} request buffer exceed the limit {self._max_chunk_count}, got {chunk_count} chunks"
89-
9083
def test_infer(self):
9184
request_header = (
9285
f"POST /v2/models/{self._model_name}/infer HTTP/1.1\r\n"
@@ -95,25 +88,15 @@ def test_infer(self):
9588

9689
self.send_chunked_request(
9790
request_header,
98-
self._max_chunk_count,
91+
self._malicious_chunk_count,
9992
"Raw request must only have 1 input (found 1) to be deduced but got 2 inputs in 'simple' model configuration",
10093
)
101-
self.send_chunked_request(
102-
request_header,
103-
self._max_chunk_count + 1,
104-
self.chunk_exceed_error("input", self._max_chunk_count + 1),
105-
)
10694

10795
def test_registry_index(self):
10896
request_header = f"POST /v2/repository/index HTTP/1.1\r\n"
10997

11098
self.send_chunked_request(
111-
request_header, self._max_chunk_count, self._parse_error
112-
)
113-
self.send_chunked_request(
114-
request_header,
115-
self._max_chunk_count + 1,
116-
self.chunk_exceed_error("registry index", self._max_chunk_count + 1),
99+
request_header, self._malicious_chunk_count, self._parse_error
117100
)
118101

119102
def test_model_control(self):
@@ -123,20 +106,10 @@ def test_model_control(self):
123106
unload_request_header = load_request_header.replace("/load", "/unload")
124107

125108
self.send_chunked_request(
126-
load_request_header, self._max_chunk_count, self._parse_error
109+
load_request_header, self._malicious_chunk_count, self._parse_error
127110
)
128111
self.send_chunked_request(
129-
load_request_header,
130-
self._max_chunk_count + 1,
131-
self.chunk_exceed_error("load model", self._max_chunk_count + 1),
132-
)
133-
self.send_chunked_request(
134-
unload_request_header, self._max_chunk_count, self._parse_error
135-
)
136-
self.send_chunked_request(
137-
unload_request_header,
138-
self._max_chunk_count + 1,
139-
self.chunk_exceed_error("unload model", self._max_chunk_count + 1),
112+
unload_request_header, self._malicious_chunk_count, self._parse_error
140113
)
141114

142115
def test_trace(self):
@@ -145,59 +118,34 @@ def test_trace(self):
145118
)
146119

147120
self.send_chunked_request(
148-
request_header, self._max_chunk_count, self._parse_error
149-
)
150-
self.send_chunked_request(
151-
request_header,
152-
self._max_chunk_count + 1,
153-
self.chunk_exceed_error("trace", self._max_chunk_count + 1),
121+
request_header, self._malicious_chunk_count, self._parse_error
154122
)
155123

156124
def test_logging(self):
157125
request_header = f"POST /v2/logging HTTP/1.1\r\n"
158126

159127
self.send_chunked_request(
160-
request_header, self._max_chunk_count, self._parse_error
161-
)
162-
self.send_chunked_request(
163-
request_header,
164-
self._max_chunk_count + 1,
165-
self.chunk_exceed_error("dynamic logging", self._max_chunk_count + 1),
128+
request_header, self._malicious_chunk_count, self._parse_error
166129
)
167130

168131
def test_system_shm_register(self):
169132
request_header = f"POST /v2/systemsharedmemory/region/test_system_shm_register/register HTTP/1.1\r\n"
170133

171134
self.send_chunked_request(
172-
request_header, self._max_chunk_count, self._parse_error
173-
)
174-
self.send_chunked_request(
175-
request_header,
176-
self._max_chunk_count + 1,
177-
self.chunk_exceed_error("register", self._max_chunk_count + 1),
135+
request_header, self._malicious_chunk_count, self._parse_error
178136
)
179137

180138
def test_cuda_shm_register(self):
181139
request_header = f"POST /v2/cudasharedmemory/region/test_cuda_shm_register/register HTTP/1.1\r\n"
182140

183141
self.send_chunked_request(
184-
request_header, self._max_chunk_count, self._parse_error
185-
)
186-
self.send_chunked_request(
187-
request_header,
188-
self._max_chunk_count + 1,
189-
self.chunk_exceed_error("register", self._max_chunk_count + 1),
142+
request_header, self._malicious_chunk_count, self._parse_error
190143
)
191144

192145
def test_generate(self):
193146
request_header = f"POST /v2/models/{self._model_name}/generate HTTP/1.1\r\n"
194147
self.send_chunked_request(
195-
request_header, self._max_chunk_count, self._parse_error
196-
)
197-
self.send_chunked_request(
198-
request_header,
199-
self._max_chunk_count + 1,
200-
self.chunk_exceed_error("generate", self._max_chunk_count + 1),
148+
request_header, self._malicious_chunk_count, self._parse_error
201149
)
202150

203151

‎qa/L0_sagemaker/sagemaker_request_many_chunks.py‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,8 @@ class SagemakerRequestManyChunksTest(unittest.TestCase):
3333
def setUp(self):
3434
self._local_host = "localhost"
3535
self._sagemaker_port = 8080
36-
self._max_chunk_count = (
37-
1000000 # defined in server/src/common.h HTTP_MAX_CHUNKS
36+
self._malicious_chunk_count = (
37+
1000000 # large enough to cause a stack overflow if using alloca()
3838
)
3939

4040
def send_chunked_request(
@@ -82,14 +82,9 @@ def test_load_model(self):
8282
)
8383
self.send_chunked_request(
8484
request_header,
85-
self._max_chunk_count,
85+
self._malicious_chunk_count,
8686
"failed to parse the request JSON buffer: Invalid value. at 0",
8787
)
88-
self.send_chunked_request(
89-
request_header,
90-
self._max_chunk_count + 1,
91-
f"Chunks in the load model request buffer exceed the limit {self._max_chunk_count}, got {self._max_chunk_count + 1} chunks",
92-
)
9388

9489

9590
if __name__ == "__main__":

‎qa/L0_sagemaker/test.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -574,7 +574,7 @@ rm -rf models && mkdir models && \
574574
export SAGEMAKER_TRITON_DEFAULT_MODEL_NAME=sm_model
575575
REQUEST_MANY_CHUNKS_PY="sagemaker_request_many_chunks.py"
576576
CLIENT_LOG="./client.sagemaker_request_many_chunks.log"
577-
SERVER_LOG="./server_request_many_chunks.log"
577+
SERVER_LOG="./server.sagemaker_request_many_chunks.log"
578578

579579
serve > $SERVER_LOG 2>&1 &
580580
SERVE_PID=$!

‎src/common.h‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,6 @@ constexpr int32_t HTTP_MAX_JSON_NESTING_DEPTH = 100;
6161
// Default maximum allowed HTTP request input size in bytes (64MB)
6262
constexpr size_t HTTP_DEFAULT_MAX_INPUT_SIZE = 1 << 26;
6363

64-
// Maximum allowed number of chunks per HTTP request (arbitrary large number).
65-
// This limit prevents potential memory exhaustion from malformed requests.
66-
constexpr int32_t HTTP_MAX_CHUNKS = 1000000;
67-
6864
/// Request parameter keys that start with a "triton_" prefix for internal use
6965
const std::vector<std::string> TRITON_RESERVED_REQUEST_PARAMS{
7066
"triton_enable_empty_final_response"};

‎src/http_server.cc‎

Lines changed: 55 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -2875,18 +2875,25 @@ HTTPAPIServer::EVRequestToJsonImpl(
28752875
std::vector<struct evbuffer_iovec> v_vec;
28762876

28772877
int n = evbuffer_peek(req->buffer_in, -1, NULL, NULL, 0);
2878-
2879-
if (n > HTTP_MAX_CHUNKS) {
2880-
return TRITONSERVER_ErrorNew(
2881-
TRITONSERVER_ERROR_INVALID_ARG,
2882-
("Chunks in the " + std::string(request_kind) +
2883-
" request buffer exceed the limit " + std::to_string(HTTP_MAX_CHUNKS) +
2884-
", got " + std::to_string(n) + " chunks")
2885-
.c_str());
2886-
}
2887-
28882878
if (n > 0) {
2889-
v_vec = std::vector<struct evbuffer_iovec>(n);
2879+
try {
2880+
v_vec = std::vector<struct evbuffer_iovec>(n);
2881+
}
2882+
catch (const std::bad_alloc& e) {
2883+
// Handle memory allocation failure
2884+
return TRITONSERVER_ErrorNew(
2885+
TRITONSERVER_ERROR_INVALID_ARG,
2886+
(std::string("Memory allocation failed for evbuffer: ") + e.what())
2887+
.c_str());
2888+
}
2889+
catch (const std::exception& e) {
2890+
// Catch any other std exceptions
2891+
return TRITONSERVER_ErrorNew(
2892+
TRITONSERVER_ERROR_INTERNAL,
2893+
(std::string("Exception while creating evbuffer vector: ") + e.what())
2894+
.c_str());
2895+
}
2896+
28902897
v = v_vec.data();
28912898
if (evbuffer_peek(req->buffer_in, -1, NULL, v, n) != n) {
28922899
return TRITONSERVER_ErrorNew(
@@ -2921,16 +2928,25 @@ HTTPAPIServer::EVBufferToInput(
29212928
std::vector<struct evbuffer_iovec> v_vec;
29222929

29232930
int n = evbuffer_peek(input_buffer, -1, NULL, NULL, 0);
2924-
if (n > HTTP_MAX_CHUNKS) {
2925-
return TRITONSERVER_ErrorNew(
2926-
TRITONSERVER_ERROR_INVALID_ARG,
2927-
("Chunks in the input buffer exceed the limit " +
2928-
std::to_string(HTTP_MAX_CHUNKS) + ", got " + std::to_string(n) +
2929-
" chunks")
2930-
.c_str());
2931-
}
29322931
if (n > 0) {
2933-
v_vec = std::vector<struct evbuffer_iovec>(n);
2932+
try {
2933+
v_vec = std::vector<struct evbuffer_iovec>(n);
2934+
}
2935+
catch (const std::bad_alloc& e) {
2936+
// Handle memory allocation failure
2937+
return TRITONSERVER_ErrorNew(
2938+
TRITONSERVER_ERROR_INVALID_ARG,
2939+
(std::string("Memory allocation failed for evbuffer: ") + e.what())
2940+
.c_str());
2941+
}
2942+
catch (const std::exception& e) {
2943+
// Catch any other std exceptions
2944+
return TRITONSERVER_ErrorNew(
2945+
TRITONSERVER_ERROR_INTERNAL,
2946+
(std::string("Exception while creating evbuffer vector: ") + e.what())
2947+
.c_str());
2948+
}
2949+
29342950
v = v_vec.data();
29352951
if (evbuffer_peek(input_buffer, -1, NULL, v, n) != n) {
29362952
return TRITONSERVER_ErrorNew(
@@ -2987,16 +3003,26 @@ HTTPAPIServer::EVBufferToRawInput(
29873003
std::vector<struct evbuffer_iovec> v_vec;
29883004

29893005
int n = evbuffer_peek(input_buffer, -1, NULL, NULL, 0);
2990-
if (n > HTTP_MAX_CHUNKS) {
2991-
return TRITONSERVER_ErrorNew(
2992-
TRITONSERVER_ERROR_INVALID_ARG,
2993-
("Chunks in the input buffer exceed the limit " +
2994-
std::to_string(HTTP_MAX_CHUNKS) + ", got " + std::to_string(n) +
2995-
" chunks")
2996-
.c_str());
2997-
}
29983006
if (n > 0) {
2999-
v_vec = std::vector<struct evbuffer_iovec>(n);
3007+
try {
3008+
v_vec = std::vector<struct evbuffer_iovec>(n);
3009+
}
3010+
catch (const std::bad_alloc& e) {
3011+
// Handle memory allocation failure
3012+
return TRITONSERVER_ErrorNew(
3013+
TRITONSERVER_ERROR_INVALID_ARG,
3014+
(std::string("Memory allocation failed for evbuffer: ") + e.what())
3015+
.c_str());
3016+
}
3017+
catch (const std::exception& e) {
3018+
// Catch any other std exceptions
3019+
return TRITONSERVER_ErrorNew(
3020+
TRITONSERVER_ERROR_INTERNAL,
3021+
(std::string("Exception while creating evbuffer vector: ") +
3022+
e.what())
3023+
.c_str());
3024+
}
3025+
30003026
v = v_vec.data();
30013027
if (evbuffer_peek(input_buffer, -1, NULL, v, n) != n) {
30023028
return TRITONSERVER_ErrorNew(

0 commit comments

Comments
 (0)