From 184fe95b097d9b737f78961719a0244a1a0821f8 Mon Sep 17 00:00:00 2001 From: namannagar89 Date: Thu, 23 Jul 2026 11:33:04 +0530 Subject: [PATCH] feat(ai-proxy): support Workload Identity / metadata server for GCP auth ai-proxy's GCP (Vertex AI) auth previously required a static service account JSON key: google-cloud-oauth signs a JWT with the private key and exchanges it for an access token. On GKE with Workload Identity this forces a long-lived key purely for APISIX. Add a metadata-server (Application Default Credentials) path: when no service account key is configured (neither `service_account_json` nor the `GCP_SERVICE_ACCOUNT` env), or when `use_metadata_server: true` is set, google-cloud-oauth fetches the access token from the GCE/GKE metadata server (`/computeMetadata/v1/instance/service-accounts/default/token`, `Metadata-Flavor: Google`). The host is overridable via `GCE_METADATA_HOST`. Existing token caching (max_ttl / expire_early_secs) is reused unchanged. - utils/google-cloud-oauth.lua: metadata-server token fetch + auto/explicit selection - ai-transport/auth.lua: pass use_metadata_server / metadata_host through - ai-proxy/schema.lua: add auth.gcp.use_metadata_server - t/plugin/ai-proxy-gcp-metadata.t: tests (default selection, token fetch, env host) - docs (en/zh) for ai-proxy and ai-proxy-multi Closes #13732 Co-Authored-By: Claude Opus 4.8 Signed-off-by: namannagar89 --- apisix/plugins/ai-proxy/schema.lua | 9 ++ apisix/plugins/ai-transport/auth.lua | 5 + apisix/utils/google-cloud-oauth.lua | 71 ++++++++++- docs/en/latest/plugins/ai-proxy-multi.md | 1 + docs/en/latest/plugins/ai-proxy.md | 1 + docs/zh/latest/plugins/ai-proxy-multi.md | 1 + docs/zh/latest/plugins/ai-proxy.md | 1 + t/plugin/ai-proxy-gcp-metadata.t | 147 +++++++++++++++++++++++ 8 files changed, 232 insertions(+), 4 deletions(-) create mode 100644 t/plugin/ai-proxy-gcp-metadata.t diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index a7833f0c1894..f61c20ff3ef0 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -44,6 +44,15 @@ local auth_schema = { type = "string", description = "GCP service account JSON content for authentication", }, + use_metadata_server = { + type = "boolean", + description = "Fetch the access token from the GCE/GKE metadata server " .. + "(Application Default Credentials / Workload Identity) " .. + "instead of a service account key. When no service account " .. + "key is configured, the metadata server is used " .. + "automatically.", + default = false, + }, max_ttl = { type = "integer", minimum = 1, diff --git a/apisix/plugins/ai-transport/auth.lua b/apisix/plugins/ai-transport/auth.lua index 147067e8f236..c6146fca31e0 100644 --- a/apisix/plugins/ai-transport/auth.lua +++ b/apisix/plugins/ai-transport/auth.lua @@ -50,6 +50,11 @@ function _M.fetch_gcp_access_token(ctx, name, gcp_conf) end auth_conf = conf end + -- When no service account key is provided, fall back to the metadata + -- server (Application Default Credentials / Workload Identity). This can + -- also be forced with use_metadata_server=true. + auth_conf.use_metadata_server = gcp_conf.use_metadata_server + auth_conf.metadata_host = gcp_conf.metadata_host local oauth = google_oauth.new(auth_conf) access_token = oauth:generate_access_token() if not access_token then diff --git a/apisix/utils/google-cloud-oauth.lua b/apisix/utils/google-cloud-oauth.lua index 340f5a6dfb39..e5555edf9a60 100644 --- a/apisix/utils/google-cloud-oauth.lua +++ b/apisix/utils/google-cloud-oauth.lua @@ -18,6 +18,7 @@ local core = require("apisix.core") local type = type local setmetatable = setmetatable +local os_getenv = os.getenv local ngx_update_time = ngx.update_time local ngx_time = ngx.time @@ -27,6 +28,15 @@ local http = require("resty.http") local jwt = require("resty.jwt") +-- Metadata server endpoint used by Application Default Credentials / GKE +-- Workload Identity to obtain an access token for the attached service account. +-- The host can be overridden with the standard GCE_METADATA_HOST env var (host +-- only, e.g. "metadata.google.internal" or "127.0.0.1:8080") or the +-- `metadata_host` config field (full base URL, mainly for testing). +local DEFAULT_METADATA_HOST = "http://metadata.google.internal" +local METADATA_TOKEN_PATH = "/computeMetadata/v1/instance/service-accounts/default/token" + + local function get_timestamp() ngx_update_time() return ngx_time() @@ -44,7 +54,52 @@ function _M.generate_access_token(self) end +local function set_access_token(self, res) + self.access_token = res.access_token + self.access_token_type = res.token_type + self.access_token_ttl = res.expires_in + self.access_token_expire_time = get_timestamp() + res.expires_in +end + + +-- Fetch an access token from the GCE/GKE metadata server. This is how +-- Application Default Credentials / Workload Identity authenticate: the token +-- is minted by the platform for the service account attached to the workload, +-- so no static service account key is required. +function _M.refresh_access_token_from_metadata_server(self) + local http_new = http.new() + local res, err = http_new:request_uri(self.metadata_host .. METADATA_TOKEN_PATH, { + method = "GET", + headers = { + ["Metadata-Flavor"] = "Google", + }, + }) + + if not res then + core.log.error("failed to fetch google access token from metadata server, ", err) + return + end + + if res.status ~= 200 then + core.log.error("failed to fetch google access token from metadata server: ", res.body) + return + end + + res, err = core.json.decode(res.body) + if not res then + core.log.error("failed to parse google metadata server response data: ", err) + return + end + + set_access_token(self, res) +end + + function _M.refresh_access_token(self) + if self.use_metadata_server then + return self:refresh_access_token_from_metadata_server() + end + local http_new = http.new() local res, err = http_new:request_uri(self.token_uri, { ssl_verify = self.ssl_verify, @@ -74,10 +129,7 @@ function _M.refresh_access_token(self) return end - self.access_token = res.access_token - self.access_token_type = res.token_type - self.access_token_ttl = res.expires_in - self.access_token_expire_time = get_timestamp() + res.expires_in + set_access_token(self, res) end @@ -100,6 +152,11 @@ end function _M.new(config, ssl_verify) + local metadata_host = config.metadata_host or os_getenv("GCE_METADATA_HOST") + if metadata_host and not core.string.has_prefix(metadata_host, "http") then + metadata_host = "http://" .. metadata_host + end + local oauth = { client_email = config.client_email, private_key = config.private_key, @@ -107,11 +164,17 @@ function _M.new(config, ssl_verify) token_uri = config.token_uri or "https://oauth2.googleapis.com/token", auth_uri = config.auth_uri or "https://accounts.google.com/o/oauth2/auth", entries_uri = config.entries_uri, + metadata_host = metadata_host or DEFAULT_METADATA_HOST, access_token = nil, access_token_type = nil, access_token_expire_time = 0, } + -- Use the metadata server (ADC / Workload Identity) when explicitly requested + -- or when no service account private key is available. This lets APISIX run + -- keyless on GKE with Workload Identity instead of mounting a static SA key. + oauth.use_metadata_server = config.use_metadata_server or (config.private_key == nil) + oauth.ssl_verify = ssl_verify if config.scope then diff --git a/docs/en/latest/plugins/ai-proxy-multi.md b/docs/en/latest/plugins/ai-proxy-multi.md index 88a1274ef491..080cd08f6b0f 100644 --- a/docs/en/latest/plugins/ai-proxy-multi.md +++ b/docs/en/latest/plugins/ai-proxy-multi.md @@ -87,6 +87,7 @@ When an instance's `provider` is set to `bedrock`, the Plugin expects requests i | instances.auth.query | object | False | | | Authentication query parameters. At least one of the `header` and `query` should be configured. | | instances.auth.gcp | object | False | | | Configuration for Google Cloud Platform (GCP) authentication. | | instances.auth.gcp.service_account_json | string | False | | | Content of the GCP service account JSON file. This can also be configured by setting the `GCP_SERVICE_ACCOUNT` environment variable. | +| instances.auth.gcp.use_metadata_server | boolean | False | false | | Fetch the access token from the GCE/GKE metadata server (Application Default Credentials / Workload Identity) instead of a service account key. When no service account key is configured (neither `service_account_json` nor `GCP_SERVICE_ACCOUNT`), the metadata server is used automatically. The metadata host can be overridden with the `GCE_METADATA_HOST` environment variable. This lets APISIX authenticate to Vertex AI keylessly on GKE with Workload Identity. | | instances.auth.gcp.max_ttl | integer | False | | minimum = 1 | Maximum TTL (in seconds) for caching the GCP access token. | | instances.auth.gcp.expire_early_secs| integer | False | 60 | minimum = 0 | Seconds to expire the access token before its actual expiration time to avoid edge cases. | | instances.auth.aws | object | False | | | AWS IAM credentials for SigV4 signing (Bedrock). Required when `provider` is `bedrock`. | diff --git a/docs/en/latest/plugins/ai-proxy.md b/docs/en/latest/plugins/ai-proxy.md index b93702ed00af..1f870c6653ba 100644 --- a/docs/en/latest/plugins/ai-proxy.md +++ b/docs/en/latest/plugins/ai-proxy.md @@ -76,6 +76,7 @@ When `provider` is set to `bedrock`, the Plugin expects requests in the [Bedrock | auth.query | object | False | | | Authentication query parameters. At least one of `header` or `query` must be configured. | | auth.gcp | object | False | | | Configuration for Google Cloud Platform (GCP) authentication. | | auth.gcp.service_account_json | string | False | | | Content of the GCP service account JSON file. This can also be configured by setting the `GCP_SERVICE_ACCOUNT` environment variable. | +| auth.gcp.use_metadata_server | boolean | False | false | | Fetch the access token from the GCE/GKE metadata server (Application Default Credentials / Workload Identity) instead of a service account key. When no service account key is configured (neither `service_account_json` nor `GCP_SERVICE_ACCOUNT`), the metadata server is used automatically. The metadata host can be overridden with the `GCE_METADATA_HOST` environment variable. This lets APISIX authenticate to Vertex AI keylessly on GKE with Workload Identity. | | auth.gcp.max_ttl | integer | False | | minimum = 1 | Maximum TTL (in seconds) for caching the GCP access token. | | auth.gcp.expire_early_secs | integer | False | 60 | minimum = 0 | Seconds to expire the access token before its actual expiration time to avoid edge cases. | | auth.aws | object | False | | | Configuration for AWS authentication. Required when `provider` is `bedrock`. | diff --git a/docs/zh/latest/plugins/ai-proxy-multi.md b/docs/zh/latest/plugins/ai-proxy-multi.md index 2e18dd9b9381..92ea064e3544 100644 --- a/docs/zh/latest/plugins/ai-proxy-multi.md +++ b/docs/zh/latest/plugins/ai-proxy-multi.md @@ -87,6 +87,7 @@ import TabItem from '@theme/TabItem'; | instances.auth.query | object | 否 | | | 身份验证查询参数。应配置 `header` 和 `query` 中的至少一个。 | | instances.auth.gcp | object | 否 | | | Google Cloud Platform (GCP) 身份验证配置。 | | instances.auth.gcp.service_account_json | string | 否 | | | GCP 服务帐户 JSON 文件的内容。也可以通过设置"GCP_SERVICE_ACCOUNT"环境变量来配置。 | +| instances.auth.gcp.use_metadata_server | boolean | 否 | false | | 从 GCE/GKE 元数据服务器（应用默认凭据 / Workload Identity）获取访问令牌，而不是使用服务帐户密钥。当未配置服务帐户密钥（既没有 `service_account_json` 也没有 `GCP_SERVICE_ACCOUNT`）时，将自动使用元数据服务器。元数据主机可通过 `GCE_METADATA_HOST` 环境变量覆盖。这使得 APISIX 可以在启用 Workload Identity 的 GKE 上无密钥地向 Vertex AI 进行身份验证。 | | instances.auth.gcp.max_ttl | integer | 否 | | minimum = 1 | 用于缓存 GCP 访问令牌的最大 TTL（以秒为单位）。 | | instances.auth.gcp.expire_early_secs| integer | 否 | 60 | minimum = 0 | 在访问令牌实际过期时间之前使其过期的秒数，以避免边缘情况。 | | instances.auth.aws | object | 否 | | | AWS 身份验证配置。当 `provider` 为 `bedrock` 时必填。 | diff --git a/docs/zh/latest/plugins/ai-proxy.md b/docs/zh/latest/plugins/ai-proxy.md index 62a1b4eae41e..2af0d8cec3d2 100644 --- a/docs/zh/latest/plugins/ai-proxy.md +++ b/docs/zh/latest/plugins/ai-proxy.md @@ -76,6 +76,7 @@ import TabItem from '@theme/TabItem'; | auth.query | object | 否 | | | 身份验证查询参数。必须配置 `header` 或 `query` 中的至少一个。 | | auth.gcp | object | 否 | | | Google Cloud Platform (GCP) 身份验证配置。 | | auth.gcp.service_account_json | string | 否 | | | GCP 服务账号 JSON 文件内容。也可以通过设置 `GCP_SERVICE_ACCOUNT` 环境变量来配置。 | +| auth.gcp.use_metadata_server | boolean | 否 | false | | 从 GCE/GKE 元数据服务器（应用默认凭据 / Workload Identity）获取访问令牌，而不是使用服务账号密钥。当未配置服务账号密钥（既没有 `service_account_json` 也没有 `GCP_SERVICE_ACCOUNT`）时，将自动使用元数据服务器。元数据主机可通过 `GCE_METADATA_HOST` 环境变量覆盖。这使得 APISIX 可以在启用 Workload Identity 的 GKE 上无密钥地向 Vertex AI 进行身份验证。 | | auth.gcp.max_ttl | integer | 否 | | ≥ 1 | GCP 访问令牌缓存的最大 TTL（秒）。 | | auth.gcp.expire_early_secs | integer | 否 | 60 | ≥ 0 | 在访问令牌实际过期之前提前过期的秒数，以避免边缘情况。 | | auth.aws | object | 否 | | | AWS 身份验证配置。当 `provider` 为 `bedrock` 时必填。 | diff --git a/t/plugin/ai-proxy-gcp-metadata.t b/t/plugin/ai-proxy-gcp-metadata.t new file mode 100644 index 000000000000..63368b6c03ee --- /dev/null +++ b/t/plugin/ai-proxy-gcp-metadata.t @@ -0,0 +1,147 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +use t::APISIX 'no_plan'; + +log_level("info"); +repeat_each(1); +no_long_string(); +no_root_location(); + +run_tests(); + +__DATA__ + +=== TEST 1: no service account key -> metadata server (Workload Identity) used by default +--- config + location /t { + content_by_lua_block { + local google_oauth = require("apisix.utils.google-cloud-oauth") + local oauth = google_oauth.new({}) + ngx.say("use_metadata_server: ", tostring(oauth.use_metadata_server)) + } + } +--- request +GET /t +--- response_body +use_metadata_server: true + + + +=== TEST 2: service account key present -> JWT flow, metadata server not used +--- config + location /t { + content_by_lua_block { + local google_oauth = require("apisix.utils.google-cloud-oauth") + local oauth = google_oauth.new({ + client_email = "sa@example.iam.gserviceaccount.com", + private_key = "-----BEGIN PRIVATE KEY-----\nfake\n-----END PRIVATE KEY-----\n", + }) + ngx.say("use_metadata_server: ", tostring(oauth.use_metadata_server)) + } + } +--- request +GET /t +--- response_body +use_metadata_server: false + + + +=== TEST 3: fetch access token from the metadata server +--- config + location = /computeMetadata/v1/instance/service-accounts/default/token { + content_by_lua_block { + local flavor = ngx.req.get_headers()["Metadata-Flavor"] + if flavor ~= "Google" then + ngx.status = 403 + ngx.say("missing Metadata-Flavor header") + return + end + ngx.header["Content-Type"] = "application/json" + ngx.say('{"access_token":"ya29.metadata-token",' + .. '"expires_in":3599,"token_type":"Bearer"}') + } + } + location /t { + content_by_lua_block { + local google_oauth = require("apisix.utils.google-cloud-oauth") + local oauth = google_oauth.new({ + metadata_host = "http://127.0.0.1:" .. ngx.var.server_port + }) + local token = oauth:generate_access_token() + ngx.say("token: ", token) + ngx.say("token_type: ", oauth.access_token_type) + ngx.say("ttl: ", oauth.access_token_ttl) + } + } +--- request +GET /t +--- response_body +token: ya29.metadata-token +token_type: Bearer +ttl: 3599 + + + +=== TEST 4: explicit use_metadata_server flag fetches from the metadata server +--- config + location = /computeMetadata/v1/instance/service-accounts/default/token { + content_by_lua_block { + ngx.header["Content-Type"] = "application/json" + ngx.say('{"access_token":"ya29.explicit",' + .. '"expires_in":100,"token_type":"Bearer"}') + } + } + location /t { + content_by_lua_block { + local google_oauth = require("apisix.utils.google-cloud-oauth") + local oauth = google_oauth.new({ + use_metadata_server = true, + metadata_host = "http://127.0.0.1:" .. ngx.var.server_port + }) + ngx.say("token: ", oauth:generate_access_token()) + } + } +--- request +GET /t +--- response_body +token: ya29.explicit + + + +=== TEST 5: metadata server host can be set via GCE_METADATA_HOST env +--- main_config +env GCE_METADATA_HOST=127.0.0.1:1984; +--- config + location = /computeMetadata/v1/instance/service-accounts/default/token { + content_by_lua_block { + ngx.header["Content-Type"] = "application/json" + ngx.say('{"access_token":"ya29.from-env",' + .. '"expires_in":42,"token_type":"Bearer"}') + } + } + location /t { + content_by_lua_block { + local google_oauth = require("apisix.utils.google-cloud-oauth") + local oauth = google_oauth.new({}) + ngx.say("token: ", oauth:generate_access_token()) + } + } +--- request +GET /t +--- response_body +token: ya29.from-env