+
+
+
+ 面向 Grok Build、Grok Web 与 Grok Console 的多账号 API 网关 +
+ ++ English | 简体中文 +
+ +
+ 
+
+
 |
+感谢 Krill AI 赞助了本项目！Krill 提供 GPT / Claude / Gemini / 多款国产模型的官方稳定极速的 API 中转服务，支持企业级定制、报销开票、7×16h 专属技术支持。更有独家适配的 WebSocket 连接，畅享极速首字速度。Krill 为本项目提供了特别优惠，使用此链接注册并在下订单时填写「grok2api」优惠码，首购套餐可享 Codex 77 折优惠！ | +
 |
+DEEIX-Chat 是一款开源可部署的 AI Chat 平台，面向需要长期、稳定、统一使用多模型能力的个人、团队与企业，将模型、对话、文件、工具调用与后台管理整合为一套可部署、可扩展的系统。点击 此处 开始部署！ | +
 |
+Right Code 是一个企业级 AI Agent 分发平台，主要提供稳定的 Claude Code、Codex、Gemini 等模型的中转服务。充值即可开票，企业、团队用户一对一对接。感谢 Right Code 提供的 Tokens 支持，点击 此处 注册并开始使用！ | +
access_token=secret owner@example.com https://grok.com/private?token=secret
` + strings.Repeat("A", 1024)) + upstreamErr := newWebMediaUpstreamError(http.StatusForbidden, body, true) + response := &http.Response{ + StatusCode: http.StatusForbidden, + ContentLength: 70000, + Header: http.Header{ + "Content-Type": []string{"text/html; charset=UTF-8"}, + "Server": []string{"cloudflare"}, + "Cf-Ray": []string{"test-ray-SIN"}, + }, + } + + adapter.logWebMediaUpstreamRejection("video_reference_upload", response, upstreamErr) + logLine := output.String() + for _, expected := range []string{ + "msg=web_media_upstream_rejected", "stage=video_reference_upload", "status=403", + "body_truncated=true", "body_prefix_sha256=", "body_kind=html", "cloudflare_challenge=true", + "content_type=\"text/html; charset=UTF-8\"", + "server=cloudflare", "cf_ray=test-ray-SIN", + } { + if !strings.Contains(logLine, expected) { + t.Fatalf("log missing %q: %s", expected, logLine) + } + } + for _, secret := range []string{"Just a moment", "challenge-secret", "access_token=secret", "owner@example.com", "token=secret", strings.Repeat("A", 256)} { + if strings.Contains(logLine, secret) { + t.Fatalf("log exposed %q: %s", secret, logLine) + } + } +} + +func TestChatModelsUseLowestSufficientTierFirst(t *testing.T) { + adapter := &Adapter{} + tests := []struct { + model string + want []account.WebTier + }{ + {model: "grok-chat-fast", want: []account.WebTier{account.WebTierBasic, account.WebTierSuper, account.WebTierHeavy}}, + {model: "grok-chat-auto", want: []account.WebTier{account.WebTierSuper, account.WebTierHeavy}}, + {model: "grok-chat-expert", want: []account.WebTier{account.WebTierSuper, account.WebTierHeavy}}, + {model: "grok-chat-heavy", want: []account.WebTier{account.WebTierHeavy}}, + } + for _, test := range tests { + got := adapter.TierOrder(test.model) + if !slices.Equal(got, test.want) { + t.Fatalf("tier order for %s = %v, want %v", test.model, got, test.want) + } + } +} + +func TestOnlyChatModelsExposeRateLimitModes(t *testing.T) { + for _, spec := range Catalog() { + if spec.Capability == modeldomain.CapabilityChat { + if !slices.Contains([]string{"auto", "fast", "expert", "heavy"}, spec.Mode) { + t.Fatalf("chat model %s has invalid quota mode %q", spec.PublicID, spec.Mode) + } + continue + } + if spec.ProtocolModel == "imagine-lite" { + if spec.Mode != "fast" { + t.Fatalf("Lite image must use fast quota mode, got %q", spec.Mode) + } + continue + } + if spec.Mode != "" { + t.Fatalf("media model %s must not expose chat quota mode %q", spec.PublicID, spec.Mode) + } + } +} + +func TestConsumeUpstreamChatFixture(t *testing.T) { + fixture := strings.Join([]string{ + `data: {"result":{"conversation":{"conversationId":"conv_1"}}}`, + `data: {"result":{"response":{"userResponse":{"responseId":"up_1"}}}}`, + `data: {"result":{"response":{"token":"thinking ","isThinking":true,"messageTag":"analysis"}}}`, + `data: {"result":{"response":{"token":"hello","isThinking":false,"messageTag":"final","webSearchResults":{"results":[{"url":"https://example.com"}]}}}}`, + `data: [DONE]`, + }, "\n") + parsed, err := consumeUpstream(strings.NewReader(fixture), nil) + if err != nil { + t.Fatal(err) + } + if parsed.ConversationID != "conv_1" || parsed.ParentID != "up_1" || parsed.Reasoning.String() != "thinking " || parsed.Text.String() != "hello" || len(parsed.SearchSources) != 1 { + t.Fatalf("parsed = %#v, text=%q reasoning=%q", parsed, parsed.Text.String(), parsed.Reasoning.String()) + } +} + +func TestWebVisibleStreamPhaseSuppressesLateReasoning(t *testing.T) { + phase := webVisibleStreamPhase{} + tests := []struct { + kind string + delta string + want bool + }{ + {kind: "reasoning", delta: "first thought", want: true}, + {kind: "text", delta: "", want: false}, + {kind: "reasoning", delta: "continued thought", want: true}, + {kind: "text", delta: "answer", want: true}, + {kind: "reasoning", delta: "late thought", want: false}, + {kind: "text", delta: " continued", want: true}, + } + for index, test := range tests { + if got := phase.Allow(test.kind, test.delta); got != test.want { + t.Fatalf("step %d Allow(%q, %q) = %v, want %v", index, test.kind, test.delta, got, test.want) + } + } +} + +func TestPreflightRejectsInBandErrorBeforeStreaming(t *testing.T) { + source := io.NopCloser(strings.NewReader(`data: {"error":{"message":"rate limited","code":8}}` + "\n")) + if _, err := preflightUpstream(source); err == nil || !strings.Contains(err.Error(), "rate limited") { + t.Fatalf("error = %v", err) + } +} + +func TestPreflightClassifiesAntiBotRejection(t *testing.T) { + source := io.NopCloser(strings.NewReader(`{"error":{"message":"Request rejected by anti-bot rules.","code":7,"details":[]}}` + "\n")) + if _, err := preflightUpstream(source); !errors.Is(err, errWebAntiBot) { + t.Fatalf("error = %v", err) + } +} + +func TestImagineRequestContainsOnlyProtocolProperties(t *testing.T) { + message := imagineRequestMessage("request", "prompt", "16:9", false, true, 8) + item := message["item"].(map[string]any) + content := item["content"].([]any)[0].(map[string]any) + properties := content["properties"].(map[string]any) + if properties["aspect_ratio"] != "16:9" || properties["enable_pro"] != true || properties["enable_nsfw"] != false || properties["num_generations"] != 8 { + t.Fatalf("properties = %#v", properties) + } + encoded := string(MarshalJSONBytes(message)) + assertForbiddenFieldsAbsent(t, encoded) +} + +func TestImagineResolutionAndBatchMapping(t *testing.T) { + tests := []struct { + resolution string + count int + pro bool + batch int + }{ + {resolution: "1k", count: 1, batch: 4}, + {resolution: "1k", count: 4, batch: 4}, + {resolution: "1k", count: 5, batch: 8}, + {resolution: "2k", count: 8, pro: true, batch: 8}, + {resolution: "2k", count: 9, pro: true, batch: 12}, + {resolution: "2k", count: 10, pro: true, batch: 12}, + } + for _, test := range tests { + config, ok := resolveImagineModel("imagine", test.resolution, test.count) + if !ok || config.Pro != test.pro || config.NativeBatchSize != test.batch || config.MaxReturnCount != 10 { + t.Fatalf("resolution=%s count=%d config=%#v", test.resolution, test.count, config) + } + } + config, _ := resolveImagineModel("imagine", "1k", 1) + if got := imagineUpstreamGenerationCount(true, 1, config); got != 1 { + t.Fatalf("streaming upstream count = %d, want 1", got) + } + if got := imagineUpstreamGenerationCount(false, 1, config); got != 4 { + t.Fatalf("non-streaming upstream count = %d, want 4", got) + } +} + +func TestImageStreamingRejectsMultipleOutputs(t *testing.T) { + response, err := (&Adapter{}).GenerateImage(context.Background(), provider.ImageGenerationRequest{ + Model: "grok-imagine-image-quality", Prompt: "cat", Count: 2, Streaming: true, + }) + if err != nil { + t.Fatal(err) + } + defer response.Body.Close() + body, _ := io.ReadAll(response.Body) + var payload map[string]any + if json.Unmarshal(body, &payload) != nil { + t.Fatalf("body=%s", body) + } + errorValue, _ := payload["error"].(map[string]any) + if response.StatusCode != http.StatusBadRequest || errorValue["message"] != "Streaming is only supported with n=1." || errorValue["type"] != "image_generation_user_error" || errorValue["param"] != "input" || errorValue["code"] != "unsupported_parameter" { + t.Fatalf("status=%d error=%#v", response.StatusCode, errorValue) + } +} + +func TestImageAspectRatioFollowsXAIContractAndSizeAlias(t *testing.T) { + for input, expected := range map[string]string{"auto": "auto", "19.5:9": "19.5:9", "9:20": "9:20", "1536x1024": "3:2", "1024x1536": "2:3"} { + got, err := resolveImageAspectRatio(input, "") + if err != nil || got != expected { + t.Fatalf("aspect ratio %q = %q, err=%v", input, got, err) + } + } + if got, err := resolveImageAspectRatio("", "1024x1024"); err != nil || got != "1:1" { + t.Fatalf("size alias = %q, err=%v", got, err) + } + if _, err := resolveImageAspectRatio("7:5", ""); err == nil { + t.Fatal("unsupported aspect ratio accepted") + } +} + +func TestImagineCollectorHandlesOutOfOrderFrames(t *testing.T) { + collector := newImagineCollector() + collector.Accept(map[string]any{"type": "json", "current_status": "completed", "image_id": "image-b", "order": 2.0, "moderated": false}) + collector.Accept(map[string]any{"type": "image", "id": "image-a", "grid_index": 0.0, "url": "https://imagine-public.x.ai/imagine-public/images/image-a.jpg"}) + collector.Accept(map[string]any{"type": "json", "current_status": "completed", "image_id": "image-a", "moderated": false}) + if collector.Done(2) { + t.Fatal("collector completed before delayed image payload arrived") + } + collector.Accept(map[string]any{"type": "image", "id": "image-b", "grid_index": 1.0, "url": "https://imagine-public.x.ai/imagine-public/images/image-b.jpg"}) + if !collector.Done(2) { + t.Fatal("collector did not complete after both image payloads arrived") + } + images := collector.Images() + if len(images) != 2 || images[0].ID != "image-a" || images[1].ID != "image-b" { + t.Fatalf("images = %#v", images) + } + ready := collector.ReadyImages() + if len(ready) != 2 || len(collector.ReadyImages()) != 0 { + t.Fatalf("ready images emitted more than once: %#v", ready) + } +} + +func TestImagineCollectorIgnoresWebSocketPreviewFrames(t *testing.T) { + collector := newImagineCollector() + collector.Accept(map[string]any{ + "type": "image", "id": "image-a", "order": 0.0, + "percentage_complete": 50.0, + "url": "https://imagine-public.x.ai/imagine-public/images/image-a.png", + "blob": "preview", + }) + previews := collector.ReadyPreviews() + if len(previews) != 1 || previews[0].ID != "image-a" || previews[0].Blob != "preview" || len(collector.ReadyPreviews()) != 0 { + t.Fatalf("previews = %#v", previews) + } + collector.Accept(map[string]any{ + "type": "json", "current_status": "completed", "image_id": "image-a", + "order": 0.0, "moderated": false, + }) + if collector.Done(1) || collector.UsableCount() != 0 || len(collector.Images()) != 0 { + t.Fatalf("preview became final: %#v", collector) + } + collector.Accept(map[string]any{ + "type": "image", "id": "image-a", "order": 0.0, + "percentage_complete": 100.0, + "url": "https://imagine-public.x.ai/imagine-public/images/image-a.jpg", + "blob": "final", + }) + if !collector.Done(1) || collector.UsableCount() != 1 { + t.Fatalf("final image not recognized: %#v", collector) + } + images := collector.Images() + if len(images) != 1 || images[0].URL != "https://imagine-public.x.ai/imagine-public/images/image-a.jpg" || images[0].Blob != "final" { + t.Fatalf("images = %#v", images) + } +} + +func TestImagineCollectorKeepsInterleavedJobsIsolated(t *testing.T) { + collector := newImagineCollector() + collector.Accept(map[string]any{ + "type": "image", "id": "image-a", "order": 0.0, "percentage_complete": 50.0, "blob": "a-preview", + }) + collector.Accept(map[string]any{ + "type": "image", "id": "image-b", "order": 1.0, "percentage_complete": 50.0, "blob": "b-preview", + }) + previews := collector.ReadyPreviews() + if len(previews) != 2 || previews[0].ID != "image-a" || previews[0].Blob != "a-preview" || previews[1].ID != "image-b" || previews[1].Blob != "b-preview" { + t.Fatalf("previews = %#v", previews) + } + collector.Accept(map[string]any{ + "type": "image", "id": "image-b", "order": 1.0, "percentage_complete": 100.0, "blob": "b-final", + }) + collector.Accept(map[string]any{"type": "json", "current_status": "completed", "image_id": "image-b", "order": 1.0, "moderated": false}) + ready := collector.ReadyImages() + if len(ready) != 1 || ready[0].ID != "image-b" || ready[0].Blob != "b-final" { + t.Fatalf("first ready = %#v", ready) + } + collector.Accept(map[string]any{ + "type": "image", "id": "image-a", "order": 0.0, "percentage_complete": 100.0, "blob": "a-final", + }) + collector.Accept(map[string]any{"type": "json", "current_status": "completed", "image_id": "image-a", "order": 0.0, "moderated": false}) + ready = collector.ReadyImages() + if len(ready) != 1 || ready[0].ID != "image-a" || ready[0].Blob != "a-final" { + t.Fatalf("second ready = %#v", ready) + } +} + +func TestImagineCollectorCanReturnRequestedSubsetBeforeNativeBatchCompletes(t *testing.T) { + collector := newImagineCollector() + for index := 0; index < 3; index++ { + id := fmt.Sprintf("image-%d", index) + collector.Accept(map[string]any{ + "type": "image", "id": id, "order": float64(index), "percentage_complete": 100.0, + "url": "https://imagine-public.x.ai/imagine-public/images/" + id + ".jpg", + }) + collector.Accept(map[string]any{ + "type": "json", "current_status": "completed", "image_id": id, + "order": float64(index), "moderated": false, + }) + } + if collector.Done(4) || collector.UsableCount() != 3 { + t.Fatalf("collector=%#v usable=%d", collector, collector.UsableCount()) + } + if collector.UsableCount() < 2 { + t.Fatal("a request for two images should already be satisfiable") + } +} + +func TestImagineCollectorSettlesModeratedSlots(t *testing.T) { + collector := newImagineCollector() + collector.Accept(map[string]any{"type": "json", "current_status": "completed", "image_id": "blocked", "moderated": true}) + if !collector.Done(1) || len(collector.Images()) != 0 { + t.Fatalf("moderated collector = %#v", collector) + } +} + +func TestGeneratedImageAssetHostsRemainStrict(t *testing.T) { + if !trustedImageAssetHost("assets.grok.com") || !trustedImageAssetHost("imagine-public.x.ai") || !trustedImageAssetHost("imgen.x.ai") || trustedImageAssetHost("example.com") { + t.Fatal("generated image host allowlist is incorrect") + } +} + +func TestImageStreamUsesOfficialOpenAIEventsWithoutTokenUsage(t *testing.T) { + adapter := &Adapter{assets: imageAssetStoreStub{}} + urlItem, err := adapter.imageDataItem(context.Background(), account.Credential{}, imagineImageValue{URL: "https://imgen.x.ai/image.jpg", Blob: "aW1hZ2U="}, "url") + if err != nil || urlItem["url"] != "https://api.example/v1/media/images/img_test" || urlItem["mime_type"] != "image/jpeg" || urlItem["revised_prompt"] != "" { + t.Fatalf("url item = %#v, err=%v", urlItem, err) + } + b64Item, err := adapter.imageDataItem(context.Background(), account.Credential{}, imagineImageValue{Blob: "aW1hZ2U="}, "b64_json") + if err != nil || b64Item["b64_json"] != "aW1hZ2U=" || b64Item["mime_type"] != "image/jpeg" { + t.Fatalf("base64 item = %#v, err=%v", b64Item, err) + } + png := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0, 'I', 'H', 'D', 'R'} + partial := openAIImageStreamEvent("image_generation.partial_image", imagineImageValue{ID: "image-a", Width: 960, Height: 960}, png, 0) + if partial["type"] != "image_generation.partial_image" || partial["partial_image_index"] != 0 || partial["size"] != "960x960" || partial["output_format"] != "png" || partial["usage"] != nil { + t.Fatalf("partial event = %#v", partial) + } + completed := openAIImageStreamEvent("image_generation.completed", imagineImageValue{ID: "image-a", Width: 960, Height: 960}, []byte("jpeg"), 0) + if completed["type"] != "image_generation.completed" || completed["partial_image_index"] != nil || completed["quality"] != "auto" || completed["usage"] != nil { + t.Fatalf("completed event = %#v", completed) + } + var output bytes.Buffer + if err := writeSSE(&output, "image_generation.partial_image", partial); err != nil { + t.Fatal(err) + } + value := output.String() + if !strings.Contains(value, "event: image_generation.partial_image") || strings.Contains(value, "image_generation.started") || strings.Contains(value, "image_generation.image.completed") || strings.Contains(value, `"usage"`) { + t.Fatalf("stream event = %q", value) + } +} + +func TestImageDataItemRetriesStorageWithoutRegenerating(t *testing.T) { + store := &imageAssetStoreRetryStub{failures: 2} + adapter := &Adapter{assets: store} + item, err := adapter.imageDataItem(context.Background(), account.Credential{ID: 42}, imagineImageValue{Blob: "aW1hZ2U="}, "url") + if err != nil { + t.Fatal(err) + } + if store.calls != mediaOutputAttempts || item["url"] != "https://api.example/v1/media/images/img_retry" { + t.Fatalf("storage retry calls=%d item=%#v", store.calls, item) + } +} + +func TestImageDataItemClassifiesExhaustedStorageFailure(t *testing.T) { + store := &imageAssetStoreRetryStub{failures: mediaOutputAttempts} + adapter := &Adapter{assets: store} + _, err := adapter.imageDataItem(context.Background(), account.Credential{ID: 42}, imagineImageValue{Blob: "aW1hZ2U="}, "url") + if err == nil || !provider.IsMediaPostProcessingError(err) || store.calls != mediaOutputAttempts { + t.Fatalf("storage failure err=%v calls=%d", err, store.calls) + } + var processingErr *provider.MediaPostProcessingError + if !errors.As(err, &processingErr) || processingErr.Stage != provider.MediaPostProcessingStorage { + t.Fatalf("storage failure classification = %#v", processingErr) + } +} + +type imageAssetStoreStub struct{} + +func (imageAssetStoreStub) SaveImage(context.Context, []byte) (mediadomain.Asset, error) { + return mediadomain.Asset{ID: "img_test", MIMEType: "image/jpeg"}, nil +} + +func (imageAssetStoreStub) PublicImageURL(string) string { + return "https://api.example/v1/media/images/img_test" +} + +type imageAssetStoreRetryStub struct { + failures int + calls int +} + +func (s *imageAssetStoreRetryStub) SaveImage(context.Context, []byte) (mediadomain.Asset, error) { + s.calls++ + if s.calls <= s.failures { + return mediadomain.Asset{}, errors.New("temporary storage failure") + } + return mediadomain.Asset{ID: "img_retry", MIMEType: "image/jpeg"}, nil +} + +func (*imageAssetStoreRetryStub) PublicImageURL(string) string { + return "https://api.example/v1/media/images/img_retry" +} + +func TestParseVideoStreamFixture(t *testing.T) { + fixture := `data: {"result":{"response":{"streamingVideoGenerationResponse":{"progress":42,"videoPostId":"post_1"}}}}` + "\n" + + `data: {"result":{"response":{"streamingVideoGenerationResponse":{"progress":100,"videoPostId":"post_1","videoUrl":"/videos/final.mp4"}}}}` + "\n" + response := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(fixture))} + progress := 0 + result, postID, err := parseVideoStream(response, func(value int) { progress = value }) + if err != nil { + t.Fatal(err) + } + if progress != 100 || postID != "post_1" || result.URL != "https://assets.grok.com/videos/final.mp4" || result.ContentType != "video/mp4" { + t.Fatalf("result = %#v, post = %q, progress = %d", result, postID, progress) + } +} + +func TestParseVideoStreamPreservesUpstreamStatus(t *testing.T) { + response := &http.Response{StatusCode: http.StatusTooManyRequests, Body: io.NopCloser(strings.NewReader("limited"))} + _, _, err := parseVideoStream(response, nil) + status, ok := provider.ErrorHTTPStatus(err) + if !ok || status != http.StatusTooManyRequests { + t.Fatalf("status = %d, ok = %v, err = %v", status, ok, err) + } +} + +func TestParseVideoConcatenatedJSONFixture(t *testing.T) { + fixture := `{"result":{"conversation":{"conversationId":"conversation_1"}}}` + + `{"result":{"response":{"streamingVideoGenerationResponse":{"videoId":"video_1","progress":1,"videoPostId":"post_1","resolutionName":"720p"}}}}` + + `{"result":{"response":{"streamingVideoGenerationResponse":{"videoId":"video_1","progress":95,"videoPostId":"post_1"}}}}` + + `{"result":{"response":{"streamingVideoGenerationResponse":{"videoId":"video_1","progress":100,"assetId":"video_1","videoPostId":"post_1","videoUrl":"users/user_1/generated/video_1/generated_video.mp4","thumbnailImageUrl":"users/user_1/generated/video_1/preview_image.jpg","moderated":false}}}}` + + `{"result":{"response":{"token":"I generated a video","isSoftStop":true}}}` + response := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(fixture))} + var values []int + result, postID, err := parseVideoStream(response, func(value int) { values = append(values, value) }) + if err != nil { + t.Fatal(err) + } + if !slices.Equal(values, []int{1, 95, 100}) || postID != "post_1" || result.URL != "https://assets.grok.com/users/user_1/generated/video_1/generated_video.mp4" || result.ContentType != "video/mp4" { + t.Fatalf("result = %#v, post = %q, progress = %#v", result, postID, values) + } +} + +func TestParseVideoStreamUsesModelResponseAttachment(t *testing.T) { + fixture := `data: {"result":{"response":{"streamingVideoGenerationResponse":{"progress":100,"videoPostId":"post_1"},"modelResponse":{"fileAttachments":["users/user_1/generated/video_1/generated_video.mp4"]}}}}` + "\n" + response := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(fixture))} + result, postID, err := parseVideoStream(response, nil) + if err != nil { + t.Fatal(err) + } + if postID != "post_1" || result.URL != "https://assets.grok.com/users/user_1/generated/video_1/generated_video.mp4" || result.ContentType != "video/mp4" { + t.Fatalf("result = %#v, post = %q", result, postID) + } +} + +func MarshalJSONBytes(value any) []byte { + data, _ := json.Marshal(value) + return data +} diff --git a/backend/internal/infra/provider/web/quota.go b/backend/internal/infra/provider/web/quota.go new file mode 100644 index 0000000000000000000000000000000000000000..f7268d4346d02660ca0a93ae403d660a97a5e66e --- /dev/null +++ b/backend/internal/infra/provider/web/quota.go @@ -0,0 +1,490 @@ +package web + +import ( + "bytes" + "context" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + "math" + "net/http" + "time" + + "github.com/chenyme/grok2api/backend/internal/domain/account" + domainegress "github.com/chenyme/grok2api/backend/internal/domain/egress" + "github.com/chenyme/grok2api/backend/internal/infra/provider" + "google.golang.org/protobuf/encoding/protowire" +) + +const weeklyQuotaMode = "weekly" + +func (a *Adapter) SyncQuota(ctx context.Context, credential account.Credential) (provider.QuotaSnapshot, error) { + windows := make([]account.QuotaWindow, 0, 2) + autoWindow, autoErr := a.SyncQuotaMode(ctx, credential, "auto") + if errors.Is(autoErr, provider.ErrUnauthorized) { + return provider.QuotaSnapshot{}, autoErr + } + if autoErr == nil { + windows = append(windows, autoWindow) + } + fastWindow, fastErr := a.SyncQuotaMode(ctx, credential, "fast") + if errors.Is(fastErr, provider.ErrUnauthorized) { + return provider.QuotaSnapshot{}, fastErr + } + if fastErr == nil { + windows = append(windows, fastWindow) + } + if len(windows) > 0 { + tier, _ := resolveWebTierFromQuota(credential.WebTier, windows, false) + // Basic/未知账号没有付费周池，避免为每次完整同步额外访问付费端点。 + // 只有模式额度已经确认付费等级时才读取 weekly 作为权威额度。 + if tier == account.WebTierSuper || tier == account.WebTierHeavy { + if weekly, weeklyErr := a.syncWeeklyCredits(ctx, credential); weeklyErr == nil { + windows = []account.QuotaWindow{weekly} + } + } + return provider.QuotaSnapshot{Tier: tier, Windows: windows, SyncedAt: time.Now().UTC()}, nil + } + // 模式端点暂不可用时，仅已确认的付费账号允许用 weekly 兜底； + // Basic/Auto 不能凭周额度探测提权，也不应制造无意义的付费端点流量。 + if credential.WebTier == account.WebTierSuper || credential.WebTier == account.WebTierHeavy { + if weekly, weeklyErr := a.syncWeeklyCredits(ctx, credential); weeklyErr == nil { + return provider.QuotaSnapshot{Tier: credential.WebTier, Windows: []account.QuotaWindow{weekly}, SyncedAt: time.Now().UTC()}, nil + } else { + return provider.QuotaSnapshot{}, weeklyErr + } + } + if fastErr != nil { + return provider.QuotaSnapshot{}, fastErr + } + return provider.QuotaSnapshot{}, autoErr +} + +func resolveWebTierFromQuota(current account.WebTier, windows []account.QuotaWindow, weeklyAvailable bool) (account.WebTier, bool) { + if len(windows) > 0 { + tier, known := inferWebTierFromQuota(windows) + if !known { + // 上游可能随时调整额度。无法识别的新形态保持 Auto，并由路由层 + // 按 Basic 的最小权限处理；不能伪造一个已确认的套餐等级。 + return account.WebTierAuto, false + } + return tier, weeklyAvailable && tier != account.WebTierBasic + } + // 周额度是付费账号信号，但无法区分 Super/Heavy。已确认的等级在模式 + // 额度暂时不可用时保留；未确认（Auto/Basic）的不能凭周额度提权。 + if current == account.WebTierHeavy || current == account.WebTierSuper { + return current, weeklyAvailable + } + return current, false +} + +// inferWebTierFromQuota 使用 Grok Web /rest/rate-limits 的真实额度形态判级。 +// 同一快照出现矛盾信号时选择较低等级，避免把 Basic 账号路由到付费能力。 +func inferWebTierFromQuota(windows []account.QuotaWindow) (account.WebTier, bool) { + detected := account.WebTierAuto + rank := map[account.WebTier]int{ + account.WebTierBasic: 1, + account.WebTierSuper: 2, + account.WebTierHeavy: 3, + } + for _, window := range windows { + candidate := account.WebTierAuto + switch window.Mode { + case "auto": + switch window.Total { + case 7, 20: + candidate = account.WebTierBasic + case 50: + candidate = account.WebTierSuper + case 150: + candidate = account.WebTierHeavy + } + case "fast": + switch window.Total { + case 30: + candidate = account.WebTierBasic + case 140: + candidate = account.WebTierSuper + case 400: + candidate = account.WebTierHeavy + } + case "heavy": + if window.Total > 0 { + candidate = account.WebTierHeavy + } + } + if candidate != account.WebTierAuto && (detected == account.WebTierAuto || rank[candidate] < rank[detected]) { + detected = candidate + } + } + return detected, detected != account.WebTierAuto +} + +func (a *Adapter) SyncQuotaMode(ctx context.Context, credential account.Credential, mode string) (account.QuotaWindow, error) { + if mode == weeklyQuotaMode { + return a.syncWeeklyCredits(ctx, credential) + } + cfg := a.config() + token, err := a.cipher.Decrypt(credential.EncryptedAccessToken) + if err != nil { + return account.QuotaWindow{}, err + } + lease, err := a.egress.AcquireCredential(ctx, domainegress.ScopeWeb, credential) + if err != nil { + return account.QuotaWindow{}, err + } + defer lease.Release() + payload, _ := json.Marshal(map[string]string{"modelName": mode}) + requestCtx, cancel := context.WithTimeout(ctx, time.Duration(cfg.QuotaTimeoutSeconds)*time.Second) + defer cancel() + endpoint := cfg.BaseURL + "/rest/rate-limits" + var response *http.Response + var body []byte + for attempt := 0; attempt < 2; attempt++ { + request, requestErr := http.NewRequestWithContext(requestCtx, http.MethodPost, endpoint, bytes.NewReader(payload)) + if requestErr != nil { + return account.QuotaWindow{}, requestErr + } + request.Header = buildHeaders(token, lease, "application/json") + applyAppHeaders(request.Header, cfg.BaseURL, cfg.BaseURL+"/") + a.applySignedStatsig(requestCtx, request, token, lease) + response, err = lease.DoDeferredForbidden(request) + if err != nil { + a.egress.Feedback(context.WithoutCancel(ctx), lease.NodeID, 0, err) + return account.QuotaWindow{}, err + } + body, err = io.ReadAll(io.LimitReader(response.Body, 4<<20)) + _ = response.Body.Close() + if err != nil { + return account.QuotaWindow{}, err + } + if response.StatusCode == http.StatusForbidden { + // Preserve definitive account-block signals before a Statsig retry can discard the first response. + if provider.IsDefinitiveAccountBlockBody(body) { + return account.QuotaWindow{}, fmt.Errorf("%w: account blocked", provider.ErrUnauthorized) + } + lease.InvalidateClearance() + if attempt == 0 && a.invalidateSignedStatsig(http.MethodPost, endpoint) { + continue + } + } + break + } + if response.StatusCode < 200 || response.StatusCode >= 300 { + if response.StatusCode == http.StatusUnauthorized { + return account.QuotaWindow{}, provider.ErrUnauthorized + } + if response.StatusCode == http.StatusForbidden && provider.IsDefinitiveAccountBlockBody(body) { + return account.QuotaWindow{}, fmt.Errorf("%w: account blocked", provider.ErrUnauthorized) + } + a.egress.Feedback(context.WithoutCancel(ctx), lease.NodeID, response.StatusCode, nil) + return account.QuotaWindow{}, fmt.Errorf("Grok Web 额度接口返回 %d", response.StatusCode) + } + a.egress.Feedback(context.WithoutCancel(ctx), lease.NodeID, response.StatusCode, nil) + var value struct { + WindowSizeSeconds int `json:"windowSizeSeconds"` + RemainingQueries int `json:"remainingQueries"` + TotalQueries int `json:"totalQueries"` + } + if err := json.Unmarshal(body, &value); err != nil { + return account.QuotaWindow{}, err + } + if value.TotalQueries <= 0 { + return account.QuotaWindow{}, fmt.Errorf("Grok Web 额度响应缺少 totalQueries") + } + if value.WindowSizeSeconds <= 0 { + value.WindowSizeSeconds = 7200 + } + now := time.Now().UTC() + resetAt := now.Add(time.Duration(value.WindowSizeSeconds) * time.Second) + return account.QuotaWindow{ + AccountID: credential.ID, Mode: mode, Remaining: max(0, value.RemainingQueries), Total: value.TotalQueries, + WindowSeconds: value.WindowSizeSeconds, ResetAt: &resetAt, SyncedAt: &now, Source: account.QuotaSourceUpstream, UpdatedAt: now, + }, nil +} + +func (a *Adapter) syncWeeklyCredits(ctx context.Context, credential account.Credential) (account.QuotaWindow, error) { + cfg := a.config() + token, err := a.cipher.Decrypt(credential.EncryptedAccessToken) + if err != nil { + return account.QuotaWindow{}, err + } + lease, err := a.egress.AcquireCredential(ctx, domainegress.ScopeWeb, credential) + if err != nil { + return account.QuotaWindow{}, err + } + defer lease.Release() + + requestCtx, cancel := context.WithTimeout(ctx, time.Duration(cfg.QuotaTimeoutSeconds)*time.Second) + defer cancel() + endpoint := cfg.BaseURL + "/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig" + request, err := http.NewRequestWithContext(requestCtx, http.MethodPost, endpoint, bytes.NewReader([]byte{0, 0, 0, 0, 0})) + if err != nil { + return account.QuotaWindow{}, err + } + request.Header = buildHeaders(token, lease, "application/grpc-web+proto") + applyAppHeaders(request.Header, cfg.BaseURL, cfg.BaseURL+"/") + request.Header.Del("x-xai-request-id") + request.Header.Set("x-grpc-web", "1") + request.Header.Set("x-user-agent", "connect-es/2.1.1") + + response, err := lease.DoDeferredForbidden(request) + if err != nil { + a.egress.Feedback(context.WithoutCancel(ctx), lease.NodeID, 0, err) + return account.QuotaWindow{}, err + } + defer response.Body.Close() + body, err := io.ReadAll(io.LimitReader(response.Body, 4<<20)) + if err != nil { + return account.QuotaWindow{}, err + } + if response.StatusCode < 200 || response.StatusCode >= 300 { + if response.StatusCode == http.StatusUnauthorized { + return account.QuotaWindow{}, provider.ErrUnauthorized + } + if response.StatusCode == http.StatusForbidden && provider.IsDefinitiveAccountBlockBody(body) { + return account.QuotaWindow{}, fmt.Errorf("%w: account blocked", provider.ErrUnauthorized) + } + if response.StatusCode == http.StatusForbidden { + lease.InvalidateClearance() + } + a.egress.Feedback(context.WithoutCancel(ctx), lease.NodeID, response.StatusCode, nil) + return account.QuotaWindow{}, fmt.Errorf("Grok Web 周额度接口返回 %d", response.StatusCode) + } + window, err := parseWeeklyCreditsResponse(body, credential.ID, time.Now().UTC()) + if err != nil { + return account.QuotaWindow{}, err + } + a.egress.Feedback(context.WithoutCancel(ctx), lease.NodeID, response.StatusCode, nil) + return window, nil +} + +func parseWeeklyCreditsResponse(body []byte, accountID uint64, syncedAt time.Time) (account.QuotaWindow, error) { + payload, err := firstGRPCWebMessage(body) + if err != nil { + return account.QuotaWindow{}, err + } + config, err := protobufMessageField(payload, 1) + if err != nil { + return account.QuotaWindow{}, fmt.Errorf("解析 Grok Web 周额度响应: %w", err) + } + var usagePercent float64 + var usagePresent bool + var periodStart, periodEnd *time.Time + breakdown := make([]account.QuotaBreakdown, 0, 8) + for len(config) > 0 { + number, fieldType, n := protowire.ConsumeTag(config) + if n < 0 { + return account.QuotaWindow{}, fmt.Errorf("周额度 protobuf tag 无效") + } + config = config[n:] + switch { + case number == 1 && fieldType == protowire.Fixed32Type: + value, consumed := protowire.ConsumeFixed32(config) + if consumed < 0 { + return account.QuotaWindow{}, fmt.Errorf("周额度使用率无效") + } + usagePercent = float64(math.Float32frombits(value)) + usagePresent = true + config = config[consumed:] + case (number == 4 || number == 5) && fieldType == protowire.BytesType: + value, consumed := protowire.ConsumeBytes(config) + if consumed < 0 { + return account.QuotaWindow{}, fmt.Errorf("周额度周期无效") + } + parsed, parseErr := parseProtoTimestamp(value) + if parseErr != nil { + return account.QuotaWindow{}, parseErr + } + if number == 4 { + periodStart = &parsed + } else { + periodEnd = &parsed + } + config = config[consumed:] + case number == 7 && fieldType == protowire.BytesType: + value, consumed := protowire.ConsumeBytes(config) + if consumed < 0 { + return account.QuotaWindow{}, fmt.Errorf("周额度产品分解无效") + } + if item, ok := parseQuotaBreakdown(value); ok { + breakdown = append(breakdown, item) + } + config = config[consumed:] + default: + consumed := protowire.ConsumeFieldValue(number, fieldType, config) + if consumed < 0 { + return account.QuotaWindow{}, fmt.Errorf("周额度 protobuf 字段无效") + } + config = config[consumed:] + } + } + if usagePresent && (math.IsNaN(usagePercent) || math.IsInf(usagePercent, 0) || usagePercent < 0 || usagePercent > 100) { + return account.QuotaWindow{}, fmt.Errorf("Grok Web 周额度响应缺少有效使用率") + } + if periodStart == nil || periodEnd == nil || !periodEnd.After(*periodStart) { + return account.QuotaWindow{}, fmt.Errorf("Grok Web 周额度响应缺少有效周期") + } + if !usagePresent && periodStart.Nanosecond() == 0 && periodEnd.Nanosecond() == 0 { + // Free accounts return a coarse entitlement period without a usage rate. + // A paid, unused weekly pool has the same rate omitted but retains its + // precise period boundaries, which represents zero percent used. + return account.QuotaWindow{}, fmt.Errorf("Grok Web 周额度响应缺少有效使用率") + } + windowSeconds := int(periodEnd.Sub(*periodStart).Seconds()) + if windowSeconds < 24*60*60 || windowSeconds > 31*24*60*60 { + return account.QuotaWindow{}, fmt.Errorf("Grok Web 周额度周期长度异常") + } + usedBasisPoints := int(math.Round(usagePercent * 100)) + return account.QuotaWindow{ + AccountID: accountID, Mode: weeklyQuotaMode, Remaining: max(0, 10000-usedBasisPoints), Total: 10000, + UsagePercent: usagePercent, Breakdown: breakdown, WindowSeconds: windowSeconds, + ResetAt: periodEnd, SyncedAt: &syncedAt, Source: account.QuotaSourceUpstream, UpdatedAt: syncedAt, + }, nil +} + +func firstGRPCWebMessage(body []byte) ([]byte, error) { + message, grpcStatus, err := parseGRPCWebFrames(body) + if err != nil { + return nil, err + } + if grpcStatus != "" && grpcStatus != "0" { + return nil, fmt.Errorf("Grok Web 周额度 gRPC 状态为 %s", grpcStatus) + } + if message == nil { + return nil, fmt.Errorf("Grok Web 周额度响应缺少消息帧") + } + return message, nil +} + +func parseGRPCWebFrames(body []byte) ([]byte, string, error) { + var message []byte + grpcStatus := "" + for len(body) > 0 { + if len(body) < 5 { + return nil, "", fmt.Errorf("gRPC-Web 响应包含不完整帧头") + } + flag := body[0] + length := int(binary.BigEndian.Uint32(body[1:5])) + body = body[5:] + if length < 0 || length > len(body) { + return nil, "", fmt.Errorf("gRPC-Web 帧长度无效") + } + payload := body[:length] + body = body[length:] + if flag&0x80 == 0 { + if flag != 0 { + return nil, "", fmt.Errorf("不支持压缩的 gRPC-Web 响应") + } + if message == nil { + message = append([]byte(nil), payload...) + } + continue + } + for _, line := range bytes.Split(payload, []byte{'\n'}) { + name, value, ok := bytes.Cut(bytes.TrimSpace(line), []byte{':'}) + if ok && string(bytes.ToLower(bytes.TrimSpace(name))) == "grpc-status" { + grpcStatus = string(bytes.TrimSpace(value)) + } + } + } + return message, grpcStatus, nil +} + +func protobufMessageField(message []byte, target protowire.Number) ([]byte, error) { + for len(message) > 0 { + number, fieldType, n := protowire.ConsumeTag(message) + if n < 0 { + return nil, fmt.Errorf("protobuf tag 无效") + } + message = message[n:] + if number == target && fieldType == protowire.BytesType { + value, consumed := protowire.ConsumeBytes(message) + if consumed < 0 { + return nil, fmt.Errorf("protobuf message 无效") + } + return value, nil + } + consumed := protowire.ConsumeFieldValue(number, fieldType, message) + if consumed < 0 { + return nil, fmt.Errorf("protobuf 字段无效") + } + message = message[consumed:] + } + return nil, fmt.Errorf("protobuf 缺少字段 %d", target) +} + +func parseProtoTimestamp(message []byte) (time.Time, error) { + var seconds int64 + var nanos int32 + for len(message) > 0 { + number, fieldType, n := protowire.ConsumeTag(message) + if n < 0 { + return time.Time{}, fmt.Errorf("protobuf timestamp tag 无效") + } + message = message[n:] + if fieldType == protowire.VarintType && (number == 1 || number == 2) { + value, consumed := protowire.ConsumeVarint(message) + if consumed < 0 { + return time.Time{}, fmt.Errorf("protobuf timestamp 值无效") + } + if number == 1 { + seconds = int64(value) + } else { + nanos = int32(value) + } + message = message[consumed:] + continue + } + consumed := protowire.ConsumeFieldValue(number, fieldType, message) + if consumed < 0 { + return time.Time{}, fmt.Errorf("protobuf timestamp 字段无效") + } + message = message[consumed:] + } + if seconds <= 0 || nanos < 0 || nanos >= int32(time.Second) { + return time.Time{}, fmt.Errorf("protobuf timestamp 范围无效") + } + return time.Unix(seconds, int64(nanos)).UTC(), nil +} + +func parseQuotaBreakdown(message []byte) (account.QuotaBreakdown, bool) { + var result account.QuotaBreakdown + var codePresent bool + for len(message) > 0 { + number, fieldType, n := protowire.ConsumeTag(message) + if n < 0 { + return account.QuotaBreakdown{}, false + } + message = message[n:] + switch { + case number == 1 && fieldType == protowire.VarintType: + value, consumed := protowire.ConsumeVarint(message) + if consumed < 0 { + return account.QuotaBreakdown{}, false + } + result.ProductCode = int(value) + codePresent = true + message = message[consumed:] + case number == 2 && fieldType == protowire.Fixed32Type: + value, consumed := protowire.ConsumeFixed32(message) + if consumed < 0 { + return account.QuotaBreakdown{}, false + } + result.UsagePercent = float64(math.Float32frombits(value)) + message = message[consumed:] + default: + consumed := protowire.ConsumeFieldValue(number, fieldType, message) + if consumed < 0 { + return account.QuotaBreakdown{}, false + } + message = message[consumed:] + } + } + if !codePresent || result.ProductCode < 0 || result.UsagePercent < 0 || result.UsagePercent > 100 || math.IsNaN(result.UsagePercent) || math.IsInf(result.UsagePercent, 0) { + return account.QuotaBreakdown{}, false + } + return result, true +} diff --git a/backend/internal/infra/provider/web/quota_test.go b/backend/internal/infra/provider/web/quota_test.go new file mode 100644 index 0000000000000000000000000000000000000000..a1ad19b0e017506d8bee38f3099e30a3dd4b9aa3 --- /dev/null +++ b/backend/internal/infra/provider/web/quota_test.go @@ -0,0 +1,442 @@ +package web + +import ( + "context" + "encoding/base64" + "encoding/hex" + "encoding/json" + "errors" + "math" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/chenyme/grok2api/backend/internal/domain/account" + egressdomain "github.com/chenyme/grok2api/backend/internal/domain/egress" + infraegress "github.com/chenyme/grok2api/backend/internal/infra/egress" + "github.com/chenyme/grok2api/backend/internal/infra/provider" + "github.com/chenyme/grok2api/backend/internal/infra/security" + "github.com/chenyme/grok2api/backend/internal/repository" +) + +const capturedWeeklyCreditsHex = "00000000630a610d0000304112001a00220c089abbccd2061080f2d1fc012a0c089ab0f1d2061080f2d1fc013a07080515000020413a070804150000803f3a020802421e0802120c089abbccd2061080f2d1fc011a0c089ab0f1d2061080f2d1fc01580162006801800000000f677270632d7374617475733a300d0a" + +func TestParseCapturedWeeklyCreditsResponse(t *testing.T) { + body, err := hex.DecodeString(capturedWeeklyCreditsHex) + if err != nil { + t.Fatal(err) + } + syncedAt := time.Date(2026, 7, 12, 14, 0, 0, 0, time.UTC) + window, err := parseWeeklyCreditsResponse(body, 42, syncedAt) + if err != nil { + t.Fatal(err) + } + if window.AccountID != 42 || window.Mode != weeklyQuotaMode || window.Total != 10000 || window.Remaining != 8900 || window.WindowSeconds != 7*24*60*60 { + t.Fatalf("window = %#v", window) + } + if math.Abs(window.UsagePercent-11) > 0.001 || window.ResetAt == nil || window.ResetAt.Unix() != 1784436762 { + t.Fatalf("usage/reset = %#v", window) + } + if len(window.Breakdown) != 3 || window.Breakdown[0].ProductCode != account.QuotaProductImagine || window.Breakdown[0].UsagePercent != 10 || window.Breakdown[1].ProductCode != account.QuotaProductChat || window.Breakdown[1].UsagePercent != 1 || window.Breakdown[2].ProductCode != account.QuotaProductBuild || window.Breakdown[2].UsagePercent != 0 { + t.Fatalf("breakdown = %#v", window.Breakdown) + } +} + +func TestParseUnusedPreciseWeeklyCreditsResponse(t *testing.T) { + body, err := hex.DecodeString("00000000480a4612001a00220c08c5d5d3d20610c0c7a1ee012a0c08c5caf8d20610c0c7a1ee01421e0802120c08c5d5d3d20610c0c7a1ee011a0c08c5caf8d20610c0c7a1ee01580162006801800000000f677270632d7374617475733a300d0a") + if err != nil { + t.Fatal(err) + } + window, err := parseWeeklyCreditsResponse(body, 42, time.Now().UTC()) + if err != nil { + t.Fatal(err) + } + if window.Total != 10000 || window.Remaining != 10000 || window.UsagePercent != 0 || window.ResetAt == nil || window.ResetAt.Nanosecond() == 0 { + t.Fatalf("window = %#v", window) + } +} + +func TestParseCoarseWeeklyCreditsResponseRemainsUnavailable(t *testing.T) { + body, err := hex.DecodeString("00000000300a2e12001a0022060880a6b6d2062a0608809bdbd2064212080212060880a6b6d2061a0608809bdbd206580162006801800000000f677270632d7374617475733a300d0a") + if err != nil { + t.Fatal(err) + } + if _, err := parseWeeklyCreditsResponse(body, 42, time.Now().UTC()); err == nil { + t.Fatal("expected coarse period without usage to be rejected") + } +} + +func TestSyncQuotaFetchesWeeklyOnlyAfterPaidTierIsConfirmed(t *testing.T) { + weeklyBody, err := hex.DecodeString(capturedWeeklyCreditsHex) + if err != nil { + t.Fatal(err) + } + var weeklyCalls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig": + weeklyCalls.Add(1) + writer.Header().Set("Content-Type", "application/grpc-web+proto") + _, _ = writer.Write(weeklyBody) + case "/rest/rate-limits": + var payload struct { + ModelName string `json:"modelName"` + } + if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { + t.Errorf("quota payload: %v", err) + } + total := map[string]int{"auto": 50, "fast": 140}[payload.ModelName] + _ = json.NewEncoder(writer).Encode(map[string]any{ + "windowSizeSeconds": 7200, "remainingQueries": total, "totalQueries": total, + }) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + cipher, err := security.NewCipher(base64.StdEncoding.EncodeToString(make([]byte, 32))) + if err != nil { + t.Fatal(err) + } + encrypted, err := cipher.Encrypt("test-sso") + if err != nil { + t.Fatal(err) + } + adapter := NewAdapter(Config{ + BaseURL: server.URL, StatsigMode: "manual", StatsigManualValue: "test-signature", + }, infraegress.NewManager(egressRepositoryStub{}, cipher), cipher, nil, nil) + snapshot, err := adapter.SyncQuota(context.Background(), account.Credential{ID: 2, WebTier: account.WebTierAuto, EncryptedAccessToken: encrypted}) + if err != nil { + t.Fatal(err) + } + if snapshot.Tier != account.WebTierSuper || len(snapshot.Windows) != 1 || snapshot.Windows[0].Mode != weeklyQuotaMode || weeklyCalls.Load() != 1 { + t.Fatalf("snapshot = %#v, weekly calls = %d", snapshot, weeklyCalls.Load()) + } +} + +func TestSyncQuotaStopsAfterFirstUnauthorizedMode(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + http.Error(writer, "unauthorized", http.StatusUnauthorized) + })) + defer server.Close() + + cipher, err := security.NewCipher(base64.StdEncoding.EncodeToString(make([]byte, 32))) + if err != nil { + t.Fatal(err) + } + encrypted, err := cipher.Encrypt("expired-sso") + if err != nil { + t.Fatal(err) + } + adapter := NewAdapter(Config{ + BaseURL: server.URL, StatsigMode: "manual", StatsigManualValue: "test-signature", + }, infraegress.NewManager(egressRepositoryStub{}, cipher), cipher, nil, nil) + _, err = adapter.SyncQuota(context.Background(), account.Credential{ID: 3, WebTier: account.WebTierAuto, EncryptedAccessToken: encrypted}) + if !errors.Is(err, provider.ErrUnauthorized) { + t.Fatalf("err = %v", err) + } + if calls.Load() != 1 { + t.Fatalf("unauthorized credential made %d quota requests", calls.Load()) + } +} + +func TestSyncQuotaBlockedForbiddenIsUnauthorized(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusForbidden) + _, _ = writer.Write([]byte(`{"code":7,"message":"User is blocked [WKE=unauthorized:blocked-user]","details":[]}`)) + })) + defer server.Close() + + cipher, err := security.NewCipher(base64.StdEncoding.EncodeToString(make([]byte, 32))) + if err != nil { + t.Fatal(err) + } + encrypted, err := cipher.Encrypt("blocked-sso") + if err != nil { + t.Fatal(err) + } + egressRepository := &recordingWebEgressRepository{node: egressdomain.Node{ + ID: 1, Name: "web", Scope: egressdomain.ScopeWeb, Enabled: true, Health: 1, + }} + adapter := NewAdapter(Config{ + BaseURL: server.URL, StatsigMode: "manual", StatsigManualValue: "test-signature", + }, infraegress.NewManager(egressRepository, cipher), cipher, nil, nil) + _, err = adapter.SyncQuota(context.Background(), account.Credential{ID: 4, WebTier: account.WebTierAuto, EncryptedAccessToken: encrypted}) + if !errors.Is(err, provider.ErrUnauthorized) { + t.Fatalf("err = %v, want ErrUnauthorized", err) + } + // manual Statsig mode does not invalidate/retry; SyncQuota stops after the first unauthorized mode. + if calls.Load() != 1 { + t.Fatalf("blocked credential made %d quota requests, want 1", calls.Load()) + } + if updates := egressRepository.UpdateCount(); updates != 0 { + t.Fatalf("blocked account changed egress health %d times", updates) + } +} + +func TestSyncQuotaBlockedForbiddenSkipsStatsigRetryInURLMode(t *testing.T) { + var rateLimitCalls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + // URL mode may also probe the base URL for Statsig meta; only count quota endpoint. + if request.URL.Path == "/rest/rate-limits" { + rateLimitCalls.Add(1) + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusForbidden) + _, _ = writer.Write([]byte(`{"code":7,"message":"User is blocked [WKE=unauthorized:blocked-user]","details":[]}`)) + return + } + http.NotFound(writer, request) + })) + defer server.Close() + + cipher, err := security.NewCipher(base64.StdEncoding.EncodeToString(make([]byte, 32))) + if err != nil { + t.Fatal(err) + } + encrypted, err := cipher.Encrypt("blocked-sso-url") + if err != nil { + t.Fatal(err) + } + // URL mode would invalidate Statsig and retry unless blocked body is classified first. + adapter := NewAdapter(Config{ + BaseURL: server.URL, StatsigMode: "url", StatsigSignerURL: "https://signer.example/sign", + }, infraegress.NewManager(egressRepositoryStub{}, cipher), cipher, nil, nil) + _, err = adapter.SyncQuota(context.Background(), account.Credential{ID: 6, WebTier: account.WebTierAuto, EncryptedAccessToken: encrypted}) + if !errors.Is(err, provider.ErrUnauthorized) { + t.Fatalf("err = %v, want ErrUnauthorized", err) + } + if rateLimitCalls.Load() != 1 { + t.Fatalf("blocked credential made %d rate-limits requests, want 1 (no Statsig retry)", rateLimitCalls.Load()) + } +} + +func TestSyncQuotaGenericForbiddenIsNotUnauthorized(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusForbidden) + _, _ = writer.Write([]byte(`{"message":"temporary rejection"}`)) + })) + defer server.Close() + + cipher, err := security.NewCipher(base64.StdEncoding.EncodeToString(make([]byte, 32))) + if err != nil { + t.Fatal(err) + } + encrypted, err := cipher.Encrypt("generic-sso") + if err != nil { + t.Fatal(err) + } + adapter := NewAdapter(Config{ + BaseURL: server.URL, StatsigMode: "manual", StatsigManualValue: "test-signature", + }, infraegress.NewManager(egressRepositoryStub{}, cipher), cipher, nil, nil) + _, err = adapter.SyncQuota(context.Background(), account.Credential{ID: 5, WebTier: account.WebTierAuto, EncryptedAccessToken: encrypted}) + if err == nil || errors.Is(err, provider.ErrUnauthorized) { + t.Fatalf("err = %v, want generic forbidden error", err) + } +} + +func TestSyncWeeklyCreditsBlockedForbiddenIsUnauthorized(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusForbidden) + _, _ = writer.Write([]byte(`{"code":"unauthorized:blocked-user","error":"User is blocked"}`)) + })) + defer server.Close() + + cipher, err := security.NewCipher(base64.StdEncoding.EncodeToString(make([]byte, 32))) + if err != nil { + t.Fatal(err) + } + encrypted, err := cipher.Encrypt("blocked-weekly-sso") + if err != nil { + t.Fatal(err) + } + adapter := NewAdapter(Config{ + BaseURL: server.URL, StatsigMode: "manual", StatsigManualValue: "test-signature", + }, infraegress.NewManager(egressRepositoryStub{}, cipher), cipher, nil, nil) + _, err = adapter.SyncQuotaMode(context.Background(), account.Credential{ID: 7, EncryptedAccessToken: encrypted}, weeklyQuotaMode) + if !errors.Is(err, provider.ErrUnauthorized) { + t.Fatalf("err = %v, want ErrUnauthorized", err) + } + if calls.Load() != 1 { + t.Fatalf("blocked weekly credential made %d requests, want 1", calls.Load()) + } +} + +func TestInferWebTierFromUpstreamQuota(t *testing.T) { + tests := []struct { + name string + windows []account.QuotaWindow + want account.WebTier + known bool + }{ + {name: "current basic", windows: []account.QuotaWindow{{Mode: "auto", Total: 7}, {Mode: "fast", Total: 30}}, want: account.WebTierBasic, known: true}, + {name: "legacy basic", windows: []account.QuotaWindow{{Mode: "auto", Total: 20}}, want: account.WebTierBasic, known: true}, + {name: "super", windows: []account.QuotaWindow{{Mode: "auto", Total: 50}, {Mode: "fast", Total: 140}}, want: account.WebTierSuper, known: true}, + {name: "heavy", windows: []account.QuotaWindow{{Mode: "auto", Total: 150}, {Mode: "fast", Total: 400}}, want: account.WebTierHeavy, known: true}, + {name: "heavy mode", windows: []account.QuotaWindow{{Mode: "heavy", Total: 20}}, want: account.WebTierHeavy, known: true}, + {name: "conflicting signal uses lower tier", windows: []account.QuotaWindow{{Mode: "auto", Total: 50}, {Mode: "fast", Total: 30}}, want: account.WebTierBasic, known: true}, + {name: "unknown", windows: []account.QuotaWindow{{Mode: "auto", Total: 9}, {Mode: "fast", Total: 31}}, want: account.WebTierAuto, known: false}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, known := inferWebTierFromQuota(test.windows) + if got != test.want || known != test.known { + t.Fatalf("tier = %q, known = %v, want %q/%v", got, known, test.want, test.known) + } + }) + } +} + +func TestResolveWebTierUsesFreshWebQuotaOverStoredTier(t *testing.T) { + basicWindows := []account.QuotaWindow{{Mode: "auto", Total: 7}, {Mode: "fast", Total: 30}} + for _, stored := range []account.WebTier{account.WebTierAuto, account.WebTierBasic, account.WebTierSuper, account.WebTierHeavy} { + tier, useWeekly := resolveWebTierFromQuota(stored, basicWindows, true) + if tier != account.WebTierBasic || useWeekly { + t.Fatalf("stored %q resolved to %q, weekly=%v", stored, tier, useWeekly) + } + } + + tier, useWeekly := resolveWebTierFromQuota(account.WebTierBasic, []account.QuotaWindow{{Mode: "auto", Total: 50}}, true) + if tier != account.WebTierSuper || !useWeekly { + t.Fatalf("super snapshot resolved to %q, weekly=%v", tier, useWeekly) + } + + tier, useWeekly = resolveWebTierFromQuota(account.WebTierHeavy, nil, true) + if tier != account.WebTierHeavy || !useWeekly { + t.Fatalf("heavy weekly fallback resolved to %q, weekly=%v", tier, useWeekly) + } + + tier, useWeekly = resolveWebTierFromQuota(account.WebTierSuper, []account.QuotaWindow{{Mode: "auto", Total: 9}}, true) + if tier != account.WebTierAuto || useWeekly { + t.Fatalf("unknown snapshot resolved to %q, weekly=%v", tier, useWeekly) + } + + tier, useWeekly = resolveWebTierFromQuota(account.WebTierSuper, nil, true) + if tier != account.WebTierSuper || !useWeekly { + t.Fatalf("super weekly fallback resolved to %q, weekly=%v", tier, useWeekly) + } + + tier, useWeekly = resolveWebTierFromQuota(account.WebTierBasic, nil, true) + if tier != account.WebTierBasic || useWeekly { + t.Fatalf("basic should not be promoted when modes unavailable: got %q, weekly=%v", tier, useWeekly) + } + + tier, useWeekly = resolveWebTierFromQuota(account.WebTierAuto, nil, true) + if tier != account.WebTierAuto || useWeekly { + t.Fatalf("auto should not be promoted when modes unavailable: got %q, weekly=%v", tier, useWeekly) + } +} + +type recordingWebEgressRepository struct { + mu sync.Mutex + node egressdomain.Node + updates int +} + +func (r *recordingWebEgressRepository) ListEgressNodes(_ context.Context, scope egressdomain.Scope, _ repository.SortQuery) ([]egressdomain.Node, error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.node.Scope != scope { + return nil, nil + } + return []egressdomain.Node{r.node}, nil +} + +func (r *recordingWebEgressRepository) GetEgressNode(context.Context, uint64) (egressdomain.Node, error) { + r.mu.Lock() + defer r.mu.Unlock() + return r.node, nil +} + +func (r *recordingWebEgressRepository) CreateEgressNode(context.Context, egressdomain.Node) (egressdomain.Node, error) { + return egressdomain.Node{}, errors.New("unsupported") +} + +func (r *recordingWebEgressRepository) UpdateEgressNode(_ context.Context, value egressdomain.Node) (egressdomain.Node, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.node = value + r.updates++ + return value, nil +} + +func (r *recordingWebEgressRepository) DeleteEgressNode(context.Context, uint64) error { + return errors.New("unsupported") +} + +func (r *recordingWebEgressRepository) UpdateCount() int { + r.mu.Lock() + defer r.mu.Unlock() + return r.updates +} + +func TestSyncQuotaCorrectsStoredSuperFromFreshWebQuota(t *testing.T) { + var weeklyCalls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig": + weeklyCalls.Add(1) + http.Error(writer, "not available", http.StatusNotFound) + case "/rest/rate-limits": + var payload struct { + ModelName string `json:"modelName"` + } + if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { + t.Errorf("quota payload: %v", err) + } + total := 0 + switch payload.ModelName { + case "auto": + total = 7 + case "fast": + total = 30 + default: + http.Error(writer, "unsupported mode", http.StatusBadRequest) + return + } + writer.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(writer).Encode(map[string]any{ + "windowSizeSeconds": 7200, "remainingQueries": total, "totalQueries": total, + }) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + cipher, err := security.NewCipher(base64.StdEncoding.EncodeToString(make([]byte, 32))) + if err != nil { + t.Fatal(err) + } + encrypted, err := cipher.Encrypt("test-sso") + if err != nil { + t.Fatal(err) + } + adapter := NewAdapter(Config{ + BaseURL: server.URL, StatsigMode: "manual", StatsigManualValue: "test-signature", + }, infraegress.NewManager(egressRepositoryStub{}, cipher), cipher, nil, nil) + snapshot, err := adapter.SyncQuota(context.Background(), account.Credential{ + ID: 1, WebTier: account.WebTierSuper, EncryptedAccessToken: encrypted, + }) + if err != nil { + t.Fatal(err) + } + if snapshot.Tier != account.WebTierBasic || len(snapshot.Windows) != 2 || snapshot.Windows[0].Mode != "auto" || snapshot.Windows[0].Total != 7 || snapshot.Windows[1].Mode != "fast" || snapshot.Windows[1].Total != 30 { + t.Fatalf("snapshot = %#v", snapshot) + } + if weeklyCalls.Load() != 0 { + t.Fatalf("basic account probed weekly endpoint %d times", weeklyCalls.Load()) + } +} diff --git a/backend/internal/infra/provider/web/security_test.go b/backend/internal/infra/provider/web/security_test.go new file mode 100644 index 0000000000000000000000000000000000000000..6b263c74110e67d9987d7e299f69071cc5375776 --- /dev/null +++ b/backend/internal/infra/provider/web/security_test.go @@ -0,0 +1,79 @@ +package web + +import ( + "net/http" + "regexp" + "strings" + "testing" + + application "github.com/chenyme/grok2api/backend/internal/application/egress" + infraegress "github.com/chenyme/grok2api/backend/internal/infra/egress" +) + +var forbiddenBrowserIdentityFields = []string{ + "grok_device_id", "x-anonuserid", "x-userid", "x-challenge", "x-signature", +} + +func TestCloudflareCookieWhitelistDropsBrowserIdentityFields(t *testing.T) { + raw := "cf_clearance=clear; __cf_bm=bm; _cfuvid=uv; cf_chl_2=challenge; grok_device_id=device; x-anonuserid=anon; x-userid=user; x-challenge=c; x-signature=s; unrelated=value" + sanitized := application.SanitizeCloudflareCookies(raw) + for _, expected := range []string{"cf_clearance=clear", "__cf_bm=bm", "_cfuvid=uv", "cf_chl_2=challenge"} { + if !strings.Contains(sanitized, expected) { + t.Fatalf("sanitized cookies missing %q: %s", expected, sanitized) + } + } + assertForbiddenFieldsAbsent(t, sanitized) +} + +func TestWebHeadersOnlyUseSSOAndCloudflareCookies(t *testing.T) { + lease := &infraegress.Lease{UserAgent: "test-agent", CFCookies: "cf_clearance=clear; x-userid=drop"} + headers := buildHeaders("sso=token-value; x-userid=drop; x-signature=drop", lease, "application/json") + cookie := headers.Get("Cookie") + if !strings.Contains(cookie, "sso=token-value") || !strings.Contains(cookie, "sso-rw=token-value") || !strings.Contains(cookie, "cf_clearance=clear") { + t.Fatalf("cookie header = %q", cookie) + } + assertForbiddenFieldsAbsent(t, cookie) + for name := range headers { + assertForbiddenFieldsAbsent(t, name) + } + for _, name := range []string{"Accept", "Accept-Language", "x-xai-request-id"} { + if headers.Get(name) == "" { + t.Fatalf("missing header %s", name) + } + } + if value := headers.Get("x-statsig-id"); value != "" { + t.Fatalf("unsigned base headers must not contain x-statsig-id: %q", value) + } + if !regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$`).MatchString(headers.Get("x-xai-request-id")) { + t.Fatalf("x-xai-request-id = %q", headers.Get("x-xai-request-id")) + } + for name := range headers { + if strings.HasPrefix(strings.ToLower(name), "sec-ch-ua") { + t.Fatalf("synthetic client hint must not be sent: %s", name) + } + } +} + +func TestAppHeadersMatchStableBrowserFetchSignals(t *testing.T) { + headers := http.Header{} + applyAppHeaders(headers, "https://grok.com", "https://grok.com/") + for name, expected := range map[string]string{ + "Origin": "https://grok.com", "Referer": "https://grok.com/", "Cache-Control": "no-cache", + "Pragma": "no-cache", "Priority": "u=1, i", "Sec-Fetch-Dest": "empty", + "Sec-Fetch-Mode": "cors", "Sec-Fetch-Site": "same-origin", + } { + if headers.Get(name) != expected { + t.Fatalf("%s = %q", name, headers.Get(name)) + } + } +} + +func assertForbiddenFieldsAbsent(t *testing.T, value string) { + t.Helper() + lower := strings.ToLower(value) + for _, forbidden := range forbiddenBrowserIdentityFields { + if strings.Contains(lower, forbidden) { + t.Fatalf("forbidden browser identity field %q found in %q", forbidden, value) + } + } +} diff --git a/backend/internal/infra/provider/web/sso_build.go b/backend/internal/infra/provider/web/sso_build.go new file mode 100644 index 0000000000000000000000000000000000000000..b46ffc68e360a53ca3c972977b28f29dc8758c54 --- /dev/null +++ b/backend/internal/infra/provider/web/sso_build.go @@ -0,0 +1,382 @@ +package web + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "sort" + "strconv" + "strings" + "time" + + accountdomain "github.com/chenyme/grok2api/backend/internal/domain/account" + egressdomain "github.com/chenyme/grok2api/backend/internal/domain/egress" + infraegress "github.com/chenyme/grok2api/backend/internal/infra/egress" + "github.com/chenyme/grok2api/backend/internal/infra/provider" + "github.com/chenyme/grok2api/backend/internal/infra/security" +) + +const ( + ssoBuildClientID = "b1a00492-073a-47ea-816f-4c329264a828" + ssoBuildScope = "openid profile email offline_access grok-cli:access api:access conversations:read conversations:write" + ssoAccountsURL = "https://accounts.x.ai/" + ssoDeviceURL = "https://auth.x.ai/oauth2/device/code" + ssoVerifyURL = "https://auth.x.ai/oauth2/device/verify" + ssoApproveURL = "https://auth.x.ai/oauth2/device/approve" + ssoTokenURL = "https://auth.x.ai/oauth2/token" + maxAuthBody = 2 << 20 +) + +type ssoBuildHTTPClient interface { + Do(*http.Request) (*http.Response, error) +} + +type ssoBuildFlow struct { + client ssoBuildHTTPClient + userAgent string + cookies map[string]string +} + +func (a *Adapter) ConvertToBuild(ctx context.Context, credential accountdomain.Credential) (provider.CredentialSeed, error) { + if credential.Provider != accountdomain.ProviderWeb || credential.AuthType != accountdomain.AuthTypeSSO { + return provider.CredentialSeed{}, fmt.Errorf("仅 Grok Web SSO 账号支持转换") + } + token, err := a.cipher.Decrypt(credential.EncryptedAccessToken) + if err != nil { + return provider.CredentialSeed{}, fmt.Errorf("解密 Grok Web SSO: %w", err) + } + token = normalizeSSOToken(token) + if token == "" { + return provider.CredentialSeed{}, provider.ErrUnauthorized + } + lease, err := a.egress.AcquireCredential(ctx, egressdomain.ScopeWeb, credential) + if err != nil { + return provider.CredentialSeed{}, err + } + defer lease.Release() + requestCtx, cancel := context.WithTimeout(ctx, 90*time.Second) + defer cancel() + flow := &ssoBuildFlow{ + client: lease, userAgent: lease.UserAgent, + cookies: map[string]string{"sso": token, "sso-rw": token}, + } + seed, err := flow.convert(requestCtx, credential) + if err != nil { + a.egress.Feedback(context.WithoutCancel(ctx), lease.NodeID, conversionStatus(err), err) + return provider.CredentialSeed{}, err + } + a.egress.Feedback(context.WithoutCancel(ctx), lease.NodeID, http.StatusOK, nil) + return seed, nil +} + +func (f *ssoBuildFlow) convert(ctx context.Context, credential accountdomain.Credential) (provider.CredentialSeed, error) { + status, finalURL, _, err := f.do(ctx, http.MethodGet, ssoAccountsURL, nil) + if err != nil { + return provider.CredentialSeed{}, err + } + if status == http.StatusUnauthorized || strings.Contains(finalURL, "sign-in") || strings.Contains(finalURL, "sign-up") { + return provider.CredentialSeed{}, provider.ErrUnauthorized + } + if status < 200 || status >= 400 { + return provider.CredentialSeed{}, fmt.Errorf("校验 Grok Web SSO 失败: %w", conversionHTTPError{status: status}) + } + + form := url.Values{"client_id": {ssoBuildClientID}, "scope": {ssoBuildScope}} + status, _, body, err := f.do(ctx, http.MethodPost, ssoDeviceURL, form) + if err != nil { + return provider.CredentialSeed{}, err + } + if status < 200 || status >= 300 { + return provider.CredentialSeed{}, fmt.Errorf("xAI Device Flow 启动失败: %w", conversionHTTPError{status: status}) + } + var device struct { + DeviceCode string `json:"device_code"` + UserCode string `json:"user_code"` + VerificationURIComplete string `json:"verification_uri_complete"` + Interval int `json:"interval"` + ExpiresIn int `json:"expires_in"` + } + if err := json.Unmarshal(body, &device); err != nil { + return provider.CredentialSeed{}, fmt.Errorf("解析 xAI Device Flow: %w", err) + } + if device.DeviceCode == "" || device.UserCode == "" || !safeXAIURL(device.VerificationURIComplete) { + return provider.CredentialSeed{}, fmt.Errorf("xAI Device Flow 返回字段不完整") + } + if device.Interval <= 0 { + device.Interval = 5 + } + if device.ExpiresIn <= 0 { + device.ExpiresIn = 1800 + } + + status, finalURL, _, err = f.do(ctx, http.MethodGet, device.VerificationURIComplete, nil) + if err != nil { + return provider.CredentialSeed{}, err + } + if status < 200 || status >= 400 { + return provider.CredentialSeed{}, fmt.Errorf("打开 Device Flow 验证页失败: %w", conversionHTTPError{status: status}) + } + status, finalURL, _, err = f.do(ctx, http.MethodPost, ssoVerifyURL, url.Values{"user_code": {device.UserCode}}) + if err != nil { + return provider.CredentialSeed{}, err + } + if status < 200 || status >= 400 { + return provider.CredentialSeed{}, fmt.Errorf("SSO 自动验证 Device Flow 失败: %w", conversionHTTPError{status: status}) + } + if !strings.Contains(finalURL, "consent") { + return provider.CredentialSeed{}, fmt.Errorf("SSO 自动验证 Device Flow 失败") + } + status, finalURL, _, err = f.do(ctx, http.MethodPost, ssoApproveURL, url.Values{ + "user_code": {device.UserCode}, "action": {"allow"}, "principal_type": {"User"}, "principal_id": {""}, + }) + if err != nil { + return provider.CredentialSeed{}, err + } + if status < 200 || status >= 400 { + return provider.CredentialSeed{}, fmt.Errorf("SSO 自动批准 Device Flow 失败: %w", conversionHTTPError{status: status}) + } + if !strings.Contains(finalURL, "done") { + return provider.CredentialSeed{}, fmt.Errorf("SSO 自动批准 Device Flow 失败") + } + + token, err := f.pollToken(ctx, device.DeviceCode, time.Duration(device.Interval)*time.Second, time.Duration(device.ExpiresIn)*time.Second) + if err != nil { + return provider.CredentialSeed{}, err + } + claims := decodeBuildClaims(firstValue(token.IDToken, token.AccessToken)) + userID := claimString(claims, "sub") + email := claimString(claims, "email") + teamID := claimString(claims, "team_id") + name := strings.TrimSpace(credential.Name) + if name == "" { + name = "Grok Web account" + } + return provider.CredentialSeed{ + Provider: accountdomain.ProviderBuild, AuthType: accountdomain.AuthTypeOAuth, + Name: firstValue(email, name+" Build", userID, "Grok Build account"), Email: email, UserID: userID, TeamID: teamID, + SourceKey: "sso-build:" + security.HashToken(token.AccessToken), OIDCClientID: ssoBuildClientID, + AccessToken: token.AccessToken, RefreshToken: token.RefreshToken, ExpiresAt: token.ExpiresAt, + }, nil +} + +type ssoBuildToken struct { + AccessToken string + RefreshToken string + IDToken string + ExpiresAt time.Time +} + +func (f *ssoBuildFlow) pollToken(ctx context.Context, deviceCode string, interval, expiresIn time.Duration) (ssoBuildToken, error) { + if interval < time.Second { + interval = time.Second + } + deadline := time.Now().Add(min(expiresIn, 75*time.Second)) + for time.Now().Before(deadline) { + timer := time.NewTimer(interval) + select { + case <-ctx.Done(): + timer.Stop() + return ssoBuildToken{}, ctx.Err() + case <-timer.C: + } + status, _, body, err := f.do(ctx, http.MethodPost, ssoTokenURL, url.Values{ + "grant_type": {"urn:ietf:params:oauth:grant-type:device_code"}, "client_id": {ssoBuildClientID}, "device_code": {deviceCode}, + }) + if err != nil { + return ssoBuildToken{}, err + } + var payload struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + IDToken string `json:"id_token"` + ExpiresIn int `json:"expires_in"` + Error string `json:"error"` + ErrorDescription string `json:"error_description"` + } + if err := json.Unmarshal(body, &payload); err != nil { + return ssoBuildToken{}, fmt.Errorf("解析 xAI OAuth Token: %w", err) + } + if status >= 200 && status < 300 && payload.AccessToken != "" { + if payload.ExpiresIn <= 0 { + payload.ExpiresIn = 3600 + } + return ssoBuildToken{AccessToken: payload.AccessToken, RefreshToken: payload.RefreshToken, IDToken: payload.IDToken, ExpiresAt: time.Now().UTC().Add(time.Duration(payload.ExpiresIn) * time.Second)}, nil + } + switch payload.Error { + case "authorization_pending": + continue + case "slow_down": + interval += 5 * time.Second + continue + case "access_denied", "expired_token": + return ssoBuildToken{}, provider.ErrAuthorizationDenied + default: + if status >= 400 { + return ssoBuildToken{}, fmt.Errorf("xAI OAuth Token 失败 (%s): %w", firstValue(payload.ErrorDescription, payload.Error), conversionHTTPError{status: status}) + } + return ssoBuildToken{}, fmt.Errorf("xAI OAuth Token 失败: %s", firstValue(payload.ErrorDescription, payload.Error, strconv.Itoa(status))) + } + } + return ssoBuildToken{}, fmt.Errorf("xAI Device Flow 轮询超时") +} + +func (f *ssoBuildFlow) do(ctx context.Context, method, endpoint string, form url.Values) (int, string, []byte, error) { + if !safeXAIURL(endpoint) { + return 0, "", nil, fmt.Errorf("xAI OAuth URL 不安全") + } + currentURL := endpoint + currentMethod := method + currentForm := form + for redirects := 0; redirects <= 8; redirects++ { + var body io.Reader + if currentForm != nil { + body = strings.NewReader(currentForm.Encode()) + } + request, err := http.NewRequestWithContext(ctx, currentMethod, currentURL, body) + if err != nil { + return 0, "", nil, err + } + request.Header.Set("Accept", "application/json, text/html;q=0.9, */*;q=0.8") + request.Header.Set("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8") + request.Header.Set("User-Agent", f.userAgent) + request.Header.Set("Cookie", f.cookieHeader()) + if currentForm != nil { + request.Header.Set("Content-Type", "application/x-www-form-urlencoded") + } + response, err := f.client.Do(request) + if err != nil { + return 0, "", nil, err + } + f.captureCookies(response) + data, readErr := io.ReadAll(io.LimitReader(response.Body, maxAuthBody+1)) + _ = response.Body.Close() + if readErr != nil { + return response.StatusCode, currentURL, nil, readErr + } + if len(data) > maxAuthBody { + return response.StatusCode, currentURL, nil, fmt.Errorf("xAI OAuth 响应超过 2 MiB") + } + if response.StatusCode < 300 || response.StatusCode > 399 { + return response.StatusCode, currentURL, data, nil + } + location := strings.TrimSpace(response.Header.Get("Location")) + if location == "" { + return response.StatusCode, currentURL, data, fmt.Errorf("xAI OAuth 重定向缺少 Location") + } + base, _ := url.Parse(currentURL) + next, err := url.Parse(location) + if err != nil { + return response.StatusCode, currentURL, data, err + } + currentURL = base.ResolveReference(next).String() + if !safeXAIURL(currentURL) { + return response.StatusCode, currentURL, data, fmt.Errorf("xAI OAuth 重定向到非受信域名") + } + if response.StatusCode == http.StatusSeeOther || ((response.StatusCode == http.StatusMovedPermanently || response.StatusCode == http.StatusFound) && currentMethod != http.MethodGet && currentMethod != http.MethodHead) { + currentMethod = http.MethodGet + currentForm = nil + } + } + return 0, currentURL, nil, fmt.Errorf("xAI OAuth 重定向次数过多") +} + +func (f *ssoBuildFlow) captureCookies(response *http.Response) { + for _, cookie := range response.Cookies() { + name := strings.TrimSpace(cookie.Name) + value := strings.TrimSpace(cookie.Value) + if name == "" || len(name) > 128 || len(value) > 16384 || strings.ContainsAny(name+value, "\r\n\x00") { + continue + } + if cookie.MaxAge < 0 { + delete(f.cookies, name) + continue + } + f.cookies[name] = value + } +} + +func (f *ssoBuildFlow) cookieHeader() string { + keys := make([]string, 0, len(f.cookies)) + for key := range f.cookies { + keys = append(keys, key) + } + sort.Strings(keys) + parts := make([]string, 0, len(keys)) + for _, key := range keys { + parts = append(parts, key+"="+f.cookies[key]) + } + return strings.Join(parts, "; ") +} + +func safeXAIURL(raw string) bool { + parsed, err := url.Parse(raw) + if err != nil || parsed.Scheme != "https" || parsed.User != nil || parsed.Hostname() == "" { + return false + } + host := strings.ToLower(parsed.Hostname()) + return host == "x.ai" || strings.HasSuffix(host, ".x.ai") +} + +func normalizeSSOToken(value string) string { + value = strings.TrimSpace(value) + if strings.HasPrefix(strings.ToLower(value), "sso=") { + value = strings.TrimSpace(value[len("sso="):]) + } + if token, _, found := strings.Cut(value, ";"); found { + value = strings.TrimSpace(token) + } + return strings.NewReplacer("\r", "", "\n", "", "\x00", "").Replace(value) +} + +func decodeBuildClaims(token string) map[string]any { + parts := strings.Split(token, ".") + if len(parts) < 2 { + return nil + } + data, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + return nil + } + var claims map[string]any + if json.Unmarshal(data, &claims) != nil { + return nil + } + return claims +} + +func claimString(claims map[string]any, key string) string { + value, _ := claims[key].(string) + return strings.TrimSpace(value) +} + +func firstValue(values ...string) string { + for _, value := range values { + if value = strings.TrimSpace(value); value != "" { + return value + } + } + return "" +} + +type conversionHTTPError struct{ status int } + +func (e conversionHTTPError) Error() string { return fmt.Sprintf("xAI OAuth HTTP %d", e.status) } + +func conversionStatus(err error) int { + var statusErr conversionHTTPError + if errors.As(err, &statusErr) { + return statusErr.status + } + if errors.Is(err, provider.ErrUnauthorized) { + return http.StatusUnauthorized + } + return 0 +} + +var _ provider.BuildCredentialConverter = (*Adapter)(nil) +var _ ssoBuildHTTPClient = (*infraegress.Lease)(nil) diff --git a/backend/internal/infra/provider/web/sso_build_test.go b/backend/internal/infra/provider/web/sso_build_test.go new file mode 100644 index 0000000000000000000000000000000000000000..0e4e54f4615d89756aecbbd1f6ccb4b4a38a7693 --- /dev/null +++ b/backend/internal/infra/provider/web/sso_build_test.go @@ -0,0 +1,65 @@ +package web + +import ( + "context" + "io" + "net/http" + "strings" + "testing" +) + +type scriptedSSOClient struct { + responses []*http.Response + requests []*http.Request +} + +func (c *scriptedSSOClient) Do(request *http.Request) (*http.Response, error) { + c.requests = append(c.requests, request) + response := c.responses[0] + c.responses = c.responses[1:] + return response, nil +} + +func TestSSOBuildFlowFollowsOnlyTrustedXAIHTTPSRedirects(t *testing.T) { + client := &scriptedSSOClient{responses: []*http.Response{ + {StatusCode: http.StatusFound, Header: http.Header{"Location": []string{"https://auth.x.ai/next"}, "Set-Cookie": []string{"session=abc; Path=/; Secure"}}, Body: io.NopCloser(strings.NewReader(""))}, + {StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader("ok"))}, + }} + flow := &ssoBuildFlow{client: client, userAgent: "test-agent", cookies: map[string]string{"sso": "secret"}} + status, finalURL, body, err := flow.do(context.Background(), http.MethodGet, ssoAccountsURL, nil) + if err != nil { + t.Fatal(err) + } + if status != http.StatusOK || finalURL != "https://auth.x.ai/next" || string(body) != "ok" { + t.Fatalf("response = %d %s %q", status, finalURL, body) + } + if len(client.requests) != 2 || client.requests[1].Header.Get("User-Agent") != "test-agent" { + t.Fatalf("requests = %#v", client.requests) + } + cookie := client.requests[1].Header.Get("Cookie") + if !strings.Contains(cookie, "sso=secret") || !strings.Contains(cookie, "session=abc") { + t.Fatalf("redirect cookies = %q", cookie) + } + + unsafe := &scriptedSSOClient{responses: []*http.Response{{StatusCode: http.StatusFound, Header: http.Header{"Location": []string{"https://example.com/steal"}}, Body: io.NopCloser(strings.NewReader(""))}}} + flow = &ssoBuildFlow{client: unsafe, userAgent: "test-agent", cookies: map[string]string{"sso": "secret"}} + if _, _, _, err := flow.do(context.Background(), http.MethodGet, ssoAccountsURL, nil); err == nil { + t.Fatal("unsafe redirect was accepted") + } +} + +func TestSSOBuildConversionSanitizesTokenAndURLs(t *testing.T) { + if token := normalizeSSOToken("sso=token-value; x-userid=drop"); token != "token-value" { + t.Fatalf("token = %q", token) + } + for _, value := range []string{"https://accounts.x.ai/", "https://auth.x.ai/oauth2/device/code"} { + if !safeXAIURL(value) { + t.Fatalf("trusted URL rejected: %s", value) + } + } + for _, value := range []string{"http://auth.x.ai/", "https://x.ai.example.com/", "https://user@auth.x.ai/"} { + if safeXAIURL(value) { + t.Fatalf("unsafe URL accepted: %s", value) + } + } +} diff --git a/backend/internal/infra/provider/web/statsig.go b/backend/internal/infra/provider/web/statsig.go new file mode 100644 index 0000000000000000000000000000000000000000..3f3e85fd7f7fd51d5356c8c20398c638167548ba --- /dev/null +++ b/backend/internal/infra/provider/web/statsig.go @@ -0,0 +1,433 @@ +package web + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "sync" + "time" + + "github.com/chenyme/grok2api/backend/internal/domain/account" + domainegress "github.com/chenyme/grok2api/backend/internal/domain/egress" + infraegress "github.com/chenyme/grok2api/backend/internal/infra/egress" + "github.com/chenyme/grok2api/backend/internal/pkg/signerurl" + "golang.org/x/net/html" + "golang.org/x/sync/singleflight" +) + +const ( + defaultStatsigSignerURL = "https://grok.wodf.de/sign" + statsigCacheTTL = time.Hour + statsigCacheMaxEntries = 4096 + statsigMetaBodyLimit = 4 << 20 + statsigResponseLimit = 4 << 10 +) + +type statsigCacheEntry struct { + value string + expiresAt time.Time +} + +type statsigSignResult struct { + value string + source string +} + +type statsigWarmTarget struct { + method string + target string +} + +type statsigSigner struct { + client *http.Client + fetchMeta func(context.Context, string, string, *infraegress.Lease) (string, error) + validateEndpoint func(context.Context, string) error + now func() time.Time + mu sync.Mutex + entries map[string]statsigCacheEntry + refreshes singleflight.Group +} + +func newStatsigSigner() *statsigSigner { + return &statsigSigner{ + client: &http.Client{ + Timeout: 12 * time.Second, + CheckRedirect: func(_ *http.Request, _ []*http.Request) error { return http.ErrUseLastResponse }, + }, + fetchMeta: fetchStatsigMetaContent, + validateEndpoint: validateStatsigSignerEndpoint, + now: time.Now, + entries: make(map[string]statsigCacheEntry), + } +} + +func (s *statsigSigner) Sign(ctx context.Context, baseURL, signerURL, token string, lease *infraegress.Lease, method, target string) (string, string, error) { + key, path, err := statsigSignatureKey(baseURL, signerURL, method, target) + if err != nil { + return "", "", err + } + if value, ok := s.cached(key, s.now().UTC()); ok { + return value, "cache", nil + } + value, err, _ := s.refreshes.Do(key, func() (any, error) { + now := s.now().UTC() + if cached, ok := s.cached(key, now); ok { + return statsigSignResult{value: cached, source: "cache"}, nil + } + fresh, refreshErr := s.freshSignature(ctx, baseURL, signerURL, token, lease, method, path) + if refreshErr != nil { + if stale, ok := s.stale(key); ok { + return statsigSignResult{value: stale, source: "stale"}, nil + } + return statsigSignResult{}, refreshErr + } + s.store(key, fresh, now.Add(statsigCacheTTL), now) + return statsigSignResult{value: fresh, source: "refresh"}, nil + }) + if err != nil { + return "", "", err + } + result := value.(statsigSignResult) + return result.value, result.source, nil +} + +// Warm 使用一次 metaContent 请求预热多个常用签名键，避免按账号或按路径重复抓取首页。 +func (s *statsigSigner) Warm(ctx context.Context, baseURL, signerURL, token string, lease *infraegress.Lease, targets []statsigWarmTarget) (int, error) { + now := s.now().UTC() + type pendingTarget struct { + key string + method string + path string + } + pending := make([]pendingTarget, 0, len(targets)) + for _, target := range targets { + key, path, err := statsigSignatureKey(baseURL, signerURL, target.method, target.target) + if err != nil { + return 0, err + } + if _, ok := s.cached(key, now); ok { + continue + } + pending = append(pending, pendingTarget{key: key, method: target.method, path: path}) + } + if len(pending) == 0 { + return 0, nil + } + meta, err := s.fetchMeta(ctx, baseURL, token, lease) + if err != nil { + return 0, err + } + warmed := 0 + for _, target := range pending { + value, signErr := s.requestSignature(ctx, signerURL, target.method, target.path, meta) + if signErr != nil { + return warmed, signErr + } + s.store(target.key, value, now.Add(statsigCacheTTL), now) + warmed++ + } + return warmed, nil +} + +func (s *statsigSigner) freshSignature(ctx context.Context, baseURL, signerURL, token string, lease *infraegress.Lease, method, path string) (string, error) { + meta, err := s.fetchMeta(ctx, baseURL, token, lease) + if err != nil { + return "", err + } + signature, err := s.requestSignature(ctx, signerURL, method, path, meta) + if err == nil { + return signature, nil + } + + meta, refreshErr := s.fetchMeta(ctx, baseURL, token, lease) + if refreshErr != nil { + return "", fmt.Errorf("刷新 Statsig metaContent: %w", refreshErr) + } + signature, retryErr := s.requestSignature(ctx, signerURL, method, path, meta) + if retryErr != nil { + return "", fmt.Errorf("Statsig 签名失败: %w", retryErr) + } + return signature, nil +} + +func (s *statsigSigner) Invalidate(baseURL, signerURL, method, target string) { + key, _, err := statsigSignatureKey(baseURL, signerURL, method, target) + if err != nil { + return + } + s.mu.Lock() + delete(s.entries, key) + s.mu.Unlock() +} + +func (s *statsigSigner) Clear() { + s.mu.Lock() + clear(s.entries) + s.mu.Unlock() +} + +func (s *statsigSigner) cached(key string, now time.Time) (string, bool) { + s.mu.Lock() + defer s.mu.Unlock() + entry, ok := s.entries[key] + if !ok || entry.value == "" || !now.Before(entry.expiresAt) { + return "", false + } + return entry.value, true +} + +func (s *statsigSigner) stale(key string) (string, bool) { + s.mu.Lock() + defer s.mu.Unlock() + entry, ok := s.entries[key] + return entry.value, ok && validStatsigID(entry.value) +} + +func (s *statsigSigner) store(key, value string, expiresAt, now time.Time) { + s.mu.Lock() + defer s.mu.Unlock() + for existingKey, entry := range s.entries { + if !now.Before(entry.expiresAt) { + delete(s.entries, existingKey) + } + } + if len(s.entries) >= statsigCacheMaxEntries { + oldestKey := "" + var oldestExpiry time.Time + for existingKey, entry := range s.entries { + if oldestKey == "" || entry.expiresAt.Before(oldestExpiry) { + oldestKey, oldestExpiry = existingKey, entry.expiresAt + } + } + delete(s.entries, oldestKey) + } + s.entries[key] = statsigCacheEntry{value: value, expiresAt: expiresAt} +} + +func statsigSignatureKey(baseURL, signerURL, method, target string) (string, string, error) { + parsed, err := url.Parse(target) + if err != nil { + return "", "", fmt.Errorf("解析 Statsig 目标地址: %w", err) + } + path := parsed.EscapedPath() + if path == "" { + path = "/" + } + method = strings.ToUpper(strings.TrimSpace(method)) + return strings.TrimRight(baseURL, "/") + "\x00" + strings.TrimSpace(signerURL) + "\x00" + method + "\x00" + path, path, nil +} + +func (s *statsigSigner) requestSignature(ctx context.Context, endpoint, method, path, metaContent string) (string, error) { + if err := s.validateEndpoint(ctx, endpoint); err != nil { + return "", err + } + payload, _ := json.Marshal(map[string]any{ + "method": strings.ToUpper(strings.TrimSpace(method)), + "path": path, + "environment": map[string]string{ + "metaContent": metaContent, + }, + }) + request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload)) + if err != nil { + return "", err + } + request.Header.Set("Content-Type", "application/json") + response, err := s.client.Do(request) + if err != nil { + return "", err + } + defer response.Body.Close() + body, err := io.ReadAll(io.LimitReader(response.Body, statsigResponseLimit+1)) + if err != nil { + return "", err + } + if len(body) > statsigResponseLimit { + return "", fmt.Errorf("签名响应超过安全上限") + } + if response.StatusCode < 200 || response.StatusCode >= 300 { + return "", fmt.Errorf("签名服务返回 %d", response.StatusCode) + } + var value struct { + StatsigID string `json:"x-statsig-id"` + } + if json.Unmarshal(body, &value) != nil || !validStatsigID(value.StatsigID) { + return "", fmt.Errorf("签名服务响应无效") + } + return value.StatsigID, nil +} + +func validateStatsigSignerEndpoint(ctx context.Context, endpoint string) error { + _ = ctx + return signerurl.Validate(endpoint) +} + +func fetchStatsigMetaContent(ctx context.Context, baseURL, token string, lease *infraegress.Lease) (string, error) { + if lease == nil { + return "", fmt.Errorf("Statsig 获取缺少出口租约") + } + requestCtx, cancel := context.WithTimeout(infraegress.WithPhysicalCallStage(ctx, "statsig_meta"), 15*time.Second) + defer cancel() + request, err := http.NewRequestWithContext(requestCtx, http.MethodGet, strings.TrimRight(baseURL, "/")+"/index", nil) + if err != nil { + return "", err + } + request.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8") + request.Header.Set("Accept-Encoding", "gzip, deflate, br, zstd") + request.Header.Set("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8") + request.Header.Set("Cache-Control", "no-cache") + request.Header.Set("Pragma", "no-cache") + request.Header.Set("Sec-Fetch-Dest", "document") + request.Header.Set("Sec-Fetch-Mode", "navigate") + request.Header.Set("Sec-Fetch-Site", "same-origin") + request.Header.Set("Upgrade-Insecure-Requests", "1") + request.Header.Set("User-Agent", lease.UserAgent) + request.Header.Set("Cookie", infraegress.BuildSSOCookie(token, lease.CFCookies)) + response, err := lease.Do(request) + if err != nil { + return "", err + } + defer response.Body.Close() + if response.StatusCode < 200 || response.StatusCode >= 300 { + return "", fmt.Errorf("Grok index 返回 %d", response.StatusCode) + } + body, err := io.ReadAll(io.LimitReader(response.Body, statsigMetaBodyLimit+1)) + if err != nil { + return "", err + } + if len(body) > statsigMetaBodyLimit { + return "", fmt.Errorf("Grok index 超过安全上限") + } + content, err := extractStatsigMetaContent(body) + if err != nil { + return "", err + } + return content, nil +} + +func extractStatsigMetaContent(body []byte) (string, error) { + tokenizer := html.NewTokenizer(bytes.NewReader(body)) + for { + switch tokenizer.Next() { + case html.ErrorToken: + if tokenizer.Err() == io.EOF { + return "", fmt.Errorf("Grok index 缺少 grok-site-verification") + } + return "", tokenizer.Err() + case html.StartTagToken, html.SelfClosingTagToken: + name, hasAttrs := tokenizer.TagName() + if !strings.EqualFold(string(name), "meta") || !hasAttrs { + continue + } + metaName := "" + content := "" + for { + key, value, more := tokenizer.TagAttr() + switch strings.ToLower(string(key)) { + case "name": + metaName = normalizeStatsigMetaName(string(value)) + case "content": + content = strings.TrimSpace(string(value)) + } + if !more { + break + } + } + if metaName == "grok-site-verification" && content != "" { + return content, nil + } + } + } +} + +func normalizeStatsigMetaName(value string) string { + value = strings.ToLower(strings.TrimSpace(value)) + return strings.NewReplacer("‐", "-", "‑", "-", "‒", "-", "–", "-", "—", "-", "―", "-").Replace(value) +} + +func validStatsigID(value string) bool { + value = strings.TrimSpace(value) + if value == "" { + return false + } + decoded, err := base64.RawStdEncoding.DecodeString(value) + if err != nil { + decoded, err = base64.StdEncoding.DecodeString(value) + } + return err == nil && len(decoded) == 70 +} + +func (a *Adapter) applySignedStatsig(ctx context.Context, request *http.Request, token string, lease *infraegress.Lease) { + if request == nil { + return + } + cfg := a.config() + request.Header.Del("x-statsig-id") + if cfg.StatsigMode == "manual" { + if value := strings.TrimSpace(cfg.StatsigManualValue); validStatsigID(value) { + request.Header.Set("x-statsig-id", value) + } + return + } + if a.statsig == nil { + return + } + value, source, err := a.statsig.Sign(ctx, cfg.BaseURL, cfg.StatsigSignerURL, token, lease, request.Method, request.URL.String()) + if err == nil { + request.Header.Set("x-statsig-id", value) + if source == "refresh" { + a.log().Info("web_statsig_refreshed", "method", request.Method, "path", request.URL.EscapedPath()) + } else if source == "stale" { + a.log().Warn("web_statsig_refresh_failed_using_stale", "method", request.Method, "path", request.URL.EscapedPath()) + } + return + } + a.log().Warn("web_statsig_fetch_failed", "method", request.Method, "path", request.URL.EscapedPath(), "error", err) +} + +// WarmStatsig 只使用一个 Web 账号和一个出口租约预热共享签名，不会逐账号访问上游。 +func (a *Adapter) WarmStatsig(ctx context.Context, credential account.Credential) (int, error) { + cfg := a.config() + if cfg.StatsigMode == "manual" { + if !validStatsigID(strings.TrimSpace(cfg.StatsigManualValue)) { + return 0, fmt.Errorf("手动 Statsig 配置无效") + } + return 0, nil + } + if a.statsig == nil { + return 0, fmt.Errorf("Statsig 签名器未初始化") + } + token, err := a.cipher.Decrypt(credential.EncryptedAccessToken) + if err != nil { + return 0, err + } + lease, err := a.egress.AcquireCredential(ctx, domainegress.ScopeWeb, credential) + if err != nil { + return 0, err + } + defer lease.Release() + baseURL := strings.TrimRight(cfg.BaseURL, "/") + return a.statsig.Warm(ctx, cfg.BaseURL, cfg.StatsigSignerURL, token, lease, []statsigWarmTarget{ + {method: http.MethodPost, target: baseURL + "/rest/app-chat/conversations/new"}, + {method: http.MethodPost, target: baseURL + "/rest/rate-limits"}, + {method: http.MethodPost, target: baseURL + "/rest/media/post/create"}, + }) +} + +func (a *Adapter) invalidateSignedStatsig(method, target string) bool { + cfg := a.config() + if cfg.StatsigMode == "url" && a.statsig != nil { + a.statsig.Invalidate(cfg.BaseURL, cfg.StatsigSignerURL, method, target) + if parsed, err := url.Parse(target); err == nil { + a.log().Info("web_statsig_invalidated", "method", method, "path", parsed.EscapedPath()) + } + return true + } + return false +} diff --git a/backend/internal/infra/provider/web/statsig_test.go b/backend/internal/infra/provider/web/statsig_test.go new file mode 100644 index 0000000000000000000000000000000000000000..b32444e6a80b86d454d4634d20c6455f2740e7e4 --- /dev/null +++ b/backend/internal/infra/provider/web/statsig_test.go @@ -0,0 +1,273 @@ +package web + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "testing" + "time" + + infraegress "github.com/chenyme/grok2api/backend/internal/infra/egress" +) + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { return f(request) } + +func TestExtractStatsigMetaContentAcceptsCurrentMetaName(t *testing.T) { + for _, name := range []string{"grok-site―verification", "grok-site-verification"} { + body := []byte(``) + value, err := extractStatsigMetaContent(body) + if err != nil || value != "meta-value" { + t.Fatalf("name=%q value=%q err=%v", name, value, err) + } + } +} + +func TestStatsigSignerSendsMethodPathAndMetaContent(t *testing.T) { + raw := make([]byte, 70) + encoded := base64.RawStdEncoding.EncodeToString(raw) + signer := newStatsigSigner() + signer.validateEndpoint = func(context.Context, string) error { return nil } + signer.client = &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + var payload struct { + Method string `json:"method"` + Path string `json:"path"` + Environment struct { + MetaContent string `json:"metaContent"` + } `json:"environment"` + } + if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { + t.Fatal(err) + } + if payload.Method != "POST" || payload.Path != "/rest/app-chat/conversations/id/responses" || payload.Environment.MetaContent != "meta-value" { + t.Fatalf("payload=%#v", payload) + } + body, _ := json.Marshal(map[string]string{"x-statsig-id": encoded}) + return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(string(body))), Header: http.Header{}}, nil + })} + value, err := signer.requestSignature(context.Background(), "https://signer.example/sign", "post", "/rest/app-chat/conversations/id/responses", "meta-value") + if err != nil || value != encoded { + t.Fatalf("value=%q err=%v", value, err) + } +} + +func TestStatsigSignerRejectsInvalidShape(t *testing.T) { + signer := newStatsigSigner() + signer.validateEndpoint = func(context.Context, string) error { return nil } + signer.client = &http.Client{Transport: roundTripFunc(func(_ *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"x-statsig-id":"invalid"}`)), Header: http.Header{}}, nil + })} + if _, err := signer.requestSignature(context.Background(), "https://signer.example/sign", "POST", "/rest/test", "meta"); err == nil { + t.Fatal("invalid signature was accepted") + } +} + +func TestValidateStatsigSignerEndpointUsesAdminURLBoundary(t *testing.T) { + for _, endpoint := range []string{ + "https://grok.wodf.de/sign", + "https://signer.example/sign", + "http://grok-signer-go:8788/sign", + "http://host.docker.internal:8788/sign", + "http://127.0.0.1:8788/sign", + "https://10.0.0.1:8443/sign", + } { + if err := validateStatsigSignerEndpoint(context.Background(), endpoint); err != nil { + t.Fatalf("endpoint %q rejected: %v", endpoint, err) + } + } + for _, endpoint := range []string{ + "http://grok.wodf.de/sign", + "https://user:pass@grok.wodf.de/sign", + "https://grok.wodf.de:8443/sign", + "https://grok.wodf.de/sign?token=value", + "http://8.8.8.8:8788/sign", + } { + if err := validateStatsigSignerEndpoint(context.Background(), endpoint); err == nil { + t.Fatalf("unsafe endpoint %q accepted", endpoint) + } + } +} + +func TestStatsigSignerClientRejectsRedirects(t *testing.T) { + signer := newStatsigSigner() + request, err := http.NewRequest(http.MethodGet, "http://grok-signer-go:8788/redirect", nil) + if err != nil { + t.Fatal(err) + } + if err := signer.client.CheckRedirect(request, []*http.Request{request}); !errors.Is(err, http.ErrUseLastResponse) { + t.Fatalf("redirect policy error = %v", err) + } +} + +func TestStatsigSignerCachesByMethodAndPathForOneHour(t *testing.T) { + var fetches int + var signedMeta []string + now := time.Date(2026, 7, 12, 12, 0, 0, 0, time.UTC) + signer := newStatsigSigner() + signer.now = func() time.Time { return now } + signer.validateEndpoint = func(context.Context, string) error { return nil } + signer.fetchMeta = func(context.Context, string, string, *infraegress.Lease) (string, error) { + fetches++ + return fmt.Sprintf("meta-%d", fetches), nil + } + signer.client = &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + var payload struct { + Environment struct { + MetaContent string `json:"metaContent"` + } `json:"environment"` + } + if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { + t.Fatal(err) + } + signedMeta = append(signedMeta, payload.Environment.MetaContent) + raw := make([]byte, 70) + raw[0] = byte(len(signedMeta)) + encoded := base64.RawStdEncoding.EncodeToString(raw) + body, _ := json.Marshal(map[string]string{"x-statsig-id": encoded}) + return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(string(body))), Header: http.Header{}}, nil + })} + first, _, err := signer.Sign(context.Background(), "https://grok.com", "https://signer.example/sign", "token-a", nil, http.MethodPost, "https://grok.com/rest/test") + if err != nil { + t.Fatal(err) + } + second, _, err := signer.Sign(context.Background(), "https://grok.com", "https://signer.example/sign", "token-b", nil, http.MethodPost, "https://grok.com/rest/test") + if err != nil { + t.Fatal(err) + } + if fetches != 1 || len(signedMeta) != 1 || first != second { + t.Fatalf("cached fetches=%d signedMeta=%v first=%q second=%q", fetches, signedMeta, first, second) + } + + now = now.Add(time.Hour) + third, _, err := signer.Sign(context.Background(), "https://grok.com", "https://signer.example/sign", "token-b", nil, http.MethodPost, "https://grok.com/rest/test") + if err != nil { + t.Fatal(err) + } + if fetches != 2 || third == first { + t.Fatalf("hourly refresh fetches=%d first=%q third=%q", fetches, first, third) + } + + signer.Invalidate("https://grok.com", "https://signer.example/sign", http.MethodPost, "https://grok.com/rest/test") + fourth, _, err := signer.Sign(context.Background(), "https://grok.com", "https://signer.example/sign", "token-a", nil, http.MethodPost, "https://grok.com/rest/test") + if err != nil { + t.Fatal(err) + } + if fetches != 3 || fourth == third { + t.Fatalf("invalidation fetches=%d third=%q fourth=%q", fetches, third, fourth) + } + + if _, _, err := signer.Sign(context.Background(), "https://grok.com", "https://signer.example/sign", "token-a", nil, http.MethodPost, "https://grok.com/rest/other"); err != nil { + t.Fatal(err) + } + if fetches != 4 { + t.Fatalf("different path reused signature: fetches=%d", fetches) + } +} + +func TestStatsigWarmupFetchesMetaOnceForSharedPaths(t *testing.T) { + var fetches, signatures int + signer := newStatsigSigner() + signer.validateEndpoint = func(context.Context, string) error { return nil } + signer.fetchMeta = func(context.Context, string, string, *infraegress.Lease) (string, error) { + fetches++ + return "shared-meta", nil + } + signer.client = &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + signatures++ + var payload struct { + Environment struct { + MetaContent string `json:"metaContent"` + } `json:"environment"` + } + if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { + t.Fatal(err) + } + if payload.Environment.MetaContent != "shared-meta" { + t.Fatalf("meta = %q", payload.Environment.MetaContent) + } + raw := make([]byte, 70) + raw[0] = byte(signatures) + body, _ := json.Marshal(map[string]string{"x-statsig-id": base64.RawStdEncoding.EncodeToString(raw)}) + return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(string(body))), Header: http.Header{}}, nil + })} + targets := []statsigWarmTarget{ + {method: http.MethodPost, target: "https://grok.com/rest/chat"}, + {method: http.MethodPost, target: "https://grok.com/rest/rate-limits"}, + {method: http.MethodPost, target: "https://grok.com/rest/media/post/create"}, + } + warmed, err := signer.Warm(context.Background(), "https://grok.com", "https://signer.example/sign", "token", nil, targets) + if err != nil { + t.Fatal(err) + } + if warmed != len(targets) || fetches != 1 || signatures != len(targets) { + t.Fatalf("warmed=%d fetches=%d signatures=%d", warmed, fetches, signatures) + } + if warmedAgain, err := signer.Warm(context.Background(), "https://grok.com", "https://signer.example/sign", "token", nil, targets); err != nil || warmedAgain != 0 || fetches != 1 { + t.Fatalf("cached warmup=%d fetches=%d err=%v", warmedAgain, fetches, err) + } +} + +func TestApplySignedStatsigUsesManualValue(t *testing.T) { + value := base64.RawStdEncoding.EncodeToString(make([]byte, 70)) + adapter := &Adapter{cfg: Config{BaseURL: "https://grok.com", StatsigMode: "manual", StatsigManualValue: value}} + request, err := http.NewRequest(http.MethodPost, "https://grok.com/rest/test", nil) + if err != nil { + t.Fatal(err) + } + adapter.applySignedStatsig(context.Background(), request, "token", nil) + if request.Header.Get("x-statsig-id") != value { + t.Fatalf("x-statsig-id = %q", request.Header.Get("x-statsig-id")) + } +} + +func TestStatsigInvalidationDoesNotReuseRejectedValue(t *testing.T) { + now := time.Date(2026, 7, 12, 12, 0, 0, 0, time.UTC) + raw := make([]byte, 70) + raw[0] = 1 + previous := base64.RawStdEncoding.EncodeToString(raw) + signer := newStatsigSigner() + signer.now = func() time.Time { return now } + key, _, err := statsigSignatureKey("https://grok.com", "https://signer.example/sign", http.MethodPost, "https://grok.com/rest/test") + if err != nil { + t.Fatal(err) + } + signer.store(key, previous, now.Add(time.Hour), now) + signer.fetchMeta = func(context.Context, string, string, *infraegress.Lease) (string, error) { + return "", errors.New("signer unavailable") + } + signer.Invalidate("https://grok.com", "https://signer.example/sign", http.MethodPost, "https://grok.com/rest/test") + value, source, err := signer.Sign(context.Background(), "https://grok.com", "https://signer.example/sign", "token", nil, http.MethodPost, "https://grok.com/rest/test") + if err == nil || value != "" || source != "" { + t.Fatalf("value=%q source=%q err=%v", value, source, err) + } +} + +func TestApplySignedStatsigNeverLeavesRandomFallback(t *testing.T) { + adapter := &Adapter{cfg: Config{BaseURL: "https://grok.com", StatsigMode: "manual", StatsigManualValue: "invalid"}} + request, err := http.NewRequest(http.MethodPost, "https://grok.com/rest/test", nil) + if err != nil { + t.Fatal(err) + } + request.Header.Set("x-statsig-id", "random-fallback") + adapter.applySignedStatsig(context.Background(), request, "token", nil) + if value := request.Header.Get("x-statsig-id"); value != "" { + t.Fatalf("x-statsig-id = %q", value) + } +} + +func TestStatsigInvalidationOnlyAppliesToURLMode(t *testing.T) { + manual := &Adapter{cfg: Config{StatsigMode: "manual"}, statsig: newStatsigSigner()} + if manual.invalidateSignedStatsig(http.MethodPost, "https://grok.com/rest/test") { + t.Fatal("manual Statsig must not be invalidated automatically") + } + urlMode := &Adapter{cfg: Config{BaseURL: "https://grok.com", StatsigMode: "url", StatsigSignerURL: "https://signer.example/sign"}, statsig: newStatsigSigner()} + if !urlMode.invalidateSignedStatsig(http.MethodPost, "https://grok.com/rest/test") { + t.Fatal("URL Statsig must be invalidated after anti-bot rejection") + } +} diff --git a/backend/internal/infra/provider/web/tools.go b/backend/internal/infra/provider/web/tools.go new file mode 100644 index 0000000000000000000000000000000000000000..bc5177c39281864e833378462ef29237181046b2 --- /dev/null +++ b/backend/internal/infra/provider/web/tools.go @@ -0,0 +1,459 @@ +package web + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "html" + "regexp" + "strings" +) + +const ( + maxFunctionTools = 128 + maxToolDescriptionSize = 16 << 10 +) + +var ( + toolNamePattern = regexp.MustCompile(`^[A-Za-z0-9_-]{1,64}$`) + toolSyntaxPattern = regexp.MustCompile(`(?i){t("auth.sessionUnavailableDescription")}
+ +