From fc600f204a46ec060c629a3f38f2d5ec69cccc5c Mon Sep 17 00:00:00 2001 From: CZJCC Date: Wed, 7 Jan 2026 19:39:20 +0800 Subject: [PATCH] =?UTF-8?q?feat(ai-proxy):=20add=20Bearer=20Token=20authen?= =?UTF-8?q?tication=20support=20for=20Bedrock=20p=E2=80=A6=20(#3305)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- plugins/wasm-go/extensions/ai-proxy/README.md | 52 +- .../wasm-go/extensions/ai-proxy/README_EN.md | 44 +- plugins/wasm-go/extensions/ai-proxy/go.mod | 2 +- plugins/wasm-go/extensions/ai-proxy/go.sum | 2 + .../wasm-go/extensions/ai-proxy/main_test.go | 8 + .../extensions/ai-proxy/provider/bedrock.go | 29 +- .../extensions/ai-proxy/test/bedrock.go | 527 ++++++++++++++++++ 7 files changed, 635 insertions(+), 29 deletions(-) create mode 100644 plugins/wasm-go/extensions/ai-proxy/test/bedrock.go diff --git a/plugins/wasm-go/extensions/ai-proxy/README.md b/plugins/wasm-go/extensions/ai-proxy/README.md index 83d468d2e..499ae5444 100644 --- a/plugins/wasm-go/extensions/ai-proxy/README.md +++ b/plugins/wasm-go/extensions/ai-proxy/README.md @@ -322,23 +322,31 @@ Google Vertex AI 所对应的 type 为 vertex。它特有的配置字段如下 #### AWS Bedrock -AWS Bedrock 所对应的 type 为 bedrock。它特有的配置字段如下： +AWS Bedrock 所对应的 type 为 bedrock。它支持两种认证方式： -| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 | -|---------------------------|--------|------|-----|------------------------------| -| `modelVersion` | string | 非必填 | - | 用于指定 Triton Server 中 model version | -| `tritonDomain` | string | 非必填 | - | Triton Server 部署的指定请求 Domain | +1. **AWS Signature V4 认证**：使用 `awsAccessKey` 和 `awsSecretKey` 进行 AWS 标准签名认证 +2. **Bearer Token 认证**：使用 `apiTokens` 配置 AWS Bearer Token（适用于 IAM Identity Center 等场景） + +**注意**：两种认证方式二选一，如果同时配置了 `apiTokens`，将优先使用 Bearer Token 认证方式。 + +它特有的配置字段如下： + +| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 | +|---------------------------|---------------|-------------------|-------|---------------------------------------------------| +| `apiTokens` | array of string | 与 ak/sk 二选一 | - | AWS Bearer Token，用于 Bearer Token 认证方式 | +| `awsAccessKey` | string | 与 apiTokens 二选一 | - | AWS Access Key，用于 AWS Signature V4 认证 | +| `awsSecretKey` | string | 与 apiTokens 二选一 | - | AWS Secret Access Key，用于 AWS Signature V4 认证 | +| `awsRegion` | string | 必填 | - | AWS 区域，例如：us-east-1 | +| `bedrockAdditionalFields` | map | 非必填 | - | Bedrock 额外模型请求参数 | #### NVIDIA Triton Interference Server NVIDIA Triton Interference Server 所对应的 type 为 triton。它特有的配置字段如下： -| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 | -|---------------------------|--------|------|-----|------------------------------| -| `awsAccessKey` | string | 必填 | - | AWS Access Key，用于身份认证 | -| `awsSecretKey` | string | 必填 | - | AWS Secret Access Key，用于身份认证 | -| `awsRegion` | string | 必填 | - | AWS 区域，例如：us-east-1 | -| `bedrockAdditionalFields` | map | 非必填 | - | Bedrock 额外模型请求参数 | +| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 | +|----------------------|--------|--------|-------|------------------------------------------| +| `tritonModelVersion` | string | 非必填 | - | 用于指定 Triton Server 中 model version | +| `tritonDomain` | string | 非必填 | - | Triton Server 部署的指定请求 Domain | ## 用法示例 @@ -2011,6 +2019,10 @@ provider: ### 使用 OpenAI 协议代理 AWS Bedrock 服务 +AWS Bedrock 支持两种认证方式： + +#### 方式一：使用 AWS Access Key/Secret Key 认证（AWS Signature V4） + **配置信息** ```yaml @@ -2018,7 +2030,21 @@ provider: type: bedrock awsAccessKey: "YOUR_AWS_ACCESS_KEY_ID" awsSecretKey: "YOUR_AWS_SECRET_ACCESS_KEY" - awsRegion: "YOUR_AWS_REGION" + awsRegion: "us-east-1" + bedrockAdditionalFields: + top_k: 200 +``` + +#### 方式二：使用 Bearer Token 认证（适用于 IAM Identity Center 等场景） + +**配置信息** + +```yaml +provider: + type: bedrock + apiTokens: + - "YOUR_AWS_BEARER_TOKEN" + awsRegion: "us-east-1" bedrockAdditionalFields: top_k: 200 ``` @@ -2027,7 +2053,7 @@ provider: ```json { - "model": "arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-5-haiku-20241022-v1:0", + "model": "us.anthropic.claude-3-5-haiku-20241022-v1:0", "messages": [ { "role": "user", diff --git a/plugins/wasm-go/extensions/ai-proxy/README_EN.md b/plugins/wasm-go/extensions/ai-proxy/README_EN.md index 0064ec851..b02812b24 100644 --- a/plugins/wasm-go/extensions/ai-proxy/README_EN.md +++ b/plugins/wasm-go/extensions/ai-proxy/README_EN.md @@ -268,14 +268,22 @@ For Vertex, the corresponding `type` is `vertex`. Its unique configuration field #### AWS Bedrock -For AWS Bedrock, the corresponding `type` is `bedrock`. Its unique configuration field is: +For AWS Bedrock, the corresponding `type` is `bedrock`. It supports two authentication methods: -| Name | Data Type | Requirement | Default | Description | -|---------------------------|-----------|-------------|---------|---------------------------------------------------------| -| `awsAccessKey` | string | Required | - | AWS Access Key used for authentication | -| `awsSecretKey` | string | Required | - | AWS Secret Access Key used for authentication | -| `awsRegion` | string | Required | - | AWS region, e.g., us-east-1 | -| `bedrockAdditionalFields` | map | Optional | - | Additional inference parameters that the model supports | +1. **AWS Signature V4 Authentication**: Uses `awsAccessKey` and `awsSecretKey` for standard AWS signature authentication +2. **Bearer Token Authentication**: Uses `apiTokens` to configure AWS Bearer Token (suitable for IAM Identity Center and similar scenarios) + +**Note**: Choose one of the two authentication methods. If `apiTokens` is configured, Bearer Token authentication will be used preferentially. + +Its unique configuration fields are: + +| Name | Data Type | Requirement | Default | Description | +|---------------------------|-----------------|--------------------------|---------|-------------------------------------------------------------------| +| `apiTokens` | array of string | Either this or ak/sk | - | AWS Bearer Token for Bearer Token authentication | +| `awsAccessKey` | string | Either this or apiTokens | - | AWS Access Key for AWS Signature V4 authentication | +| `awsSecretKey` | string | Either this or apiTokens | - | AWS Secret Access Key for AWS Signature V4 authentication | +| `awsRegion` | string | Required | - | AWS region, e.g., us-east-1 | +| `bedrockAdditionalFields` | map | Optional | - | Additional inference parameters that the model supports | ## Usage Examples @@ -1779,13 +1787,31 @@ provider: ``` ### Utilizing OpenAI Protocol Proxy for AWS Bedrock Services + +AWS Bedrock supports two authentication methods: + +#### Method 1: Using AWS Access Key/Secret Key Authentication (AWS Signature V4) + **Configuration Information** ```yaml provider: type: bedrock awsAccessKey: "YOUR_AWS_ACCESS_KEY_ID" awsSecretKey: "YOUR_AWS_SECRET_ACCESS_KEY" - awsRegion: "YOUR_AWS_REGION" + awsRegion: "us-east-1" + bedrockAdditionalFields: + top_k: 200 +``` + +#### Method 2: Using Bearer Token Authentication (suitable for IAM Identity Center and similar scenarios) + +**Configuration Information** +```yaml +provider: + type: bedrock + apiTokens: + - "YOUR_AWS_BEARER_TOKEN" + awsRegion: "us-east-1" bedrockAdditionalFields: top_k: 200 ``` @@ -1793,7 +1819,7 @@ provider: **Request Example** ```json { - "model": "arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-5-haiku-20241022-v1:0", + "model": "us.anthropic.claude-3-5-haiku-20241022-v1:0", "messages": [ { "role": "user", diff --git a/plugins/wasm-go/extensions/ai-proxy/go.mod b/plugins/wasm-go/extensions/ai-proxy/go.mod index c21dc8be2..27e50d663 100644 --- a/plugins/wasm-go/extensions/ai-proxy/go.mod +++ b/plugins/wasm-go/extensions/ai-proxy/go.mod @@ -8,7 +8,7 @@ toolchain go1.24.4 require ( github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20251103120604-77e9cce339d2 - github.com/higress-group/wasm-go v1.0.7-0.20251209122854-7e766df5675c + github.com/higress-group/wasm-go v1.0.9-0.20251226032831-95da539a1ec7 github.com/stretchr/testify v1.9.0 github.com/tidwall/gjson v1.18.0 ) diff --git a/plugins/wasm-go/extensions/ai-proxy/go.sum b/plugins/wasm-go/extensions/ai-proxy/go.sum index 9d45243f7..c917827c9 100644 --- a/plugins/wasm-go/extensions/ai-proxy/go.sum +++ b/plugins/wasm-go/extensions/ai-proxy/go.sum @@ -6,6 +6,8 @@ github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20251103120604-77e9cce339d2 h1 github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20251103120604-77e9cce339d2/go.mod h1:tRI2LfMudSkKHhyv1uex3BWzcice2s/l8Ah8axporfA= github.com/higress-group/wasm-go v1.0.7-0.20251209122854-7e766df5675c h1:DdVPyaMHSYBqO5jwB9Wl3PqsBGIf4u29BHMI0uIVB1Y= github.com/higress-group/wasm-go v1.0.7-0.20251209122854-7e766df5675c/go.mod h1:uKVYICbRaxTlKqdm8E0dpjbysxM8uCPb9LV26hF3Km8= +github.com/higress-group/wasm-go v1.0.9-0.20251226032831-95da539a1ec7 h1:ddAkPFIIf6isVQNymS6+X6QO51/WV0Af4Afb9a2z9TE= +github.com/higress-group/wasm-go v1.0.9-0.20251226032831-95da539a1ec7/go.mod h1:uKVYICbRaxTlKqdm8E0dpjbysxM8uCPb9LV26hF3Km8= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= diff --git a/plugins/wasm-go/extensions/ai-proxy/main_test.go b/plugins/wasm-go/extensions/ai-proxy/main_test.go index f5a5c0f20..c7accee46 100644 --- a/plugins/wasm-go/extensions/ai-proxy/main_test.go +++ b/plugins/wasm-go/extensions/ai-proxy/main_test.go @@ -128,3 +128,11 @@ func TestGeneric(t *testing.T) { test.RunGenericOnHttpRequestHeadersTests(t) test.RunGenericOnHttpRequestBodyTests(t) } + +func TestBedrock(t *testing.T) { + test.RunBedrockParseConfigTests(t) + test.RunBedrockOnHttpRequestHeadersTests(t) + test.RunBedrockOnHttpRequestBodyTests(t) + test.RunBedrockOnHttpResponseHeadersTests(t) + test.RunBedrockOnHttpResponseBodyTests(t) +} diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/bedrock.go b/plugins/wasm-go/extensions/ai-proxy/provider/bedrock.go index 6dbf4f94f..85b341a2a 100644 --- a/plugins/wasm-go/extensions/ai-proxy/provider/bedrock.go +++ b/plugins/wasm-go/extensions/ai-proxy/provider/bedrock.go @@ -43,8 +43,11 @@ const ( type bedrockProviderInitializer struct{} func (b *bedrockProviderInitializer) ValidateConfig(config *ProviderConfig) error { - if len(config.awsAccessKey) == 0 || len(config.awsSecretKey) == 0 { - return errors.New("missing bedrock access authentication parameters") + hasAkSk := len(config.awsAccessKey) > 0 && len(config.awsSecretKey) > 0 + hasApiToken := len(config.apiTokens) > 0 + + if !hasAkSk && !hasApiToken { + return errors.New("missing bedrock access authentication parameters: either apiTokens or (awsAccessKey + awsSecretKey) is required") } if len(config.awsRegion) == 0 { return errors.New("missing bedrock region parameters") @@ -634,6 +637,13 @@ func (b *bedrockProvider) OnRequestHeaders(ctx wrapper.HttpContext, apiName ApiN func (b *bedrockProvider) TransformRequestHeaders(ctx wrapper.HttpContext, apiName ApiName, headers http.Header) { util.OverwriteRequestHostHeader(headers, fmt.Sprintf(bedrockDefaultDomain, b.config.awsRegion)) + + // If apiTokens is configured, set Bearer token authentication here + // This follows the same pattern as other providers (qwen, zhipuai, etc.) + // AWS SigV4 authentication is handled in setAuthHeaders because it requires the request body + if len(b.config.apiTokens) > 0 { + util.OverwriteRequestAuthorizationHeader(headers, "Bearer "+b.config.GetApiTokenInUse(ctx)) + } } func (b *bedrockProvider) OnRequestBody(ctx wrapper.HttpContext, apiName ApiName, body []byte) (types.Action, error) { @@ -659,18 +669,18 @@ func (b *bedrockProvider) TransformResponseBody(ctx wrapper.HttpContext, apiName case ApiNameChatCompletion: return b.onChatCompletionResponseBody(ctx, body) case ApiNameImageGeneration: - return b.onImageGenerationResponseBody(ctx, body) + return b.onImageGenerationResponseBody(body) } return nil, errUnsupportedApiName } -func (b *bedrockProvider) onImageGenerationResponseBody(ctx wrapper.HttpContext, body []byte) ([]byte, error) { +func (b *bedrockProvider) onImageGenerationResponseBody(body []byte) ([]byte, error) { bedrockResponse := &bedrockImageGenerationResponse{} if err := json.Unmarshal(body, bedrockResponse); err != nil { log.Errorf("unable to unmarshal bedrock image gerneration response: %v", err) return nil, fmt.Errorf("unable to unmarshal bedrock image generation response: %v", err) } - response := b.buildBedrockImageGenerationResponse(ctx, bedrockResponse) + response := b.buildBedrockImageGenerationResponse(bedrockResponse) return json.Marshal(response) } @@ -710,7 +720,7 @@ func (b *bedrockProvider) buildBedrockImageGenerationRequest(origRequest *imageG return requestBytes, err } -func (b *bedrockProvider) buildBedrockImageGenerationResponse(ctx wrapper.HttpContext, bedrockResponse *bedrockImageGenerationResponse) *imageGenerationResponse { +func (b *bedrockProvider) buildBedrockImageGenerationResponse(bedrockResponse *bedrockImageGenerationResponse) *imageGenerationResponse { data := make([]imageGenerationData, len(bedrockResponse.Images)) for i, image := range bedrockResponse.Images { data[i] = imageGenerationData{ @@ -1138,6 +1148,13 @@ func chatMessage2BedrockMessage(chatMessage chatMessage) bedrockMessage { } func (b *bedrockProvider) setAuthHeaders(body []byte, headers http.Header) { + // Bearer token authentication is already set in TransformRequestHeaders + // This function only handles AWS SigV4 authentication which requires the request body + if len(b.config.apiTokens) > 0 { + return + } + + // Use AWS Signature V4 authentication t := time.Now().UTC() amzDate := t.Format("20060102T150405Z") dateStamp := t.Format("20060102") diff --git a/plugins/wasm-go/extensions/ai-proxy/test/bedrock.go b/plugins/wasm-go/extensions/ai-proxy/test/bedrock.go new file mode 100644 index 000000000..766e11e06 --- /dev/null +++ b/plugins/wasm-go/extensions/ai-proxy/test/bedrock.go @@ -0,0 +1,527 @@ +package test + +import ( + "encoding/json" + "testing" + + "github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types" + "github.com/higress-group/wasm-go/pkg/test" + "github.com/stretchr/testify/require" +) + +// Test config: Basic Bedrock config with AWS Access Key/Secret Key (AWS Signature V4) +var basicBedrockConfig = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "bedrock", + "awsAccessKey": "test-ak-for-unit-test", + "awsSecretKey": "test-sk-for-unit-test", + "awsRegion": "us-east-1", + "modelMapping": map[string]string{ + "*": "anthropic.claude-3-5-haiku-20241022-v1:0", + }, + }, + }) + return data +}() + +// Test config: Bedrock config with Bearer Token authentication +var bedrockApiTokenConfig = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "bedrock", + "apiTokens": []string{ + "test-token-for-unit-test", + }, + "awsRegion": "us-east-1", + "modelMapping": map[string]string{ + "*": "anthropic.claude-3-5-haiku-20241022-v1:0", + }, + }, + }) + return data +}() + +// Test config: Bedrock config with multiple Bearer Tokens +var bedrockMultiTokenConfig = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "bedrock", + "apiTokens": []string{ + "test-token-1-for-unit-test", + "test-token-2-for-unit-test", + }, + "awsRegion": "us-west-2", + "modelMapping": map[string]string{ + "gpt-4": "anthropic.claude-3-opus-20240229-v1:0", + "*": "anthropic.claude-3-haiku-20240307-v1:0", + }, + }, + }) + return data +}() + +// Test config: Bedrock config with additional fields +var bedrockWithAdditionalFieldsConfig = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "bedrock", + "awsAccessKey": "test-ak-for-unit-test", + "awsSecretKey": "test-sk-for-unit-test", + "awsRegion": "us-east-1", + "bedrockAdditionalFields": map[string]interface{}{ + "top_k": 200, + }, + "modelMapping": map[string]string{ + "*": "anthropic.claude-3-5-haiku-20241022-v1:0", + }, + }, + }) + return data +}() + +// Test config: Invalid config - missing both apiTokens and ak/sk +var bedrockInvalidConfigMissingAuth = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "bedrock", + "awsRegion": "us-east-1", + "modelMapping": map[string]string{ + "*": "anthropic.claude-3-5-haiku-20241022-v1:0", + }, + }, + }) + return data +}() + +// Test config: Invalid config - missing region +var bedrockInvalidConfigMissingRegion = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "bedrock", + "apiTokens": []string{ + "test-token-for-unit-test", + }, + "modelMapping": map[string]string{ + "*": "anthropic.claude-3-5-haiku-20241022-v1:0", + }, + }, + }) + return data +}() + +// Test config: Invalid config - only has access key without secret key +var bedrockInvalidConfigPartialAkSk = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "bedrock", + "awsAccessKey": "test-ak-for-unit-test", + "awsRegion": "us-east-1", + "modelMapping": map[string]string{ + "*": "anthropic.claude-3-5-haiku-20241022-v1:0", + }, + }, + }) + return data +}() + +func RunBedrockParseConfigTests(t *testing.T) { + test.RunGoTest(t, func(t *testing.T) { + // Test basic Bedrock config with AWS Signature V4 authentication + t.Run("basic bedrock config with ak/sk", func(t *testing.T) { + host, status := test.NewTestHost(basicBedrockConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + config, err := host.GetMatchConfig() + require.NoError(t, err) + require.NotNil(t, config) + }) + + // Test Bedrock config with Bearer Token authentication + t.Run("bedrock config with api token", func(t *testing.T) { + host, status := test.NewTestHost(bedrockApiTokenConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + config, err := host.GetMatchConfig() + require.NoError(t, err) + require.NotNil(t, config) + }) + + // Test Bedrock config with multiple tokens + t.Run("bedrock config with multiple tokens", func(t *testing.T) { + host, status := test.NewTestHost(bedrockMultiTokenConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + config, err := host.GetMatchConfig() + require.NoError(t, err) + require.NotNil(t, config) + }) + + // Test Bedrock config with additional fields + t.Run("bedrock config with additional fields", func(t *testing.T) { + host, status := test.NewTestHost(bedrockWithAdditionalFieldsConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + config, err := host.GetMatchConfig() + require.NoError(t, err) + require.NotNil(t, config) + }) + + // Test invalid config - missing authentication + t.Run("bedrock invalid config missing auth", func(t *testing.T) { + host, status := test.NewTestHost(bedrockInvalidConfigMissingAuth) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusFailed, status) + }) + + // Test invalid config - missing region + t.Run("bedrock invalid config missing region", func(t *testing.T) { + host, status := test.NewTestHost(bedrockInvalidConfigMissingRegion) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusFailed, status) + }) + + // Test invalid config - partial ak/sk (only access key, no secret key) + t.Run("bedrock invalid config partial ak/sk", func(t *testing.T) { + host, status := test.NewTestHost(bedrockInvalidConfigPartialAkSk) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusFailed, status) + }) + }) +} + +func RunBedrockOnHttpRequestHeadersTests(t *testing.T) { + test.RunTest(t, func(t *testing.T) { + // Test Bedrock request headers processing with AWS Signature V4 + t.Run("bedrock chat completion request headers with ak/sk", func(t *testing.T) { + host, status := test.NewTestHost(basicBedrockConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // Set request headers + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + + require.Equal(t, types.HeaderStopIteration, action) + + // Verify request headers + requestHeaders := host.GetRequestHeaders() + require.NotNil(t, requestHeaders) + + // Verify Host is changed to Bedrock service domain + hostValue, hasHost := test.GetHeaderValue(requestHeaders, ":authority") + require.True(t, hasHost, "Host header should exist") + require.Contains(t, hostValue, "bedrock-runtime.us-east-1.amazonaws.com", "Host should be changed to Bedrock service domain") + }) + + // Test Bedrock request headers processing with Bearer Token + t.Run("bedrock chat completion request headers with api token", func(t *testing.T) { + host, status := test.NewTestHost(bedrockApiTokenConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // Set request headers + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + + require.Equal(t, types.HeaderStopIteration, action) + + // Verify request headers + requestHeaders := host.GetRequestHeaders() + require.NotNil(t, requestHeaders) + + // Verify Host is changed to Bedrock service domain + hostValue, hasHost := test.GetHeaderValue(requestHeaders, ":authority") + require.True(t, hasHost, "Host header should exist") + require.Contains(t, hostValue, "bedrock-runtime.us-east-1.amazonaws.com", "Host should be changed to Bedrock service domain") + }) + }) +} + +func RunBedrockOnHttpRequestBodyTests(t *testing.T) { + test.RunTest(t, func(t *testing.T) { + // Test Bedrock request body processing with Bearer Token authentication + t.Run("bedrock chat completion request body with api token", func(t *testing.T) { + host, status := test.NewTestHost(bedrockApiTokenConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // Set request headers + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + require.Equal(t, types.HeaderStopIteration, action) + + // Set request body + requestBody := `{ + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "Hello, how are you?" + } + ], + "temperature": 0.7 + }` + + action = host.CallOnHttpRequestBody([]byte(requestBody)) + require.Equal(t, types.ActionContinue, action) + + // Verify request headers for Bearer Token authentication + requestHeaders := host.GetRequestHeaders() + require.NotNil(t, requestHeaders) + + // Verify Authorization header uses Bearer token + authValue, hasAuth := test.GetHeaderValue(requestHeaders, "Authorization") + require.True(t, hasAuth, "Authorization header should exist") + require.Contains(t, authValue, "Bearer ", "Authorization should use Bearer token") + require.Contains(t, authValue, "test-token-for-unit-test", "Authorization should contain the configured token") + + // Verify path is transformed to Bedrock format + pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path") + require.True(t, hasPath, "Path header should exist") + require.Contains(t, pathValue, "/model/", "Path should contain Bedrock model path") + require.Contains(t, pathValue, "/converse", "Path should contain converse endpoint") + }) + + // Test Bedrock request body processing with AWS Signature V4 authentication + t.Run("bedrock chat completion request body with ak/sk", func(t *testing.T) { + host, status := test.NewTestHost(basicBedrockConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // Set request headers + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + require.Equal(t, types.HeaderStopIteration, action) + + // Set request body + requestBody := `{ + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "Hello, how are you?" + } + ], + "temperature": 0.7 + }` + + action = host.CallOnHttpRequestBody([]byte(requestBody)) + require.Equal(t, types.ActionContinue, action) + + // Verify request headers for AWS Signature V4 authentication + requestHeaders := host.GetRequestHeaders() + require.NotNil(t, requestHeaders) + + // Verify Authorization header uses AWS Signature + authValue, hasAuth := test.GetHeaderValue(requestHeaders, "Authorization") + require.True(t, hasAuth, "Authorization header should exist") + require.Contains(t, authValue, "AWS4-HMAC-SHA256", "Authorization should use AWS4-HMAC-SHA256 signature") + require.Contains(t, authValue, "Credential=", "Authorization should contain Credential") + require.Contains(t, authValue, "Signature=", "Authorization should contain Signature") + + // Verify X-Amz-Date header exists + dateValue, hasDate := test.GetHeaderValue(requestHeaders, "X-Amz-Date") + require.True(t, hasDate, "X-Amz-Date header should exist for AWS Signature V4") + require.NotEmpty(t, dateValue, "X-Amz-Date should not be empty") + + // Verify path is transformed to Bedrock format + pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path") + require.True(t, hasPath, "Path header should exist") + require.Contains(t, pathValue, "/model/", "Path should contain Bedrock model path") + require.Contains(t, pathValue, "/converse", "Path should contain converse endpoint") + }) + + // Test Bedrock streaming request + t.Run("bedrock streaming request", func(t *testing.T) { + host, status := test.NewTestHost(bedrockApiTokenConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // Set request headers + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + require.Equal(t, types.HeaderStopIteration, action) + + // Set streaming request body + requestBody := `{ + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "Hello" + } + ], + "stream": true + }` + + action = host.CallOnHttpRequestBody([]byte(requestBody)) + require.Equal(t, types.ActionContinue, action) + + // Verify path is transformed to Bedrock streaming format + requestHeaders := host.GetRequestHeaders() + pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path") + require.True(t, hasPath, "Path header should exist") + require.Contains(t, pathValue, "/model/", "Path should contain Bedrock model path") + require.Contains(t, pathValue, "/converse-stream", "Path should contain converse-stream endpoint for streaming") + }) + }) +} + +func RunBedrockOnHttpResponseHeadersTests(t *testing.T) { + test.RunTest(t, func(t *testing.T) { + // Test Bedrock response headers processing + t.Run("bedrock response headers", func(t *testing.T) { + host, status := test.NewTestHost(bedrockApiTokenConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // Set request headers + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + require.Equal(t, types.HeaderStopIteration, action) + + // Set request body + requestBody := `{ + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "Hello" + } + ] + }` + action = host.CallOnHttpRequestBody([]byte(requestBody)) + require.Equal(t, types.ActionContinue, action) + + // Process response headers + action = host.CallOnHttpResponseHeaders([][2]string{ + {":status", "200"}, + {"Content-Type", "application/json"}, + {"X-Amzn-Requestid", "test-request-id-12345"}, + }) + require.Equal(t, types.ActionContinue, action) + + // Verify response headers + responseHeaders := host.GetResponseHeaders() + require.NotNil(t, responseHeaders) + + // Verify status code + statusValue, hasStatus := test.GetHeaderValue(responseHeaders, ":status") + require.True(t, hasStatus, "Status header should exist") + require.Equal(t, "200", statusValue, "Status should be 200") + }) + }) +} + +func RunBedrockOnHttpResponseBodyTests(t *testing.T) { + test.RunTest(t, func(t *testing.T) { + // Test Bedrock response body processing + t.Run("bedrock response body", func(t *testing.T) { + host, status := test.NewTestHost(bedrockApiTokenConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // Set request headers + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + require.Equal(t, types.HeaderStopIteration, action) + + // Set request body + requestBody := `{ + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "Hello" + } + ] + }` + action = host.CallOnHttpRequestBody([]byte(requestBody)) + require.Equal(t, types.ActionContinue, action) + + // Set response property to ensure IsResponseFromUpstream() returns true + host.SetProperty([]string{"response", "code_details"}, []byte("via_upstream")) + + // Process response headers (must include :status 200 for body processing) + action = host.CallOnHttpResponseHeaders([][2]string{ + {":status", "200"}, + {"Content-Type", "application/json"}, + }) + require.Equal(t, types.ActionContinue, action) + + // Process response body (Bedrock format) + responseBody := `{ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "Hello! How can I help you today?" + } + ] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 10, + "outputTokens": 15, + "totalTokens": 25 + } + }` + + action = host.CallOnHttpResponseBody([]byte(responseBody)) + require.Equal(t, types.ActionContinue, action) + + // Verify response body is transformed to OpenAI format + transformedResponseBody := host.GetResponseBody() + require.NotNil(t, transformedResponseBody) + + var responseMap map[string]interface{} + err := json.Unmarshal(transformedResponseBody, &responseMap) + require.NoError(t, err) + + // Verify choices exist in transformed response + choices, exists := responseMap["choices"] + require.True(t, exists, "Choices should exist in response body") + require.NotNil(t, choices, "Choices should not be nil") + + // Verify usage exists + usage, exists := responseMap["usage"] + require.True(t, exists, "Usage should exist in response body") + require.NotNil(t, usage, "Usage should not be nil") + }) + }) +}