ratelimit_service_403_html_test.go 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "net/http"
  6. "testing"
  7. "github.com/Wei-Shaw/sub2api/internal/config"
  8. "github.com/stretchr/testify/require"
  9. )
  10. // countingOpenAI403CounterCache 在既有桩的基础上记录递增次数。
  11. // HTML 403 不仅不能处罚账号,连计数都不能加——否则后续真实的账号级 403
  12. // 会踩着这些"白涨"的计数提前触发永久禁用。
  13. type countingOpenAI403CounterCache struct {
  14. openAI403CounterCacheStub
  15. increments int
  16. }
  17. func (s *countingOpenAI403CounterCache) IncrementOpenAI403Count(ctx context.Context, accountID int64, window int) (int64, error) {
  18. s.increments++
  19. return s.openAI403CounterCacheStub.IncrementOpenAI403Count(ctx, accountID, window)
  20. }
  21. type openAI403TestHarness struct {
  22. svc *RateLimitService
  23. repo *rateLimitAccountRepoStub
  24. counter *countingOpenAI403CounterCache
  25. blocker *runtimeBlockRecorder
  26. account *Account
  27. }
  28. func newOpenAI403TestHarness(t *testing.T, accountID int64, counts ...int64) *openAI403TestHarness {
  29. t.Helper()
  30. repo := &rateLimitAccountRepoStub{}
  31. counter := &countingOpenAI403CounterCache{openAI403CounterCacheStub: openAI403CounterCacheStub{counts: counts}}
  32. blocker := &runtimeBlockRecorder{}
  33. svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  34. svc.SetOpenAI403CounterCache(counter)
  35. svc.SetAccountRuntimeBlocker(blocker)
  36. return &openAI403TestHarness{
  37. svc: svc,
  38. repo: repo,
  39. counter: counter,
  40. blocker: blocker,
  41. account: &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeOAuth},
  42. }
  43. }
  44. func (h *openAI403TestHarness) handle(body string) bool {
  45. return h.svc.HandleUpstreamError(
  46. context.Background(), h.account, http.StatusForbidden, http.Header{}, []byte(body),
  47. )
  48. }
  49. func (h *openAI403TestHarness) requireNoAccountPenalty(t *testing.T) {
  50. t.Helper()
  51. require.Equal(t, 0, h.repo.setErrorCalls, "端点级 403 不得永久禁用账号")
  52. require.Equal(t, 0, h.repo.tempCalls, "端点级 403 不得把账号设为临时不可调度")
  53. require.Empty(t, h.blocker.accounts, "端点级 403 不得触发调度阻断通知")
  54. require.Equal(t, 0, h.counter.increments, "端点级 403 不得递增连续 403 计数")
  55. }
  56. // issue #5334:无效的 /v1/responses 子路径被转发后,上游代理在到达 OpenAI API
  57. // 之前回 HTML 403 页面。这是链路/端点级响应,不是账号凭据失效的证据。
  58. const openAI403HTMLBody = "<!DOCTYPE html>\n<html><head><title>403 Forbidden</title></head>" +
  59. "<body><h1>403 Forbidden</h1></body></html>"
  60. func TestHandleUpstreamError_OpenAIHTML403DoesNotPenalizeAccount(t *testing.T) {
  61. cases := []struct {
  62. name string
  63. body string
  64. }{
  65. {"doctype_prefixed", openAI403HTMLBody},
  66. {"bare_html_tag", "<html><body>403 Forbidden</body></html>"},
  67. {"leading_whitespace_and_uppercase", "\n\t <!DOCTYPE HTML><html><body>Forbidden</body></html>"},
  68. }
  69. for _, tc := range cases {
  70. t.Run(tc.name, func(t *testing.T) {
  71. h := newOpenAI403TestHarness(t, 501, 1)
  72. shouldDisable := h.handle(tc.body)
  73. require.False(t, shouldDisable, "HTML 403 不得判定账号应下线")
  74. h.requireNoAccountPenalty(t)
  75. })
  76. }
  77. }
  78. // 一个持有 API Key 的调用方反复打无效子路径时,既有实现会在第
  79. // openAI403DisableThreshold 次把账号永久禁用。修复后连续多少次都不该升级。
  80. func TestHandleUpstreamError_OpenAIHTML403RepeatedNeverEscalates(t *testing.T) {
  81. h := newOpenAI403TestHarness(t, 502, 1, 2, 3, 4, 5)
  82. for i := 0; i < openAI403DisableThreshold+2; i++ {
  83. require.False(t, h.handle(openAI403HTMLBody), "第 %d 次 HTML 403 仍不得判定账号应下线", i+1)
  84. }
  85. h.requireNoAccountPenalty(t)
  86. }
  87. // 对照不变式:真正的结构化 JSON 403 是账号级证据,处罚链路必须原样保留。
  88. // 缺了这组断言,上面的跳过逻辑一旦写宽就会把真实的封号 403 也放过去。
  89. func TestHandleUpstreamError_OpenAIStructured403StillPenalizes(t *testing.T) {
  90. t.Run("first_hit_temp_unschedulable", func(t *testing.T) {
  91. h := newOpenAI403TestHarness(t, 503, 1)
  92. require.True(t, h.handle(`{"error":{"message":"Your account is not authorized"}}`))
  93. require.Equal(t, 1, h.counter.increments)
  94. require.Equal(t, 1, h.repo.tempCalls)
  95. require.Equal(t, 0, h.repo.setErrorCalls)
  96. require.Contains(t, h.repo.lastTempReason, "Your account is not authorized")
  97. require.Len(t, h.blocker.accounts, 1)
  98. })
  99. t.Run("threshold_disables", func(t *testing.T) {
  100. h := newOpenAI403TestHarness(t, 504, int64(openAI403DisableThreshold))
  101. require.True(t, h.handle(`{"error":{"message":"workspace forbidden by policy"}}`))
  102. require.Equal(t, 1, h.repo.setErrorCalls)
  103. require.Contains(t, h.repo.lastErrorMsg, "workspace forbidden by policy")
  104. })
  105. // 非 HTML 的非结构化响应(纯文本网关错误)不在本次放行范围内,维持原有处罚。
  106. t.Run("plain_text_body_unchanged", func(t *testing.T) {
  107. h := newOpenAI403TestHarness(t, 505, 1)
  108. require.True(t, h.handle("Forbidden"))
  109. require.Equal(t, 1, h.repo.tempCalls)
  110. })
  111. }
  112. // 作用域守卫:放行只针对 OpenAI 平台。其他平台的 403 处理不受影响。
  113. func TestHandleUpstreamError_HTML403OnOtherPlatformsUnchanged(t *testing.T) {
  114. for _, platform := range []string{PlatformAnthropic, PlatformGemini} {
  115. t.Run(platform, func(t *testing.T) {
  116. repo := &rateLimitAccountRepoStub{}
  117. svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  118. account := &Account{ID: 506, Platform: platform, Type: AccountTypeAPIKey}
  119. shouldDisable := svc.HandleUpstreamError(
  120. context.Background(), account, http.StatusForbidden, http.Header{}, []byte(openAI403HTMLBody),
  121. )
  122. require.True(t, shouldDisable)
  123. require.Equal(t, 1, repo.setErrorCalls, "其他平台保持原有 SetError 行为")
  124. })
  125. }
  126. }
  127. func TestIsHTMLResponse(t *testing.T) {
  128. cases := []struct {
  129. name string
  130. body string
  131. want bool
  132. }{
  133. {"doctype_lower", "<!doctype html><html></html>", true},
  134. {"doctype_upper", "<!DOCTYPE HTML>", true},
  135. {"bare_html", "<html lang=\"en\">", true},
  136. {"leading_whitespace", "\n\n <html>", true},
  137. {"json_error", `{"error":{"message":"forbidden"}}`, false},
  138. {"plain_text", "Forbidden", false},
  139. {"empty", "", false},
  140. // XML/SVG 之类不是 HTML,不在放行范围内。
  141. {"xml_declaration", `<?xml version="1.0"?><error/>`, false},
  142. }
  143. for _, tc := range cases {
  144. t.Run(tc.name, func(t *testing.T) {
  145. require.Equal(t, tc.want, isHTMLResponse([]byte(tc.body)))
  146. })
  147. }
  148. }