Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/account_header_override_test.go

@@ -0,0 +1,346 @@
+//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},
+		{"kimi apikey", PlatformKimi, AccountTypeAPIKey, true},
+		{"zhipu apikey", PlatformZhipu, AccountTypeAPIKey, true},
+		{"deepseek apikey", PlatformDeepseek, AccountTypeAPIKey, true},
+		{"anthropic oauth", PlatformAnthropic, AccountTypeOAuth, false},
+		{"openai oauth", PlatformOpenAI, AccountTypeOAuth, false},
+		{"kimi oauth", PlatformKimi, AccountTypeOAuth, false},
+		{"zhipu oauth", PlatformZhipu, AccountTypeOAuth, false},
+		{"deepseek oauth", PlatformDeepseek, 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)
+	})
+}