Files
sub2api/backend/internal/service/claude_code_validator_test.go
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

577 lines
20 KiB
Go
Raw Permalink 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.
package service
import (
"context"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/stretchr/testify/require"
)
const claudeCodeMetadataUserIDJSON = `{"device_id":"0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef","account_uuid":"","session_id":"aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"}`
func TestClaudeCodeValidator_ProbeBypass(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/1.2.3 (darwin; arm64)")
req = req.WithContext(context.WithValue(req.Context(), ctxkey.IsMaxTokensOneHaikuRequest, true))
ok := validator.Validate(req, map[string]any{
"model": "claude-haiku-4-5",
"max_tokens": 1,
})
require.True(t, ok)
}
func TestClaudeCodeValidator_ProbeBypassRequiresUA(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "curl/8.0.0")
req = req.WithContext(context.WithValue(req.Context(), ctxkey.IsMaxTokensOneHaikuRequest, true))
ok := validator.Validate(req, map[string]any{
"model": "claude-haiku-4-5",
"max_tokens": 1,
})
require.False(t, ok)
}
func TestClaudeCodeValidator_MessagesWithoutProbeStillNeedStrictValidation(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/1.2.3 (darwin; arm64)")
ok := validator.Validate(req, map[string]any{
"model": "claude-haiku-4-5",
"max_tokens": 1,
})
require.False(t, ok)
}
func TestClaudeCodeValidator_CountTokensPathUAOnly(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages/count_tokens", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
})
require.True(t, ok)
}
func TestClaudeCodeValidator_CountTokensPathRequiresUA(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages/count_tokens", nil)
req.Header.Set("User-Agent", "curl/8.0.0")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
})
require.False(t, ok)
}
func TestClaudeCodeValidator_MessagesPathFullValid(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)")
req.Header.Set("X-App", "claude-code")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
"stream": true,
"system": []any{
map[string]any{
"type": "text",
"text": "You are Claude Code, Anthropic's official CLI for Claude.",
},
},
"metadata": map[string]any{
"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
},
})
require.True(t, ok)
}
func TestClaudeCodeValidator_BillingBlockRecognizedWithoutIdentityPrompt(t *testing.T) {
// 真实抓取的完整安全监视器 system prompt(不含身份 prose)。
monitorPrompt, err := os.ReadFile("testdata/security_monitor_system_prompt.txt")
require.NoError(t, err)
validator := NewClaudeCodeValidator()
// 前提:完整监视器正文经 Dice 相似度远低于阈值,无法被身份 prose 机制识别——
// 故下面 Validate 的放行只可能来自计费归因块识别。
require.Less(t, validator.bestSimilarityScore(string(monitorPrompt)), systemPromptThreshold)
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.162 (external, cli)")
req.Header.Set("X-App", "cli")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
// Claude Code 安全监视器子请求:不携带身份 prose,但 system 数组携带计费归因块
// cc_entrypoint=cli,应据此识别为 Claude Code 客户端。
ok := validator.Validate(req, map[string]any{
"model": "claude-3-5-haiku-20241022",
"system": []any{
map[string]any{
"type": "text",
"text": "x-anthropic-billing-header: cc_version=2.1.162.884; cc_entrypoint=cli; cch=d8726;",
},
map[string]any{
"type": "text",
"text": string(monitorPrompt),
},
},
"metadata": map[string]any{
"user_id": claudeCodeMetadataUserIDJSON,
},
})
require.True(t, ok)
}
func TestClaudeCodeValidator_SecurityMonitorWithoutBillingBlock(t *testing.T) {
monitorPrompt, err := os.ReadFile("testdata/security_monitor_system_prompt.txt")
require.NoError(t, err)
validHeaders := map[string]string{
"User-Agent": "claude-cli/2.1.220 (external, cli)",
"X-App": "cli",
"anthropic-beta": "claude-code-20250219",
"anthropic-version": "2023-06-01",
}
validBody := func(prompt string) map[string]any {
return map[string]any{
"model": "claude-haiku-4-5-20251001",
"system": []any{
map[string]any{"type": "text", "text": prompt},
},
"metadata": map[string]any{"user_id": claudeCodeMetadataUserIDJSON},
}
}
// 真实 CLI(2.1.220)在监视器提示词之后追加的独立会话上下文块(脱敏),
// 随会话/环境变化,服务端不可控(见 issue #5152 抓包)。
sessionContext := "\n\n## Session Context\n\n- **User identity**: testuser\n" +
"- **Working directory**: /home/testuser/project\n- **Platform**: linux"
tests := []struct {
name string
headers map[string]string
body map[string]any
wantAccept bool
}{
{
name: "official classifier request",
headers: validHeaders,
body: validBody(string(monitorPrompt)),
wantAccept: true,
},
{
name: "classifier output with category element",
headers: validHeaders,
body: validBody(strings.Replace(
string(monitorPrompt),
"<block>yes</block><reason>",
"<block>yes</block><category>Exact BLOCK Rule Name</category><reason>",
1,
)),
wantAccept: true,
},
{
name: "non-Claude user agent",
headers: map[string]string{
"User-Agent": "curl/8.0.0",
"X-App": "cli",
"anthropic-beta": "claude-code-20250219",
"anthropic-version": "2023-06-01",
},
body: validBody(string(monitorPrompt)),
},
{
name: "missing X-App",
headers: map[string]string{
"User-Agent": validHeaders["User-Agent"],
"anthropic-beta": validHeaders["anthropic-beta"],
"anthropic-version": validHeaders["anthropic-version"],
},
body: validBody(string(monitorPrompt)),
},
{
name: "missing anthropic-beta",
headers: map[string]string{
"User-Agent": validHeaders["User-Agent"],
"X-App": validHeaders["X-App"],
"anthropic-version": validHeaders["anthropic-version"],
},
body: validBody(string(monitorPrompt)),
},
{
name: "missing anthropic-version",
headers: map[string]string{
"User-Agent": validHeaders["User-Agent"],
"X-App": validHeaders["X-App"],
"anthropic-beta": validHeaders["anthropic-beta"],
},
body: validBody(string(monitorPrompt)),
},
{
name: "missing metadata",
headers: validHeaders,
body: map[string]any{
"model": "claude-haiku-4-5-20251001",
"system": []any{map[string]any{"type": "text", "text": string(monitorPrompt)}},
},
},
{
name: "invalid metadata user ID",
headers: validHeaders,
body: func() map[string]any {
body := validBody(string(monitorPrompt))
body["metadata"] = map[string]any{"user_id": "invalid"}
return body
}(),
},
{
name: "unrelated prompt",
headers: validHeaders,
body: validBody("You are a different security classifier for coding agents."),
wantAccept: false,
},
{
name: "opening sentence alone",
headers: validHeaders,
body: validBody(claudeCodeSecurityMonitorPromptPrefix),
wantAccept: false,
},
{
name: "opening sentence plus arbitrary altered suffix",
headers: validHeaders,
body: validBody(claudeCodeSecurityMonitorPromptPrefix + "\n\n" +
strings.Repeat("This is arbitrary altered classifier content. ", 300)),
wantAccept: false,
},
{
// 回归 issue #5152:真实分类器请求携带 2 个 system entry
//(监视器提示词 + 追加的会话上下文块),不得因 entry 数量拒识。
name: "classifier with trailing session context entry",
headers: validHeaders,
body: func() map[string]any {
body := validBody(string(monitorPrompt))
system, ok := body["system"].([]any)
require.True(t, ok)
body["system"] = append(system, map[string]any{
"type": "text",
"text": sessionContext,
})
return body
}(),
wantAccept: true,
},
{
name: "classifier with leading session context entry",
headers: validHeaders,
body: func() map[string]any {
body := validBody(string(monitorPrompt))
system, ok := body["system"].([]any)
require.True(t, ok)
body["system"] = append([]any{map[string]any{
"type": "text",
"text": sessionContext,
}}, system...)
return body
}(),
wantAccept: true,
},
{
name: "session context entry alone",
headers: validHeaders,
body: validBody(sessionContext),
wantAccept: false,
},
{
// 篡改后的长提示词(marker 缺失)即便带上会话上下文块也不得放行。
name: "tampered classifier with session context entry",
headers: validHeaders,
body: func() map[string]any {
body := validBody(strings.ReplaceAll(
string(monitorPrompt), "## HARD BLOCK", "## ALTERED BLOCK"))
system, ok := body["system"].([]any)
require.True(t, ok)
body["system"] = append(system, map[string]any{
"type": "text",
"text": sessionContext,
})
return body
}(),
wantAccept: false,
},
}
validator := NewClaudeCodeValidator()
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
for name, value := range tt.headers {
req.Header.Set(name, value)
}
require.Equal(t, tt.wantAccept, validator.Validate(req, tt.body))
})
}
}
func TestClaudeCodeValidator_BillingBlockVSCodeEntrypointRecognized(t *testing.T) {
// 回归:Claude Code 在 VSCode 扩展内运行时,计费块入口为 cc_entrypoint=claude-vscode
// 而非 cli。其安全监视器子请求同样不携带身份 prose,此前写死 cc_entrypoint=cli 的
// 快速通道无法识别它,导致 claude_code_only 分组误拒。入口值不应作为识别条件。
monitorPrompt, err := os.ReadFile("testdata/security_monitor_system_prompt.txt")
require.NoError(t, err)
validator := NewClaudeCodeValidator()
// 前提:完整监视器正文经 Dice 相似度远低于阈值,放行只可能来自计费归因块识别。
require.Less(t, validator.bestSimilarityScore(string(monitorPrompt)), systemPromptThreshold)
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.181 (external, claude-vscode, agent-sdk/0.3.181)")
req.Header.Set("X-App", "cli")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
"system": []any{
map[string]any{
"type": "text",
"text": "x-anthropic-billing-header: cc_version=2.1.181.f17; cc_entrypoint=claude-vscode;",
},
map[string]any{
"type": "text",
"text": string(monitorPrompt),
},
},
"metadata": map[string]any{
"user_id": claudeCodeMetadataUserIDJSON,
},
})
require.True(t, ok)
}
func TestClaudeCodeValidator_BillingBlockWithoutEntrypointFallsThrough(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.162 (external, cli)")
req.Header.Set("X-App", "cli")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
// 计费块前缀命中但完全没有 cc_entrypoint= 字段,且无身份 prose
// 不应凭前缀放行,应落回 Dice 检查并失败。验证 cc_entrypoint= 字段的存在仍是必要条件。
ok := validator.Validate(req, map[string]any{
"model": "claude-3-5-haiku-20241022",
"system": []any{
map[string]any{
"type": "text",
"text": "x-anthropic-billing-header: cc_version=2.1.162.884; cch=d8726;",
},
map[string]any{
"type": "text",
"text": "Some unrelated system prompt that does not resemble Claude Code.",
},
},
"metadata": map[string]any{
"user_id": claudeCodeMetadataUserIDJSON,
},
})
require.False(t, ok)
}
func TestClaudeCodeValidator_BillingBlockStillRequiresClaudeCodeUA(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "curl/8.0.0")
req.Header.Set("X-App", "cli")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
// 计费块无法绕过 UA 校验:非 claude-cli 客户端在 Step 1 即被拒。
ok := validator.Validate(req, map[string]any{
"model": "claude-3-5-haiku-20241022",
"system": []any{
map[string]any{
"type": "text",
"text": "x-anthropic-billing-header: cc_version=2.1.162.884; cc_entrypoint=cli; cch=d8726;",
},
},
})
require.False(t, ok)
}
// 新版 Claude Code CLI 已取消 cch=... 签名字段,billing block 形如
// `x-anthropic-billing-header: cc_version=...; cc_entrypoint=cli;`(无 cch)。
// 检测依赖前缀 + cc_entrypoint=cli,不依赖 cch,故无身份 prose 的子请求仍应被识别。
// 这同时覆盖了本仓 mimicry 注入的新格式 block(见 buildBillingAttributionText)。
func TestClaudeCodeValidator_BillingBlockRecognizedWithoutCCH(t *testing.T) {
monitorPrompt, err := os.ReadFile("testdata/security_monitor_system_prompt.txt")
require.NoError(t, err)
validator := NewClaudeCodeValidator()
require.Less(t, validator.bestSimilarityScore(string(monitorPrompt)), systemPromptThreshold)
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.162 (external, cli)")
req.Header.Set("X-App", "cli")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
ok := validator.Validate(req, map[string]any{
"model": "claude-3-5-haiku-20241022",
"system": []any{
map[string]any{
"type": "text",
// 注意:无 cch 段,对齐新版 CLI 与本仓新的注入格式。
"text": "x-anthropic-billing-header: cc_version=2.1.162.884; cc_entrypoint=cli;",
},
map[string]any{
"type": "text",
"text": string(monitorPrompt),
},
},
"metadata": map[string]any{
"user_id": claudeCodeMetadataUserIDJSON,
},
})
require.True(t, ok, "无 cch 的新版 billing block 仍应被识别为 Claude Code")
}
// 安全回归:去掉 cch 后检测并未放松——非 claude-cli UA 即便携带无 cch 的 billing block
// 仍在 Step 1 被拒,ClaudeCodeOnly group 不会因此被仿冒绕过。
func TestClaudeCodeValidator_NoCCHBlockStillRequiresClaudeCodeUA(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "curl/8.0.0")
req.Header.Set("X-App", "cli")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
ok := validator.Validate(req, map[string]any{
"model": "claude-3-5-haiku-20241022",
"system": []any{
map[string]any{
"type": "text",
"text": "x-anthropic-billing-header: cc_version=2.1.162.884; cc_entrypoint=cli;",
},
},
})
require.False(t, ok)
}
func TestClaudeCodeValidator_MessagesPathRejectsNonClaudeCodeUA(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "curl/8.0.0")
req.Header.Set("X-App", "claude-code")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
"stream": true,
"system": []any{
map[string]any{
"type": "text",
"text": "You are Claude Code, Anthropic's official CLI for Claude.",
},
},
"metadata": map[string]any{
"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
},
})
require.False(t, ok)
}
func TestClaudeCodeValidator_MessagesPathWithoutSystemPromptStillRejected(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
req.Header.Set("User-Agent", "claude-cli/2.1.156 (Claude Code)")
req.Header.Set("X-App", "claude-code")
req.Header.Set("anthropic-beta", "claude-code-20250219")
req.Header.Set("anthropic-version", "2023-06-01")
ok := validator.Validate(req, map[string]any{
"model": "claude-opus-4-8",
"stream": true,
"messages": []any{
map[string]any{"role": "user", "content": "hello"},
},
"metadata": map[string]any{
"user_id": "user_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa_account__session_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
},
})
require.False(t, ok)
}
func TestClaudeCodeValidator_NonMessagesPathUAOnly(t *testing.T) {
validator := NewClaudeCodeValidator()
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/models", nil)
req.Header.Set("User-Agent", "claude-cli/1.2.3 (darwin; arm64)")
ok := validator.Validate(req, nil)
require.True(t, ok)
}
func TestExtractVersion(t *testing.T) {
v := NewClaudeCodeValidator()
tests := []struct {
ua string
want string
}{
{"claude-cli/2.1.22 (darwin; arm64)", "2.1.22"},
{"claude-cli/1.0.0", "1.0.0"},
{"Claude-CLI/3.10.5 (linux; x86_64)", "3.10.5"}, // 大小写不敏感
{"curl/8.0.0", ""}, // 非 Claude CLI
{"", ""}, // 空字符串
{"claude-cli/", ""}, // 无版本号
{"claude-cli/2.1.22-beta", "2.1.22"}, // 带后缀仍提取主版本号
}
for _, tt := range tests {
got := v.ExtractVersion(tt.ua)
require.Equal(t, tt.want, got, "ExtractVersion(%q)", tt.ua)
}
}
func TestCompareVersions(t *testing.T) {
tests := []struct {
a, b string
want int
}{
{"2.1.0", "2.1.0", 0}, // 相等
{"2.1.1", "2.1.0", 1}, // patch 更大
{"2.0.0", "2.1.0", -1}, // minor 更小
{"3.0.0", "2.99.99", 1}, // major 更大
{"1.0.0", "2.0.0", -1}, // major 更小
{"0.0.1", "0.0.0", 1}, // patch 差异
{"", "1.0.0", -1}, // 空字符串 vs 正常版本
{"v2.1.0", "2.1.0", 0}, // v 前缀处理
}
for _, tt := range tests {
got := CompareVersions(tt.a, tt.b)
require.Equal(t, tt.want, got, "CompareVersions(%q, %q)", tt.a, tt.b)
}
}
func TestSetGetClaudeCodeVersion(t *testing.T) {
ctx := context.Background()
require.Equal(t, "", GetClaudeCodeVersion(ctx), "empty context should return empty string")
ctx = SetClaudeCodeVersion(ctx, "2.1.63")
require.Equal(t, "2.1.63", GetClaudeCodeVersion(ctx))
}