Block Vertex AI redirects

Author: mattwittwer

qa/L0_vertex_ai/test.sh

@@ -809,7 +809,35 @@ 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 operation is restricted\n***"
+    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 regardless of restricted header configuration (TRI-715 RCE vector)
+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 (TRI-715 DoS vector)
+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
@@ -856,11 +884,39 @@ if [ "$code" != "200" ]; then
     RET=1
 fi
 
+# 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 valid restricted header to pass restriction check (request may still fail for non-auth reasons)\n***"
+    echo -e "\n***\n*** Failed. Expected model statistics passes with restricted header\n***"
     RET=1
 fi
 
@@ -877,6 +933,51 @@ set -e
 kill $SERVER_PID
 wait $SERVE_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

src/vertex_ai_server.cc

@@ -196,32 +196,38 @@ VertexAiAPIServer::Handle(evhtp_request_t* req)
       } else if (RE2::FullMatch(
                      redirect_endpoint, systemsharedmemory_regex_, &region,
                      &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_, &region,
                      &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;
       }
     }
   }