Files
sub2api/backend/internal/service/account_header_override_test.go
T
shaw 221581400b feat(grok): 支持账号级自定义上游地址与请求头覆写
将账号级请求头覆写从 anthropic/openai 的 api_key 账号扩展到 Grok 的
api_key 与 oauth 账号,并放开 Grok OAuth 账号的自定义上游地址(仅作用于
转发端点,授权与 token 刷新链路不变)。

后端:
- IsHeaderOverrideEligible 扩展到 Grok(api_key+oauth),禁止名单新增
  x-grok-conv-id(逐请求会话路由头)。
- 所有 Grok 上游请求路径接线 ApplyHeaderOverrides(Responses/Chat 桥/
  媒体/配额探测/billing 探测/连通性测试),统一置于内置默认头之后。
- GetGrokBaseURL/GetGrokMediaBaseURL 放开 OAuth 自定义地址:官方地址
  视同未定制回落官方网关,仅显式第三方 host 改发转发流量。
- 非官方 host 允许任意 path 前缀,官方 host 仍强制 /v1。
- OAuth base_url 校验按 host 判定官方/自定义,自定义 host 恒受运营方
  URL 策略约束,不受 XAI_ALLOW_UNSAFE_URL_OVERRIDES 调试开关放宽。
- 给 GrokQuotaService 注入 config,使配额/billing 探测与转发共用同一
  URL 策略。

前端:
- Edit 模态为 Grok OAuth 账号新增「自定义上游地址」开关。
- 请求头表单新增 JSON 快速导入与按 JSON 一键复制(复用 useClipboard,
  兼容非安全上下文 HTTP 页面)。
- 门控改为平台×类型判定,Bulk 批量 base_url 增加格式校验,zh/en 文案同步。
2026-07-15 20:30:23 +08:00

341 lines
12 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//go:build unit
package service
import (
"net/http"
"strings"
"testing"
"github.com/stretchr/testify/require"
)
func headerOverrideTestAccount(platform, accountType string, credentials map[string]any) *Account {
return &Account{
Platform: platform,
Type: accountType,
Credentials: credentials,
}
}
func TestIsHeaderOverrideEligible(t *testing.T) {
tests := []struct {
name string
platform string
accType string
want bool
}{
{"anthropic apikey", PlatformAnthropic, AccountTypeAPIKey, true},
{"openai apikey", PlatformOpenAI, AccountTypeAPIKey, true},
{"anthropic oauth", PlatformAnthropic, AccountTypeOAuth, false},
{"openai oauth", PlatformOpenAI, AccountTypeOAuth, false},
{"gemini apikey", PlatformGemini, AccountTypeAPIKey, false},
{"grok apikey", PlatformGrok, AccountTypeAPIKey, true},
{"grok oauth", PlatformGrok, AccountTypeOAuth, true},
{"antigravity apikey", PlatformAntigravity, AccountTypeAPIKey, false},
{"anthropic bedrock", PlatformAnthropic, AccountTypeBedrock, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
acc := headerOverrideTestAccount(tt.platform, tt.accType, nil)
require.Equal(t, tt.want, acc.IsHeaderOverrideEligible())
})
}
var nilAccount *Account
require.False(t, nilAccount.IsHeaderOverrideEligible())
require.False(t, nilAccount.IsHeaderOverrideEnabled())
require.Nil(t, nilAccount.GetHeaderOverrides())
}
func TestIsHeaderOverrideEnabled(t *testing.T) {
acc := headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
})
require.True(t, acc.IsHeaderOverrideEnabled())
// 未配置 / 非 bool / false 均视为未启用
require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, nil).IsHeaderOverrideEnabled())
require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: "true",
}).IsHeaderOverrideEnabled())
require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: false,
}).IsHeaderOverrideEnabled())
// 不符合平台/类型条件时即使配置了 true 也不启用
require.False(t, headerOverrideTestAccount(PlatformAnthropic, AccountTypeOAuth, map[string]any{
credKeyHeaderOverrideEnabled: true,
}).IsHeaderOverrideEnabled())
require.False(t, headerOverrideTestAccount(PlatformGemini, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
}).IsHeaderOverrideEnabled())
}
func TestGetHeaderOverrides(t *testing.T) {
acc := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
"User-Agent": "my-agent/1.0", // 大写 key 归一化为小写
" X-App ": "cli", // 名称去空白
"x-empty": "", // 空 value(模板占位)跳过
"authorization": "Bearer leaked", // 禁止覆写的头跳过
"bad name": "value", // 非法 header 名跳过
"x-padded": " padded ", // value 去空白
},
})
overrides := acc.GetHeaderOverrides()
require.Equal(t, map[string]string{
"user-agent": "my-agent/1.0",
"x-app": "cli",
"x-padded": "padded",
}, overrides)
// 未启用时返回 nil
disabled := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrides: map[string]any{"user-agent": "x"},
})
require.Nil(t, disabled.GetHeaderOverrides())
// 启用但全部为空 value 时返回 nil
empty := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{"user-agent": ""},
})
require.Nil(t, empty.GetHeaderOverrides())
// 未经 Normalize 落库的超长数据 / WebSocket 握手头在应用时被防御性跳过
oversizedValue := strings.Repeat("a", maxHeaderOverrideValueLength+1)
defensive := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
"x-big": oversizedValue,
"sec-websocket-key": "forged",
"content-type": "application/json", // 名单扩充前落库的数据也要被拦截
"x-claude-code-session-id": "pinned-session",
"x-ok": "ok",
},
})
require.Equal(t, map[string]string{"x-ok": "ok"}, defensive.GetHeaderOverrides())
}
func TestApplyHeaderOverrides(t *testing.T) {
acc := headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
"user-agent": "override-agent/2.0",
"anthropic-beta": "custom-beta-1",
"x-custom": "custom-value",
},
})
h := http.Header{}
// 模拟转发链路:canonical key 与 wire casing 原样 key 混合存在
h.Set("User-Agent", "claude-cli/2.1.161 (external, cli)")
h["anthropic-beta"] = []string{"claude-code-20250219,oauth-2025-04-20"} // 非 canonical 原样 key
h.Set("Content-Type", "application/json")
acc.ApplyHeaderOverrides(h)
// user-agent 覆盖且只有一个值(已知头恢复 wire casing)
require.Equal(t, []string{"override-agent/2.0"}, h["User-Agent"])
// anthropic-beta:非 canonical 旧值被清除,写入 wire casing(小写)
require.Equal(t, []string{"custom-beta-1"}, h["anthropic-beta"])
require.Empty(t, h["Anthropic-Beta"])
// 新增头(未知头以小写原样键写入,与转发链路 wire casing 约定一致)
require.Equal(t, []string{"custom-value"}, h["x-custom"])
require.Equal(t, "custom-value", getHeaderRaw(h, "x-custom"))
// 未覆写的头不受影响
require.Equal(t, "application/json", h.Get("Content-Type"))
// 覆盖后不存在任何大小写重复
count := 0
for k := range h {
if k == "anthropic-beta" || k == "Anthropic-Beta" {
count++
}
}
require.Equal(t, 1, count)
}
func TestApplyHeaderOverridesNoOpPaths(t *testing.T) {
baseline := func() http.Header {
h := http.Header{}
h.Set("User-Agent", "orig")
return h
}
// OAuth 账号:即使配置了覆写也不生效
oauth := headerOverrideTestAccount(PlatformAnthropic, AccountTypeOAuth, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{"user-agent": "hacked"},
})
h := baseline()
oauth.ApplyHeaderOverrides(h)
require.Equal(t, "orig", h.Get("User-Agent"))
// 未启用开关
off := headerOverrideTestAccount(PlatformAnthropic, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrides: map[string]any{"user-agent": "hacked"},
})
h = baseline()
off.ApplyHeaderOverrides(h)
require.Equal(t, "orig", h.Get("User-Agent"))
// 禁止覆写的头(authorization / x-api-key / host 等)不会被应用
blocked := headerOverrideTestAccount(PlatformOpenAI, AccountTypeAPIKey, map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
"Authorization": "Bearer evil",
"X-Api-Key": "evil",
"Host": "evil.example.com",
"Content-Length": "0",
},
})
h = http.Header{}
h.Set("Authorization", "Bearer real-key")
blocked.ApplyHeaderOverrides(h)
require.Equal(t, "Bearer real-key", h.Get("Authorization"))
require.Empty(t, h.Get("X-Api-Key"))
require.Empty(t, h.Get("Host"))
// nil header 不 panic
blocked.ApplyHeaderOverrides(nil)
}
func TestNormalizeHeaderOverrideCredentials(t *testing.T) {
t.Run("nil credentials no-op", func(t *testing.T) {
require.NoError(t, NormalizeHeaderOverrideCredentials(nil))
})
t.Run("missing keys no-op", func(t *testing.T) {
creds := map[string]any{"api_key": "sk-xxx"}
require.NoError(t, NormalizeHeaderOverrideCredentials(creds))
_, exists := creds[credKeyHeaderOverrides]
require.False(t, exists)
})
t.Run("normalizes names and values", func(t *testing.T) {
creds := map[string]any{
credKeyHeaderOverrideEnabled: true,
credKeyHeaderOverrides: map[string]any{
" User-Agent ": " my-agent ",
"X-App": "",
"": "", // 完全空行被丢弃
},
}
require.NoError(t, NormalizeHeaderOverrideCredentials(creds))
require.Equal(t, map[string]any{
"user-agent": "my-agent",
"x-app": "",
}, creds[credKeyHeaderOverrides])
})
t.Run("accepts map[string]string input", func(t *testing.T) {
creds := map[string]any{
credKeyHeaderOverrides: map[string]string{"X-App": "cli"},
}
require.NoError(t, NormalizeHeaderOverrideCredentials(creds))
require.Equal(t, map[string]any{"x-app": "cli"}, creds[credKeyHeaderOverrides])
})
t.Run("rejects non-bool enabled", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrideEnabled: "yes",
})
require.Error(t, err)
})
t.Run("rejects non-object overrides", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: []any{"user-agent"},
})
require.Error(t, err)
})
t.Run("rejects non-string value", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{"x-app": 123},
})
require.Error(t, err)
})
t.Run("rejects invalid header name", func(t *testing.T) {
for _, name := range []string{"bad name", "bad:name", "bad\nname", "值"} {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{name: "v"},
})
require.Error(t, err, "name %q should be rejected", name)
}
})
t.Run("rejects empty name with value", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{" ": "v"},
})
require.Error(t, err)
})
t.Run("rejects blocked headers", func(t *testing.T) {
for _, name := range []string{
"Authorization", "x-api-key", "Host", "content-length", "Transfer-Encoding",
"connection", "accept-encoding", "Sec-WebSocket-Key", "session_id",
"conversation_id", "x-codex-turn-state", "chatgpt-account-id",
"Content-Type", "Cookie", "x-goog-api-key",
"X-Claude-Code-Session-Id", "x-client-request-id",
} {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{name: "v"},
})
require.Error(t, err, "blocked header %q should be rejected", name)
}
})
t.Run("allows tab inside value", func(t *testing.T) {
creds := map[string]any{
credKeyHeaderOverrides: map[string]any{"x-app": "a\tb"},
}
require.NoError(t, NormalizeHeaderOverrideCredentials(creds))
require.Equal(t, map[string]any{"x-app": "a\tb"}, creds[credKeyHeaderOverrides])
})
t.Run("rejects invalid value", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{"x-app": "bad\nvalue"},
})
require.Error(t, err)
})
t.Run("rejects duplicate names case-insensitively", func(t *testing.T) {
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{
"User-Agent": "a",
"user-agent": "b",
},
})
require.Error(t, err)
})
t.Run("rejects too many entries", func(t *testing.T) {
entries := make(map[string]any, maxHeaderOverrideEntries+1)
for i := 0; i <= maxHeaderOverrideEntries; i++ {
entries["x-h-"+string(rune('a'+i%26))+string(rune('a'+(i/26)%26))+string(rune('a'+(i/676)%26))] = "v"
}
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: entries,
})
require.Error(t, err)
})
t.Run("rejects oversized value", func(t *testing.T) {
big := make([]byte, maxHeaderOverrideValueLength+1)
for i := range big {
big[i] = 'a'
}
err := NormalizeHeaderOverrideCredentials(map[string]any{
credKeyHeaderOverrides: map[string]any{"x-app": string(big)},
})
require.Error(t, err)
})
}