fix: Update Vertex AI endpoint to use HTTP restriction config

docs/customization_guide/deploy.md

@@ -1,5 +1,5 @@
 <!--
-# Copyright (c) 2020-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -226,6 +226,10 @@ for more information.
 > [!Note]
 > Restricting access can be used to limit exposure to model
 > control APIs to trusted users.
+> When Vertex AI endpoint support is enabled, this setting also
+> applies to redirected Vertex AI requests.
+> If Triton is built without `TRITON_ENABLE_HTTP`, Vertex AI uses
+> default unrestricted API settings.
 
 ##### `--allow-sagemaker <boolean> default False`
 

docs/customization_guide/inference_protocols.md

@@ -163,6 +163,11 @@ protocol group (GRPC) or restricted API group (HTTP):
 --http-restricted-api=<API_1>,API_2>,...:<restricted-key>=<restricted-value>
 ```
 
+When Vertex AI endpoint support is enabled, `--http-restricted-api`
+also applies to redirected Vertex AI requests.
+If Triton is built without `TRITON_ENABLE_HTTP`, Vertex AI falls back
+to default unrestricted API settings.
+
 The option can be specified multiple times to specifies multiple groups of
 protocols or APIs with different restriction settings.
 

qa/L0_vertex_ai/test.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -44,7 +44,7 @@ export CUDA_VISIBLE_DEVICES=0
 
 RET=0
 
-rm -rf multi_models single_model
+rm -rf multi_models single_model restricted_single_model
 rm -f *.log
 rm -f *.out
 
@@ -68,6 +68,11 @@ mkdir multi_models && \
 mkdir single_model && \
     cp -r multi_models/addsub single_model/.
 
+# Set up single-model Python repository used by restricted API regression
+mkdir -p restricted_single_model/identity_fp32/1 && \
+    cp ../python_models/identity_fp32/config.pbtxt restricted_single_model/identity_fp32/ && \
+    cp ../python_models/identity_fp32/model.py restricted_single_model/identity_fp32/1/
+
 # Use Vertex AI's health endpoint to check server status
 # Wait until server health endpoint shows ready. Sets WAIT_RET to 0 on
 # success, 1 on failure
@@ -134,7 +139,7 @@ if [ "$SERVER_PID" == "0" ]; then
     exit 1
 fi
 kill $SERVER_PID
-wait $SERVE_PID
+wait $SERVER_PID
 set +e
 # Expect no message regarding Vertex AI as it is disabled
 grep "failed to start Vertex AI service" $SERVER_LOG
@@ -157,7 +162,7 @@ if [ "$SERVER_PID" == "0" ]; then
     exit 1
 fi
 kill $SERVER_PID
-wait $SERVE_PID
+wait $SERVER_PID
 set +e
 grep "Started Vertex AI HTTPService at" $SERVER_LOG
 if [ $? -ne 0 ]; then
@@ -186,7 +191,7 @@ if [ "$SERVER_PID" == "0" ]; then
     exit 1
 fi
 kill $SERVER_PID
-wait $SERVE_PID
+wait $SERVER_PID
 set +e
 grep "Started Vertex AI HTTPService at" $SERVER_LOG
 if [ $? -ne 0 ]; then
@@ -220,7 +225,7 @@ if [ "$SERVER_PID" == "0" ]; then
     exit 1
 fi
 kill $SERVER_PID
-wait $SERVE_PID
+wait $SERVER_PID
 set +e
 # Expect no message regarding Vertex AI as it is disabled
 grep "failed to start Vertex AI service" $SERVER_LOG
@@ -331,7 +336,7 @@ fi
 set -e
 
 kill $SERVER_PID
-wait $SERVE_PID
+wait $SERVER_PID
 
 #
 # AIP_STORAGE_URI / AIP_HTTP_PORT
@@ -370,7 +375,7 @@ fi
 set -e
 
 kill $SERVER_PID
-wait $SERVE_PID
+wait $SERVER_PID
 
 #
 # default model
@@ -453,7 +458,7 @@ else
 fi
 set -e
 kill $SERVER_PID
-wait $SERVE_PID
+wait $SERVER_PID
 
 # Test default model as well as multi model
 SERVER_ARGS="--vertex-ai-default-model=addsub"
@@ -712,7 +717,133 @@ fi
 set -e
 
 kill $SERVER_PID
-wait $SERVE_PID
+wait $SERVER_PID
+
+#
+# Restricted API regression for Vertex redirect
+#
+unset_vertex_variables
+export AIP_PREDICT_ROUTE="/predict"
+export AIP_HEALTH_ROUTE="/health"
+
+SERVER_LOG="vertex_restricted_api_testing_server.log"
+SERVER_ARGS="--log-verbose=1 --allow-vertex-ai=true \
+  --model-repository=restricted_single_model \
+  --vertex-ai-default-model=identity_fp32 \
+  --http-restricted-api=metadata,model-config,model-repository,statistics,shared-memory:X-Vertex-Restricted=secret"
+
+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
+
+# Baseline infer remains available without restricted header
+set +e
+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":[42.0]}],"outputs":[{"name":"OUTPUT0"}]}' localhost:8080/predict`
+if [ "$code" != "200" ]; then
+    cat ./curl.out
+    echo -e "\n***\n*** Failed. Expected /predict inference succeeds without restricted header\n***"
+    RET=1
+else
+    grep "OUTPUT0" ./curl.out
+    if [ $? -ne 0 ]; then
+        cat ./curl.out
+        echo -e "\n***\n*** Failed. Expected inference output is returned\n***"
+        RET=1
+    fi
+fi
+set -e
+
+# Redirected management 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"
+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
+    if [ "$code" != "403" ]; then
+        cat ./curl.out
+        echo -e "\n***\n*** Failed. Expected ${redirect_endpoint} is restricted\n***"
+        RET=1
+    fi
+done
+set -e
+
+# Restricted header should allow 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"
+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
+    fi
+done
+
+# Wrong restricted header should reject the request
+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"
+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: invalid" localhost:8080/predict`
+    else
+        code=`curl -s -w %{http_code} -o ./curl.out -X POST -H "X-Vertex-Ai-Triton-Redirect: ${redirect_endpoint}" \
+        -H "X-Vertex-Restricted: invalid" -H "Content-Type: application/json" -d '{"key":"test_shm","byte_size":1024}' localhost:8080/predict`
+    fi
+    if [ "$code" != "403" ]; then
+        cat ./curl.out
+        echo -e "\n***\n*** Failed. Expected invalid restricted header value is rejected for ${redirect_endpoint}\n***"
+        RET=1
+    fi
+done
+set -e
+
+kill $SERVER_PID
+wait $SERVER_PID
 
 if [ $RET -eq 0 ]; then
     echo -e "\n***\n*** Test Passed\n***"

src/main.cc

@@ -1,4 +1,4 @@
-// Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -204,10 +204,23 @@ StartVertexAiService(
     triton::server::TraceManager* trace_manager,
     const std::shared_ptr<triton::server::SharedMemoryManager>& shm_manager)
 {
+  size_t max_input_size = triton::server::HTTP_DEFAULT_MAX_INPUT_SIZE;
+  triton::server::RestrictedFeatures restricted_apis{};
+#ifdef TRITON_ENABLE_HTTP
+  // Reuse HTTP server settings for Vertex AI endpoint behavior. In
+  // particular, --http-restricted-api and --http-max-input-size also apply
+  // to Vertex AI requests. Without TRITON_ENABLE_HTTP, Vertex AI falls back to
+  // default input size and unrestricted APIs (no command-line configuration for
+  // restricted APIs).
+  max_input_size = g_triton_params.http_max_input_size_;
+  restricted_apis = g_triton_params.http_restricted_apis_;
+#endif  // TRITON_ENABLE_HTTP
+
   TRITONSERVER_Error* err = triton::server::VertexAiAPIServer::Create(
       server, trace_manager, shm_manager, g_triton_params.vertex_ai_port_,
       g_triton_params.vertex_ai_address_, g_triton_params.vertex_ai_thread_cnt_,
-      g_triton_params.vertex_ai_default_model_, service);
+      max_input_size, restricted_apis, g_triton_params.vertex_ai_default_model_,
+      service);
   if (err == nullptr) {
     err = (*service)->Start();
   }

src/vertex_ai_server.cc

@@ -1,4 +1,4 @@
-// Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -42,10 +42,12 @@ VertexAiAPIServer::VertexAiAPIServer(
     const std::shared_ptr<SharedMemoryManager>& shm_manager, const int32_t port,
     const std::string address, const int thread_cnt,
     const std::string& prediction_route, const std::string& health_route,
-    const std::string& default_model_name)
+    const std::string& default_model_name, const size_t max_input_size,
+    const RestrictedFeatures& restricted_apis)
     : HTTPAPIServer(
           server, trace_manager, shm_manager, port, false /* reuse_port */,
-          address, "" /* header_forward_pattern */, thread_cnt),
+          address, "" /* header_forward_pattern */, thread_cnt, max_input_size,
+          restricted_apis),
       prediction_regex_(prediction_route), health_regex_(health_route),
       health_mode_("ready"), model_name_(default_model_name),
       model_version_str_("")
@@ -269,6 +271,7 @@ VertexAiAPIServer::Create(
     triton::server::TraceManager* trace_manager,
     const std::shared_ptr<SharedMemoryManager>& shm_manager, const int32_t port,
     const std::string address, const int thread_cnt,
+    const size_t max_input_size, const RestrictedFeatures& restricted_apis,
     std::string default_model_name, std::unique_ptr<HTTPServer>* http_server)
 {
   auto predict_route = GetEnvironmentVariableOrDefault("AIP_PREDICT_ROUTE", "");
@@ -344,7 +347,8 @@ VertexAiAPIServer::Create(
 
   http_server->reset(new VertexAiAPIServer(
       server, trace_manager, shm_manager, port, address, thread_cnt,
-      predict_route, health_route, default_model_name));
+      predict_route, health_route, default_model_name, max_input_size,
+      restricted_apis));
 
   const std::string addr = address + ":" + std::to_string(port);
   LOG_INFO << "Started Vertex AI HTTPService at " << addr;

src/vertex_ai_server.h

@@ -1,4 +1,4 @@
-// Copyright 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -37,6 +37,7 @@ class VertexAiAPIServer : public HTTPAPIServer {
       triton::server::TraceManager* trace_manager,
       const std::shared_ptr<SharedMemoryManager>& smb_manager,
       const int32_t port, const std::string address, const int thread_cnt,
+      const size_t max_input_size, const RestrictedFeatures& restricted_apis,
       std::string default_model_name,
       std::unique_ptr<HTTPServer>* vertex_ai_server);
 
@@ -47,7 +48,8 @@ class VertexAiAPIServer : public HTTPAPIServer {
       const std::shared_ptr<SharedMemoryManager>& shm_manager,
       const int32_t port, const std::string address, const int thread_cnt,
       const std::string& prediction_route, const std::string& health_route,
-      const std::string& default_model_name);
+      const std::string& default_model_name, const size_t max_input_size,
+      const RestrictedFeatures& restricted_apis);
 
   void Handle(evhtp_request_t* req) override;