From 47bcd55d1268225d820ab6f76a75f5aac9c1da9a Mon Sep 17 00:00:00 2001 From: Matthew Wittwer Date: Fri, 6 Mar 2026 00:09:31 +0000 Subject: [PATCH 1/4] Block Vertex AI redirects --- qa/L0_vertex_ai/test.sh | 157 ++++++++++++++++++++++++++++++++-------- src/vertex_ai_server.cc | 22 ++++-- 2 files changed, 142 insertions(+), 37 deletions(-) diff --git a/qa/L0_vertex_ai/test.sh b/qa/L0_vertex_ai/test.sh index 7a2dc86a2a..6514ff8aab 100755 --- a/qa/L0_vertex_ai/test.sh +++ b/qa/L0_vertex_ai/test.sh @@ -760,62 +760,116 @@ else fi set -e -# Redirected management APIs should be blocked without restricted header +# Redirected read-only APIs should be blocked without restricted header set +e for redirect_endpoint in \ "v2" \ "v2/models/identity_fp32/config" \ "v2/repository/index" \ "v2/systemsharedmemory/status" \ - "v2/cudasharedmemory/status" \ - "v2/systemsharedmemory/region/test/register" + "v2/cudasharedmemory/status" do rm -f ./curl.out - if [ "$redirect_endpoint" != "v2/systemsharedmemory/region/test/register" ]; then - code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: ${redirect_endpoint}" localhost:8080/predict` - else - code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: ${redirect_endpoint}" \ - -H "Content-Type: application/json" -d '{"key":"test_shm","byte_size":1024}' localhost:8080/predict` - fi + code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: ${redirect_endpoint}" localhost:8080/predict` if [ "$code" != "403" ]; then cat ./curl.out echo -e "\n***\n*** Failed. Expected ${redirect_endpoint} is restricted\n***" RET=1 fi done + +# Mutating shared memory operations are blocked through redirect +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: v2/systemsharedmemory/region/test/register" -H "Content-Type: application/json" -d '{"key":"test_shm","byte_size":1024}' localhost:8080/predict` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected shared memory register is blocked via redirect\n***" + RET=1 +fi + +# Model load redirect is unconditionally blocked through the prediction endpoint +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: v2/repository/models/identity_fp32/load" localhost:8080/predict` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected model load via redirect is blocked\n***" + RET=1 +fi + +# Model unload redirect is unconditionally blocked through the prediction endpoint +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: v2/repository/models/identity_fp32/unload" localhost:8080/predict` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected model unload via redirect is blocked\n***" + RET=1 +fi + +# Statistics redirect without restricted header should be blocked +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: v2/models/identity_fp32/stats" localhost:8080/predict` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected model statistics via redirect is restricted\n***" + RET=1 +fi set -e -# Restricted header should allow handler invocation +# Restricted header should allow read-only handler invocation set +e for redirect_endpoint in \ "v2" \ "v2/models/identity_fp32/config" \ "v2/repository/index" \ "v2/systemsharedmemory/status" \ - "v2/cudasharedmemory/status" \ - "v2/systemsharedmemory/region/test/register" + "v2/cudasharedmemory/status" do rm -f ./curl.out - if [ "$redirect_endpoint" != "v2/systemsharedmemory/region/test/register" ]; then - code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: ${redirect_endpoint}" \ - -H "X-Vertex-Restricted: secret" localhost:8080/predict` - if [ "$code" != "200" ]; then - cat ./curl.out - echo -e "\n***\n*** Failed. Expected ${redirect_endpoint} passes with restricted header\n***" - RET=1 - fi - else - code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: ${redirect_endpoint}" \ - -H "X-Vertex-Restricted: secret" -H "Content-Type: application/json" -d '{"key":"test_shm","byte_size":1024}' localhost:8080/predict` - # This request may fail with a different error: "Unable to open shared memory region: 'test_shm'" - if [ "$code" == "403" ]; then - cat ./curl.out - echo -e "\n***\n*** Failed. Expected ${redirect_endpoint} passes with restricted header\n***" - RET=1 - fi + code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: ${redirect_endpoint}" \ + -H "X-Vertex-Restricted: secret" localhost:8080/predict` + if [ "$code" != "200" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected ${redirect_endpoint} passes with restricted header\n***" + RET=1 fi done +# Mutating shared memory operations remain blocked even with valid header +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: v2/systemsharedmemory/region/test/register" -H "X-Vertex-Restricted: secret" -H "Content-Type: application/json" -d '{"key":"test_shm","byte_size":1024}' localhost:8080/predict` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected shared memory register is blocked via redirect even with valid header\n***" + RET=1 +fi + +# Model load remains blocked even with valid restricted header +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: v2/repository/models/identity_fp32/load" -H "X-Vertex-Restricted: secret" localhost:8080/predict` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected model load is blocked via redirect even with valid header\n***" + RET=1 +fi + +# Model unload remains blocked even with valid restricted header +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: v2/repository/models/identity_fp32/unload" -H "X-Vertex-Restricted: secret" localhost:8080/predict` +if [ "$code" != "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected model unload is blocked via redirect even with valid header\n***" + RET=1 +fi + +# Statistics redirect with restricted header should pass +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: v2/models/identity_fp32/stats" -H "X-Vertex-Restricted: secret" localhost:8080/predict` +if [ "$code" == "403" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected model statistics passes with restricted header\n***" + RET=1 +fi + # Wrong restricted header should reject the request set +e for redirect_endpoint in \ @@ -845,6 +899,51 @@ set -e kill $SERVER_PID wait $SERVER_PID +# +# HTTP max input size enforcement on Vertex AI endpoint +# +unset_vertex_variables +export AIP_PREDICT_ROUTE="/predict" +export AIP_HEALTH_ROUTE="/health" + +SERVER_LOG="vertex_max_input_size_server.log" +SERVER_ARGS="--allow-vertex-ai=true --model-repository=restricted_single_model --vertex-ai-default-model=identity_fp32 --http-max-input-size=256" +run_server_nowait +vertex_ai_wait_for_server_ready $SERVER_PID 10 +if [ "$WAIT_RET" != "0" ]; then + echo -e "\n***\n*** Failed to start $SERVER\n***" + kill $SERVER_PID + wait $SERVER_PID + cat $SERVER_LOG + exit 1 +fi + +set +e + +# Small payload under 256 bytes should succeed +rm -f ./curl.out +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "Content-Type: application/json" -d'{"inputs":[{"name":"INPUT0","datatype":"FP32","shape":[1,1],"data":[1.0]}],"outputs":[{"name":"OUTPUT0"}]}' localhost:8080/predict` +if [ "$code" != "200" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected small payload to succeed on Vertex AI endpoint (got $code)\n***" + RET=1 +fi + +# Large payload over 256 bytes should be rejected +rm -f ./curl.out +LARGE_PAYLOAD='{"inputs":[{"name":"INPUT0","datatype":"FP32","shape":[1,16],"data":[1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0,13.0,14.0,15.0,16.0]}],"outputs":[{"name":"OUTPUT0"}]}' +code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "Content-Type: application/json" -d"$LARGE_PAYLOAD" localhost:8080/predict` +if [ "$code" == "200" ]; then + cat ./curl.out + echo -e "\n***\n*** Failed. Expected oversized payload to be rejected on Vertex AI endpoint\n***" + RET=1 +fi + +set -e + +kill $SERVER_PID +wait $SERVE_PID + if [ $RET -eq 0 ]; then echo -e "\n***\n*** Test Passed\n***" else diff --git a/src/vertex_ai_server.cc b/src/vertex_ai_server.cc index 9bfbe9033d..e1eb0a833c 100644 --- a/src/vertex_ai_server.cc +++ b/src/vertex_ai_server.cc @@ -196,32 +196,38 @@ VertexAiAPIServer::Handle(evhtp_request_t* req) } else if (RE2::FullMatch( redirect_endpoint, systemsharedmemory_regex_, ®ion, &action)) { - // system shared memory if (action == "status") { req->method = htp_method_GET; + HandleSystemSharedMemory(req, region, action); + return; } - HandleSystemSharedMemory(req, region, action); + // register/unregister are mutating operations that must not be + // reachable through the prediction endpoint. + LOG_VERBOSE(1) << "Vertex AI redirect blocked: " << redirect_endpoint; + evhtp_send_reply(req, EVHTP_RES_FORBIDDEN); return; } else if (RE2::FullMatch( redirect_endpoint, cudasharedmemory_regex_, ®ion, &action)) { - // cuda shared memory if (action == "status") { req->method = htp_method_GET; + HandleCudaSharedMemory(req, region, action); + return; } - HandleCudaSharedMemory(req, region, action); + LOG_VERBOSE(1) << "Vertex AI redirect blocked: " << redirect_endpoint; + evhtp_send_reply(req, EVHTP_RES_FORBIDDEN); return; } else if (RE2::FullMatch( redirect_endpoint, modelcontrol_regex_, &repo_name, &kind, &model_name, &action)) { - // model repository if (kind == "index") { HandleRepositoryIndex(req, repo_name); return; - } else if (kind.find("models", 0) == 0) { - HandleRepositoryControl(req, repo_name, model_name, action); - return; } + // Model load/unload must not be reachable through the prediction endpoint. + LOG_VERBOSE(1) << "Vertex AI redirect blocked: " << redirect_endpoint; + evhtp_send_reply(req, EVHTP_RES_FORBIDDEN); + return; } } } From 5e5f440ba990e18b324171ede3621050882bfe20 Mon Sep 17 00:00:00 2001 From: Matthew Wittwer Date: Fri, 6 Mar 2026 00:21:59 +0000 Subject: [PATCH 2/4] cleanup comments --- qa/L0_vertex_ai/test.sh | 14 ++++---------- src/vertex_ai_server.cc | 8 +++++--- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/qa/L0_vertex_ai/test.sh b/qa/L0_vertex_ai/test.sh index 6514ff8aab..f419f5c6bb 100755 --- a/qa/L0_vertex_ai/test.sh +++ b/qa/L0_vertex_ai/test.sh @@ -699,20 +699,14 @@ else fi set -e -# repository control (expect error) +# repository control via redirect is blocked unconditionally rm -f ./curl.out set +e code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: v2/repository/models/subadd/unload" localhost:8080/predict` -if [ "$code" == "200" ]; then +if [ "$code" != "403" ]; then cat ./curl.out - echo -e "\n***\n*** Test Failed\n***" + echo -e "\n***\n*** Failed. Expected model unload via redirect to return 403 (got $code)\n***" RET=1 -else - grep "explicit model load / unload is not allowed" ./curl.out - if [ $? -ne 0 ]; then - echo -e "\n***\n*** Failed. Expected error on model control\n***" - RET=1 - fi fi set -e @@ -907,7 +901,7 @@ export AIP_PREDICT_ROUTE="/predict" export AIP_HEALTH_ROUTE="/health" SERVER_LOG="vertex_max_input_size_server.log" -SERVER_ARGS="--allow-vertex-ai=true --model-repository=restricted_single_model --vertex-ai-default-model=identity_fp32 --http-max-input-size=256" +SERVER_ARGS="--allow-vertex-ai=true --model-repository=restricted_single_model --vertex-ai-default-model=identity_fp32 --http-max-input-size=128" run_server_nowait vertex_ai_wait_for_server_ready $SERVER_PID 10 if [ "$WAIT_RET" != "0" ]; then diff --git a/src/vertex_ai_server.cc b/src/vertex_ai_server.cc index e1eb0a833c..ec4fd6802f 100644 --- a/src/vertex_ai_server.cc +++ b/src/vertex_ai_server.cc @@ -201,8 +201,9 @@ VertexAiAPIServer::Handle(evhtp_request_t* req) HandleSystemSharedMemory(req, region, action); return; } - // register/unregister are mutating operations that must not be - // reachable through the prediction endpoint. + // Only read-only status queries are permitted through redirect. + // Mutating operations (register/unregister) require the core + // HTTP endpoint. LOG_VERBOSE(1) << "Vertex AI redirect blocked: " << redirect_endpoint; evhtp_send_reply(req, EVHTP_RES_FORBIDDEN); return; @@ -224,7 +225,8 @@ VertexAiAPIServer::Handle(evhtp_request_t* req) HandleRepositoryIndex(req, repo_name); return; } - // Model load/unload must not be reachable through the prediction endpoint. + // Only repository index queries are permitted through redirect. + // Model load/unload requires the core HTTP endpoint. LOG_VERBOSE(1) << "Vertex AI redirect blocked: " << redirect_endpoint; evhtp_send_reply(req, EVHTP_RES_FORBIDDEN); return; From 6f9fb12ffb3f302e0a269d05579323bad30a3284 Mon Sep 17 00:00:00 2001 From: Matthew Wittwer Date: Fri, 6 Mar 2026 20:22:31 +0000 Subject: [PATCH 3/4] complete rebase --- qa/L0_vertex_ai/test.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/qa/L0_vertex_ai/test.sh b/qa/L0_vertex_ai/test.sh index f419f5c6bb..2c2cb74a3d 100755 --- a/qa/L0_vertex_ai/test.sh +++ b/qa/L0_vertex_ai/test.sh @@ -914,7 +914,7 @@ fi set +e -# Small payload under 256 bytes should succeed +# Small payload under 128 bytes should succeed rm -f ./curl.out code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "Content-Type: application/json" -d'{"inputs":[{"name":"INPUT0","datatype":"FP32","shape":[1,1],"data":[1.0]}],"outputs":[{"name":"OUTPUT0"}]}' localhost:8080/predict` if [ "$code" != "200" ]; then @@ -923,7 +923,7 @@ if [ "$code" != "200" ]; then RET=1 fi -# Large payload over 256 bytes should be rejected +# Large payload over 128 bytes should be rejected rm -f ./curl.out LARGE_PAYLOAD='{"inputs":[{"name":"INPUT0","datatype":"FP32","shape":[1,16],"data":[1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0,13.0,14.0,15.0,16.0]}],"outputs":[{"name":"OUTPUT0"}]}' code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "Content-Type: application/json" -d"$LARGE_PAYLOAD" localhost:8080/predict` @@ -936,7 +936,7 @@ fi set -e kill $SERVER_PID -wait $SERVE_PID +wait $SERVER_PID if [ $RET -eq 0 ]; then echo -e "\n***\n*** Test Passed\n***" From 93df24cb252bbceee3c15a664843d91d74944629 Mon Sep 17 00:00:00 2001 From: Matthew Wittwer Date: Fri, 6 Mar 2026 20:39:44 +0000 Subject: [PATCH 4/4] cleanup readability --- qa/L0_vertex_ai/test.sh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/qa/L0_vertex_ai/test.sh b/qa/L0_vertex_ai/test.sh index 2c2cb74a3d..c8da5ed57d 100755 --- a/qa/L0_vertex_ai/test.sh +++ b/qa/L0_vertex_ai/test.sh @@ -901,7 +901,10 @@ export AIP_PREDICT_ROUTE="/predict" export AIP_HEALTH_ROUTE="/health" SERVER_LOG="vertex_max_input_size_server.log" -SERVER_ARGS="--allow-vertex-ai=true --model-repository=restricted_single_model --vertex-ai-default-model=identity_fp32 --http-max-input-size=128" +SERVER_ARGS="--allow-vertex-ai=true \ + --model-repository=restricted_single_model \ + --vertex-ai-default-model=identity_fp32 \ + --http-max-input-size=128" run_server_nowait vertex_ai_wait_for_server_ready $SERVER_PID 10 if [ "$WAIT_RET" != "0" ]; then