From ed51efd8bcd2fd11158bbd3a4860d11f5f06228b Mon Sep 17 00:00:00 2001 From: nyoung Date: Thu, 8 Oct 2026 15:13:23 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E7=BD=91=E5=85=B3=E4=B8=8E=E6=8E=A7?= =?UTF-8?q?=E5=88=B6=E5=8F=B0=E4=BB=A3=E7=A0=81=E5=AE=A1=E6=9F=A5=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=E6=89=B9=E6=AC=A1=EF=BC=88=E5=90=AB=E5=9B=9E=E5=BD=92?= =?UTF-8?q?=E6=B5=8B=E8=AF=95=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Chat 接口拒绝 n>1（此前静默丢弃多余 choice），返回 400 结构化错误 - Responses 接口透传 text.format（json_schema / json_object），三种接口协议对齐 - TLS 指纹统一升级到 Chrome_150（internal/httpx.DefaultProfile），与 UA 匹配 - metrics 落盘失败不再静默：Save() 返回错误并打日志，临时文件失败会清理 - GET /api/config 脱敏 admin.password 与 turnstile.api_key（控制台从不展示， 不应随 5s 轮询下发）；PUT 改走 ReplaceJSON，先校验后替换、失败回滚 - set_limits / set_turnstile 改指针部分更新：缺字段保持现状，不会把并发、 限流或求解开关按零值清掉 - 读请求体统一 64MB 上限；配置写盘原子替换 + 0600 权限 - main.go 改为 run() 返回退出码：日志落盘、指标落盘、xray 子进程回收全部 挂 defer，任何退出路径（含启动失败）都会执行，杜绝孤儿 xray - proxypool / turnstile / upstream / openai / config / httpx / logbuf 等 逐项补红绿回归测试，新增 14 个测试文件，race 下全绿 --- cmd/dsfree2api/main.go | 26 ++- cmd/dsfree2api/main_test.go | 64 +++++++ internal/admin/actions.go | 33 ++-- internal/admin/actions_test.go | 102 +++++++++++ internal/admin/server.go | 40 ++-- internal/api/api_test.go | 49 +++++ internal/api/chat.go | 15 +- internal/api/messages.go | 70 +++++-- internal/api/metrics_test.go | 99 ++++++++++ internal/api/middleware_test.go | 61 ++++++ internal/api/regression_test.go | 94 ++++++++++ internal/api/responses.go | 87 +++++++-- internal/api/responses_format_test.go | 55 ++++++ internal/api/server.go | 63 ++++++- internal/config/config.go | 118 +++++++++++- internal/config/regression_test.go | 190 +++++++++++++++++++ internal/httpx/client.go | 46 ++++- internal/httpx/regression_test.go | 123 +++++++++++++ internal/logbuf/logbuf.go | 12 +- internal/logbuf/logbuf_test.go | 43 +++++ internal/metrics/metrics.go | 36 +++- internal/metrics/metrics_test.go | 37 +++- internal/openai/prompt.go | 29 ++- internal/openai/regression_test.go | 81 ++++++++ internal/openai/toolcalls.go | 8 +- internal/openai/types.go | 42 ++++- internal/proxypool/manager.go | 65 +++++-- internal/proxypool/regression_test.go | 115 ++++++++++++ internal/turnstile/regression_test.go | 99 ++++++++++ internal/turnstile/solver.go | 43 ++++- internal/upstream/client.go | 100 +++++++--- internal/upstream/regression_test.go | 255 ++++++++++++++++++++++++++ internal/upstream/sse.go | 17 +- 33 files changed, 2159 insertions(+), 158 deletions(-) create mode 100644 cmd/dsfree2api/main_test.go create mode 100644 internal/admin/actions_test.go create mode 100644 internal/api/metrics_test.go create mode 100644 internal/api/middleware_test.go create mode 100644 internal/api/regression_test.go create mode 100644 internal/api/responses_format_test.go create mode 100644 internal/config/regression_test.go create mode 100644 internal/httpx/regression_test.go create mode 100644 internal/logbuf/logbuf_test.go create mode 100644 internal/openai/regression_test.go create mode 100644 internal/proxypool/regression_test.go create mode 100644 internal/turnstile/regression_test.go create mode 100644 internal/upstream/regression_test.go diff --git a/cmd/dsfree2api/main.go b/cmd/dsfree2api/main.go index 1a5207a..5ad84b8 100644 --- a/cmd/dsfree2api/main.go +++ b/cmd/dsfree2api/main.go @@ -28,6 +28,13 @@ import ( ) func main() { + // 退出码交给 run 的返回值决定：所有清理都挂在 run 的 defer 上，任何 + // 退出路径（含服务器启动失败）都会执行日志落盘、指标落盘和 xray 子 + // 进程回收。直接 os.Exit 会跳过 defer，留下孤儿 xray。 + os.Exit(run()) +} + +func run() int { var ( configPath = flag.String("config", "config.toml", "path to config.toml") host = flag.String("host", "", "override bind host") @@ -39,13 +46,13 @@ func main() { if *showVer { fmt.Printf("dsfree2api %s\n", admin.Version) - return + return 0 } cfg, err := config.Load(*configPath) if err != nil { fmt.Fprintf(os.Stderr, "config error: %v\n", err) - os.Exit(1) + return 1 } if *host != "" { cfg.Server.Host = *host @@ -55,17 +62,17 @@ func main() { } if err := cfg.Validate(); err != nil { fmt.Fprintf(os.Stderr, "config invalid: %v\n", err) - os.Exit(1) + return 1 } if *checkOnly { fmt.Println("config OK") - return + return 0 } dataDir, err := cfg.DataDir() if err != nil { fmt.Fprintf(os.Stderr, "cannot create data dir: %v\n", err) - os.Exit(1) + return 1 } level := parseLevel(cfg.Server.LogLevel) @@ -85,10 +92,10 @@ func main() { // Proxy pool: per-site egress endpoints (Xray share links, subscriptions, // plain http/socks5) with sticky selection. No-op unless enabled. + // Close 会先停后台循环再同步杀掉 xray 子进程，确保进程退出前回收。 pool := proxypool.New(cfg, dataDir, logger) - poolCtx, poolCancel := context.WithCancel(context.Background()) - defer poolCancel() - pool.Start(poolCtx) + pool.Start(context.Background()) + defer pool.Close() up := upstream.New(cfg, ts, logger, pool) @@ -161,7 +168,7 @@ func main() { select { case err := <-errCh: logger.Error("server failed", "error", err) - os.Exit(1) + return 1 case s := <-sig: logger.Info("shutting down", "signal", s.String()) } @@ -174,6 +181,7 @@ func main() { } met.Close() logger.Info("bye") + return 0 } // logSink opens the daily file log (default /logs/dsfree2api.log) diff --git a/cmd/dsfree2api/main_test.go b/cmd/dsfree2api/main_test.go new file mode 100644 index 0000000..3183c48 --- /dev/null +++ b/cmd/dsfree2api/main_test.go @@ -0,0 +1,64 @@ +package main + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "github.com/nyoungo/dsfree2api/internal/config" +) + +// 回归：HTTP 服务启动失败（errCh 分支）时 run() 必须返回退出码 1，且所有 +// defer 照常执行 —— metrics.json 只有 met.Close() 才会写出来。旧实现 +// 在这里直接 os.Exit(1)，defer 全被跳过：指标/日志不落盘，xray 子进程 +// 也收不回来。 +// +// 触发方式用一个不可绑定的地址（RFC 5737 文档段，任何机器都不可能配在 +// 本地网卡上），而不是占用端口：Windows 上 Go 的 listener 带 SO_REUSEADDR， +// 第二个同端口绑定照样成功，根本走不到失败分支。 +func TestServerBindFailureRunsDefers(t *testing.T) { + bin := filepath.Join(t.TempDir(), "dsfree2api") + if runtime.GOOS == "windows" { + bin += ".exe" + } + if out, err := exec.Command("go", "build", "-o", bin, ".").CombinedOutput(); err != nil { + t.Fatalf("build: %v\n%s", err, out) + } + + cfg, err := config.Load("../../config.example.toml") + if err != nil { + t.Fatalf("load example config: %v", err) + } + dataDir := t.TempDir() + cfg.Runtime.DataDir = dataDir + cfg.Admin.Enabled = false + cfgPath := filepath.Join(t.TempDir(), "config.toml") + if err := cfg.SaveTo(cfgPath); err != nil { + t.Fatalf("save temp config: %v", err) + } + + // CommandContext + 超时：万一子进程不按预期退出，到点强杀，绝不留 + // 孤儿进程（早期版本用裸 exec.Command，遗留过一个跑在随机端口上的 + // 实例直到人工发现）。 + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + out, err := exec.CommandContext(ctx, bin, "-config", cfgPath, "-host", "192.0.2.1").CombinedOutput() + if ctx.Err() == context.DeadlineExceeded { + t.Fatalf("子进程 60s 未退出（已强杀）\noutput:\n%s", out) + } + ee, ok := err.(*exec.ExitError) + if !ok || ee.ExitCode() != 1 { + t.Fatalf("exit = %v, want code 1\noutput:\n%s", err, out) + } + if !strings.Contains(string(out), "server failed") { + t.Fatalf("没走到 server failed 分支\noutput:\n%s", out) + } + if _, err := os.Stat(filepath.Join(dataDir, "metrics.json")); err != nil { + t.Fatalf("metrics.json 没落盘，defer 被跳过: %v\noutput:\n%s", err, out) + } +} diff --git a/internal/admin/actions.go b/internal/admin/actions.go index 21a39ca..89c819b 100644 --- a/internal/admin/actions.go +++ b/internal/admin/actions.go @@ -198,21 +198,27 @@ func (s *Server) handleActions(w http.ResponseWriter, r *http.Request) { s.save(w) case "set_limits": + // 部分更新：只覆盖 payload 里出现的字段，缺字段不能按零值写回， + // 否则一次只改 cross_site 的请求会把并发闸门（0 = 不限）和限流清零。 var p struct { - SlowStart float64 `json:"slow_start"` - MaxConcurrent int `json:"max_concurrent"` - RatePerMinute float64 `json:"rate_per_minute"` - CrossSite *bool `json:"cross_site_failover"` - AutoRefresh *bool `json:"auto_refresh"` - RefreshRetries *int `json:"refresh_retries"` + SlowStart *float64 `json:"slow_start"` + MaxConcurrent *int `json:"max_concurrent"` + RatePerMinute *float64 `json:"rate_per_minute"` + CrossSite *bool `json:"cross_site_failover"` + AutoRefresh *bool `json:"auto_refresh"` + RefreshRetries *int `json:"refresh_retries"` } _ = json.Unmarshal(body, &p) s.cfg.Lock() - if p.SlowStart >= 0 { - s.cfg.Proxy.SlowStartSeconds = p.SlowStart + if p.SlowStart != nil && *p.SlowStart >= 0 { + s.cfg.Proxy.SlowStartSeconds = *p.SlowStart + } + if p.MaxConcurrent != nil { + s.cfg.Limits.MaxConcurrentPerSite = *p.MaxConcurrent + } + if p.RatePerMinute != nil { + s.cfg.Limits.RatePerMinute = *p.RatePerMinute } - s.cfg.Limits.MaxConcurrentPerSite = p.MaxConcurrent - s.cfg.Limits.RatePerMinute = p.RatePerMinute if p.CrossSite != nil { s.cfg.Upstream.CrossSiteFailover = *p.CrossSite } @@ -227,7 +233,7 @@ func (s *Server) handleActions(w http.ResponseWriter, r *http.Request) { case "set_turnstile": var p struct { - Enabled bool `json:"enabled"` + Enabled *bool `json:"enabled"` Provider *string `json:"provider"` APIURL string `json:"api_url"` APIKey string `json:"api_key"` @@ -309,7 +315,10 @@ func (s *Server) handleActions(w http.ResponseWriter, r *http.Request) { if p.WarmCheckSeconds != nil { tc.WarmCheckSeconds = clampInt(*p.WarmCheckSeconds, 5, 3600) } - tc.Enabled = p.Enabled + // 部分更新：payload 不带 enabled 时保持现状，不能把求解器关掉。 + if p.Enabled != nil { + tc.Enabled = *p.Enabled + } if tc.Enabled { // provider-specific requirements mirror config.Validate() switch tc.ProviderValue() { diff --git a/internal/admin/actions_test.go b/internal/admin/actions_test.go new file mode 100644 index 0000000..1f66526 --- /dev/null +++ b/internal/admin/actions_test.go @@ -0,0 +1,102 @@ +package admin + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/nyoungo/dsfree2api/internal/config" + "github.com/nyoungo/dsfree2api/internal/turnstile" +) + +func testConfig(t *testing.T) *config.Config { + t.Helper() + raw, err := os.ReadFile("../../config.example.toml") + if err != nil { + t.Fatalf("read example config: %v", err) + } + dst := filepath.Join(t.TempDir(), "config.toml") + if err := os.WriteFile(dst, raw, 0o644); err != nil { + t.Fatalf("write temp config: %v", err) + } + cfg, err := config.Load(dst) + if err != nil { + t.Fatalf("load config: %v", err) + } + return cfg +} + +func newActionServer(t *testing.T, cfg *config.Config) *Server { + t.Helper() + return New(cfg, nil, nil, nil, turnstile.New(cfg.Turnstile, cfg.Upstream.UserAgent), nil, nil, nil) +} + +func postAction(t *testing.T, s *Server, payload string) *httptest.ResponseRecorder { + t.Helper() + req := httptest.NewRequest(http.MethodPost, "/api/actions", strings.NewReader(payload)) + w := httptest.NewRecorder() + s.handleActions(w, req) + return w +} + +// 回归：set_limits 是部分更新 —— payload 里没带的字段必须保持原值。旧实现把 +// max_concurrent/rate_per_minute 当零值字段，缺字段就把它们写成 0： +// 并发闸门（max<=0 = 不限）和按 Key 限流（limit<=0 = 不限）会被静默关掉， +// slow_start 同样被清零。 +func TestSetLimitsPartialPayloadKeepsExistingValues(t *testing.T) { + cfg := testConfig(t) + cfg.Limits.MaxConcurrentPerSite = 4 + cfg.Limits.RatePerMinute = 7 + cfg.Proxy.SlowStartSeconds = 8 + s := newActionServer(t, cfg) + + w := postAction(t, s, `{"action":"set_limits","cross_site_failover":false}`) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } + cfg.RLock() + got := struct { + MaxConcurrent int + Rate float64 + SlowStart float64 + CrossSite bool + }{cfg.Limits.MaxConcurrentPerSite, cfg.Limits.RatePerMinute, cfg.Proxy.SlowStartSeconds, cfg.Upstream.CrossSiteFailover} + cfg.RUnlock() + + if got.MaxConcurrent != 4 { + t.Errorf("max_concurrent = %d, want 4（缺字段被清零会让并发闸门失效）", got.MaxConcurrent) + } + if got.Rate != 7 { + t.Errorf("rate_per_minute = %v, want 7", got.Rate) + } + if got.SlowStart != 8 { + t.Errorf("slow_start = %v, want 8", got.SlowStart) + } + if got.CrossSite { + t.Errorf("cross_site_failover = true, want false（payload 里的值要生效）") + } +} + +// 回归：set_turnstile 同样是部分更新 —— 不带 enabled 不能把求解器关掉。 +func TestSetTurnstilePartialPayloadKeepsEnabled(t *testing.T) { + cfg := testConfig(t) + cfg.Turnstile.Enabled = true + cfg.Turnstile.Provider = config.ProviderAPI + cfg.Turnstile.APIURL = "https://solve.example.invalid/" + cfg.Turnstile.Retries = 5 + s := newActionServer(t, cfg) + + w := postAction(t, s, `{"action":"set_turnstile","retries":3}`) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } + if !cfg.Turnstile.Enabled { + t.Errorf("turnstile.enabled 被部分 payload 关掉了") + } + if cfg.Turnstile.Retries != 3 { + t.Errorf("retries = %d, want 3", cfg.Turnstile.Retries) + } +} diff --git a/internal/admin/server.go b/internal/admin/server.go index c78a46d..680fc84 100644 --- a/internal/admin/server.go +++ b/internal/admin/server.go @@ -202,6 +202,9 @@ func (s *Server) handleOverview(w http.ResponseWriter, _ *http.Request) { }{s.cfg.Proxy.URL, s.cfg.Proxy.FallbackURLs, s.cfg.Proxy.SlowStartSeconds, s.cfg.Limits.MaxConcurrentPerSite, s.cfg.Limits.RatePerMinute, s.cfg.Upstream.CrossSiteFailover} turnstileCfg := s.cfg.Turnstile + // The console never renders the solver key (its field is "type to + // replace"), so it must not ship every 5s with the dashboard poll either. + turnstileCfg.APIKey = "" keys := append([]string(nil), s.cfg.Security.APIKeys...) s.cfg.RUnlock() @@ -249,8 +252,27 @@ func routeList(c *upstream.Client) []map[string]string { func (s *Server) handleGetConfig(w http.ResponseWriter, _ *http.Request) { s.cfg.RLock() - defer s.cfg.RUnlock() - writeJSON(w, http.StatusOK, s.cfg) + raw, err := json.Marshal(s.cfg) + s.cfg.RUnlock() + if err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]any{"error": err.Error()}) + return + } + var out map[string]any + if err := json.Unmarshal(raw, &out); err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]any{"error": err.Error()}) + return + } + // The console never displays these two secrets (both inputs are "leave + // blank to keep"), so there is no reason to hand them out; an empty value + // round-tripping through PUT is interpreted as "unchanged". + if admin, ok := out["admin"].(map[string]any); ok { + admin["password"] = "" + } + if ts, ok := out["turnstile"].(map[string]any); ok { + ts["api_key"] = "" + } + writeJSON(w, http.StatusOK, out) } func (s *Server) handlePutConfig(w http.ResponseWriter, r *http.Request) { @@ -259,17 +281,9 @@ func (s *Server) handlePutConfig(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusBadRequest, map[string]any{"error": err.Error()}) return } - s.cfg.Lock() - err = json.Unmarshal(body, s.cfg) - if err == nil { - err = s.cfg.Reindex() - } - s.cfg.Unlock() - if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]any{"error": err.Error()}) - return - } - if err := s.cfg.Validate(); err != nil { + // ReplaceJSON validates first and rolls back on rejection, so a bad + // payload can never leave half-applied state in the running config. + if err := s.cfg.ReplaceJSON(body); err != nil { writeJSON(w, http.StatusBadRequest, map[string]any{"error": err.Error()}) return } diff --git a/internal/api/api_test.go b/internal/api/api_test.go index 6561029..7c30dd1 100644 --- a/internal/api/api_test.go +++ b/internal/api/api_test.go @@ -186,6 +186,29 @@ func anthropicErrType(body string) string { return "" } +// openaiErrType 解析 OpenAI 风格错误体 {"error":{"type":...}}。 +func openaiErrType(body string) string { + var out struct { + Error struct { + Type string `json:"type"` + } `json:"error"` + } + _ = json.Unmarshal([]byte(body), &out) + return out.Error.Type +} + +// openaiErrMsg 取出 error.message（JSON 会把 ">" 转义成 >， +// 所以断言文案必须走反序列化，不能拿原文子串匹配）。 +func openaiErrMsg(body string) string { + var out struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } + _ = json.Unmarshal([]byte(body), &out) + return out.Error.Message +} + func TestMessagesRequiresAuthInAnthropicFormat(t *testing.T) { srv := newTestServer(t, nil) resp, body := post(t, srv.URL+"/v1/messages", nil, `{}`) @@ -234,3 +257,29 @@ func TestResponsesValidation(t *testing.T) { t.Fatalf("unknown model status = %d body=%s, want 404", resp.StatusCode, body) } } + +// 回归：n>1 必须 400 —— 网关只产 1 条 choice，静默忽略会让客户端以为 +// 拿到了 n 条候选。校验排在模型查找之前（模型不存在也先报参数错）。 +func TestChatRejectsNGreaterThanOne(t *testing.T) { + srv := newTestServer(t, nil) + h := map[string]string{"Authorization": "Bearer " + testKey, "Content-Type": "application/json"} + + resp, body := post(t, srv.URL+"/v1/chat/completions", h, + `{"model":"no-such-model","messages":[{"role":"user","content":"hi"}],"n":2}`) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("n=2 status = %d body=%s, want 400", resp.StatusCode, body) + } + if !strings.Contains(openaiErrMsg(body), "n > 1") { + t.Errorf("error message should mention n > 1, body=%s", body) + } + if got := openaiErrType(body); got != "invalid_request_error" { + t.Errorf("error type = %q, want invalid_request_error; body=%s", got, body) + } + + // n=1（含省略）照常放行到模型查找，返回 404 + resp, body = post(t, srv.URL+"/v1/chat/completions", h, + `{"model":"no-such-model","messages":[{"role":"user","content":"hi"}],"n":1}`) + if resp.StatusCode != http.StatusNotFound { + t.Fatalf("n=1 status = %d body=%s, want 404", resp.StatusCode, body) + } +} diff --git a/internal/api/chat.go b/internal/api/chat.go index 7fa3a2b..f3575a0 100644 --- a/internal/api/chat.go +++ b/internal/api/chat.go @@ -30,6 +30,14 @@ func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusBadRequest, "messages is required", "invalid_request_error") return } + // 网关一次只回 1 条 choice；n>1 若静默忽略，客户端会以为拿到了 + // n 条候选 —— 明确拒绝。校验要排在模型查找之前。 + if req.N != nil && *req.N > 1 { + writeError(w, http.StatusBadRequest, + "n > 1 is not supported: this gateway returns a single choice", + "invalid_request_error") + return + } s.cfg.RLock() m, ok := s.cfg.Models[req.Model] site := "" @@ -148,7 +156,8 @@ func (s *Server) completeChat(ctx context.Context, w http.ResponseWriter, req op } func buildCompletion(req openai.ChatCompletionRequest, prompt, body string) *openai.ChatCompletionResponse { - message := openai.ChoiceMessage{Content: body} + // Role 的 json tag 没有 omitempty，漏设会返回空串；与流式路径保持一致。 + message := openai.ChoiceMessage{Role: "assistant", Content: body} if len(req.Tools) > 0 { remaining, calls, ok := openai.TryParseToolCalls(body) if ok { @@ -233,8 +242,8 @@ func (s *Server) streamChat(ctx context.Context, w http.ResponseWriter, req open _ = write(chunk(openai.ChoiceDelta{Role: "assistant"}, nil)) - s.met.IncInFlight(1) - defer s.met.IncInFlight(-1) + // in_flight 由 handleChat 统一计一次，streamChat 不再重复计数， + // 否则控制台的并发数会按 2 倍显示。 err := s.chatWithContinue(ctx, req.Model, prompt, cont, info, func() { if firstDelta.IsZero() { diff --git a/internal/api/messages.go b/internal/api/messages.go index bacfe7f..a22bf2c 100644 --- a/internal/api/messages.go +++ b/internal/api/messages.go @@ -16,8 +16,7 @@ import ( func (s *Server) handleMessages(w http.ResponseWriter, r *http.Request) { var req anthropic.Request - if err := json.NewDecoder(r.Body).Decode(&req); err != nil { - writeAnthropicError(w, http.StatusBadRequest, "invalid JSON body: "+err.Error(), "invalid_request_error") + if !decodeBody(w, r, &req, writeAnthropicError) { return } if req.Model == "" { @@ -128,17 +127,6 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req } msgID := anthropic.NewID("msg_") - _ = writeEvent("message_start", map[string]any{ - "type": "message_start", - "message": map[string]any{ - "id": msgID, "type": "message", "role": "assistant", "model": req.Model, - "content": []any{}, "stop_reason": nil, "stop_sequence": nil, - "usage": map[string]any{ - "input_tokens": estimateUsage(prompt, "").PromptTokens, - "output_tokens": 0, - }, - }, - }) // With tools the raw tool-call JSON is converted into streamed tool_use // blocks: content_block_start once id+name are known, then one @@ -147,11 +135,36 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req if len(tools) > 0 { args = openai.NewArgStreamer() } - if args == nil { - _ = writeEvent("content_block_start", map[string]any{ - "type": "content_block_start", "index": 0, - "content_block": map[string]any{"type": "text", "text": ""}, - }) + + // The opening events are held back until the upstream produces output: + // that keeps the HTTP status free for a failure with nothing streamed yet, + // exactly like the non-streaming path. + var started bool + beginStream := func() error { + if started { + return nil + } + started = true + if err := writeEvent("message_start", map[string]any{ + "type": "message_start", + "message": map[string]any{ + "id": msgID, "type": "message", "role": "assistant", "model": req.Model, + "content": []any{}, "stop_reason": nil, "stop_sequence": nil, + "usage": map[string]any{ + "input_tokens": estimateUsage(prompt, "").PromptTokens, + "output_tokens": 0, + }, + }, + }); err != nil { + return err + } + if args == nil { + return writeEvent("content_block_start", map[string]any{ + "type": "content_block_start", "index": 0, + "content_block": map[string]any{"type": "text", "text": ""}, + }) + } + return nil } // Anthropic content blocks are strictly sequential: close the open @@ -170,6 +183,9 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req if d.Index < curIdx { return nil } + if err := beginStream(); err != nil { + return err + } if d.ID != "" { ids[d.Index] = d.ID } @@ -234,6 +250,9 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req firstDelta = time.Now() } }, func(seg string) error { + if err := beginStream(); err != nil { + return err + } text.WriteString(seg) if args == nil { return writeEvent("content_block_delta", map[string]any{ @@ -264,6 +283,12 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req Status: metrics.StatusError, Error: err.Error(), }) s.log.Warn("messages stream error", "model", req.Model, "error", err) + if !started { + // nothing on the wire yet — report a real status like the + // non-streaming path does + writeAnthropicUpstreamError(w, err) + return + } writeErr("api_error", err.Error()) return } @@ -272,10 +297,19 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req Model: req.Model, ServedBy: servedBy, Site: servedSite, Stream: true, DurationMs: duration.Milliseconds(), Status: metrics.StatusEmpty, Error: "empty response", }) + if !started { + writeAnthropicError(w, http.StatusInternalServerError, emptyResponseMessage, "api_error") + return + } writeErr("api_error", emptyResponseMessage) return } + // Success: the opening events must precede every event below. + if err := beginStream(); err != nil { + return + } + usage := estimateUsage(prompt, body) outText, calls := splitToolOutput(body, tools) diff --git a/internal/api/metrics_test.go b/internal/api/metrics_test.go new file mode 100644 index 0000000..3f6f813 --- /dev/null +++ b/internal/api/metrics_test.go @@ -0,0 +1,99 @@ +package api + +import ( + "io" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/nyoungo/dsfree2api/internal/config" + "github.com/nyoungo/dsfree2api/internal/logbuf" + "github.com/nyoungo/dsfree2api/internal/metrics" + "github.com/nyoungo/dsfree2api/internal/turnstile" + "github.com/nyoungo/dsfree2api/internal/upstream" +) + +func newTestServerWithConfig(t *testing.T, cfg *config.Config) *Server { + t.Helper() + met := metrics.New(t.TempDir()) + met.Start() + t.Cleanup(met.Close) + log := slog.New(slog.NewTextHandler(io.Discard, nil)) + ts := turnstile.New(cfg.Turnstile, cfg.Upstream.UserAgent) + up := upstream.New(cfg, ts, log, nil) + return New(cfg, up, met, logbuf.New(10), ts, log) +} + +// 回归：一次流式请求只能计一次 in_flight。handleChat 已经计过，streamChat +// 里又计一次，控制台的并发数会按 2 倍显示。 +func TestStreamingRequestCountsInFlightOnce(t *testing.T) { + // 假站点：每个请求拖 300ms，好让测试在请求进行中采样并发数 + hang := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(300 * time.Millisecond) + w.Header().Set("Content-Type", "text/html") + _, _ = w.Write([]byte("nope")) + })) + t.Cleanup(hang.Close) + + cfg, err := config.Load(examplePath) + if err != nil { + t.Fatalf("load config: %v", err) + } + cfg.Turnstile.Enabled = false + cfg.Upstream.AutoRefresh = false + cfg.Upstream.RefreshRetries = 0 + cfg.Upstream.CrossSiteFailover = false + // 所有站点都指向假站点，避免 failover 打到公网 + for _, st := range cfg.Sites { + st.BaseURL = hang.URL + st.AJAXURL = hang.URL + "/wp-admin/admin-ajax.php" + } + var modelID string + for id, m := range cfg.Models { + if m.Enabled { + modelID = id + break + } + } + if modelID == "" { + t.Fatalf("example config 里没有可用模型") + } + + s := newTestServerWithConfig(t, cfg) + + body := `{"model":"` + modelID + `","messages":[{"role":"user","content":"hi"}],"stream":true}` + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body)) + w := httptest.NewRecorder() + + done := make(chan struct{}) + go func() { + defer close(done) + s.handleChat(w, req) + }() + + var max int64 + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() +loop: + for { + if v := s.met.InFlight(); v > max { + max = v + } + select { + case <-done: + break loop + case <-ticker.C: + } + } + <-done + + if max > 1 { + t.Fatalf("一次流式请求的 in_flight 峰值 = %d, want 1", max) + } + if max == 0 { + t.Fatalf("请求期间 in_flight 一直是 0，测试没抓到窗口") + } +} diff --git a/internal/api/middleware_test.go b/internal/api/middleware_test.go new file mode 100644 index 0000000..aa420f7 --- /dev/null +++ b/internal/api/middleware_test.go @@ -0,0 +1,61 @@ +package api + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/nyoungo/dsfree2api/internal/config" +) + +func newMiddlewareServer(t *testing.T) *Server { + t.Helper() + cfg, err := config.Load(examplePath) + if err != nil { + t.Fatalf("load config: %v", err) + } + return newTestServerWithConfig(t, cfg) +} + +// 回归：SSE 已经开流（200 + 若干字节）之后再 panic，不能又写一个 500 的 +// JSON 进响应体 —— 状态码改不了，只会把 {"error"...} 拼在事件流后面， +// 客户端解析流直接失败。 +func TestRecoverPanicAfterHeadersKeepsBodyIntact(t *testing.T) { + s := newMiddlewareServer(t) + h := s.recoverPanic(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream; charset=utf-8") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("data: {\"a\":1}\n\n")) + panic("boom") + })) + + w := httptest.NewRecorder() + h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil)) + + body := w.Body.String() + if strings.Contains(body, "internal server error") { + t.Fatalf("已开流的响应被追加了 500 JSON: %q", body) + } + if !strings.Contains(body, `data: {"a":1}`) { + t.Fatalf("原有事件流丢了: %q", body) + } +} + +// 回归：还没发出响应头就 panic，必须照常返回 500 + 错误对象。 +func TestRecoverPanicBeforeHeadersWrites500(t *testing.T) { + s := newMiddlewareServer(t) + h := s.recoverPanic(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + panic("boom") + })) + + w := httptest.NewRecorder() + h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil)) + + if w.Code != http.StatusInternalServerError { + t.Fatalf("status = %d, want 500", w.Code) + } + if !strings.Contains(w.Body.String(), "internal server error") { + t.Fatalf("body = %q, want the error object", w.Body.String()) + } +} diff --git a/internal/api/regression_test.go b/internal/api/regression_test.go new file mode 100644 index 0000000..c25715c --- /dev/null +++ b/internal/api/regression_test.go @@ -0,0 +1,94 @@ +package api + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/nyoungo/dsfree2api/internal/openai" +) + +// 回归：非流式 /v1/chat/completions 的 choices[0].message.role 必须是 +// "assistant"。Role 的 json tag 没有 omitempty（types.go:163），漏设就会 +// 返回空串；流式路径与非流式路径必须一致。 +func TestBuildCompletionRoleIsAssistant(t *testing.T) { + resp := buildCompletion(openai.ChatCompletionRequest{Model: "m"}, "p", "hello") + if len(resp.Choices) != 1 { + t.Fatalf("choices = %d", len(resp.Choices)) + } + got := resp.Choices[0].Message.Role + if got != "assistant" { + raw, _ := json.Marshal(resp.Choices[0].Message) + t.Fatalf("choices[0].message.role = %q (want \"assistant\"), raw=%s", got, raw) + } +} + +// 回归：/v1/responses 的 usage 必须用 Responses API 的字段名 +// （input_tokens/output_tokens），而不是 Chat Completions 的 +// prompt_tokens/completion_tokens —— 后者会让客户端读到 None。 +// 官方文档："Chat Completions reports prompt_tokens, completion_tokens, and +// total_tokens. Responses reports input_tokens, output_tokens, and total_tokens." +func TestResponsesUsageFieldNames(t *testing.T) { + m := responsesUsage(openai.Usage{PromptTokens: 1, CompletionTokens: 2, TotalTokens: 3}) + raw, err := json.Marshal(m) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var got map[string]any + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatalf("unmarshal: %v", err) + } + for _, want := range []string{"input_tokens", "output_tokens", "total_tokens"} { + if _, ok := got[want]; !ok { + t.Fatalf("Responses API usage 缺少 %s，实际字段: %v", want, keys(got)) + } + } + for _, bad := range []string{"prompt_tokens", "completion_tokens"} { + if _, ok := got[bad]; ok { + t.Fatalf("Responses API usage 不应包含 Chat 字段 %s，实际字段: %v", bad, keys(got)) + } + } +} + +// 回归：请求体必须有上限（api_keys 为空时这些端点未鉴权），超限要返回 413 +// 而不是把整个 body 读进内存。 +func TestDecodeBodyRejectsOversizedPayload(t *testing.T) { + payload := []byte(`{"model":"m","messages":[{"role":"user","content":"` + + strings.Repeat("a", maxRequestBodyBytes+4096) + `"}]}`) + r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(payload)) + w := httptest.NewRecorder() + + var req openai.ChatCompletionRequest + if decodeBody(w, r, &req, writeError) { + t.Fatalf("超过上限的请求体被接受了") + } + if w.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("status = %d, want %d", w.Code, http.StatusRequestEntityTooLarge) + } +} + +// 回归：未超限的正常请求体仍然可以解码。 +func TestDecodeBodyAcceptsNormalPayload(t *testing.T) { + r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"m","messages":[{"role":"user","content":"hi"}]}`)) + w := httptest.NewRecorder() + + var req openai.ChatCompletionRequest + if !decodeBody(w, r, &req, writeError) { + t.Fatalf("正常请求体被拒绝: %s", w.Body.String()) + } + if req.Model != "m" { + t.Fatalf("model = %q", req.Model) + } +} + +func keys(m map[string]any) []string { + out := make([]string, 0, len(m)) + for k := range m { + out = append(out, k) + } + return out +} diff --git a/internal/api/responses.go b/internal/api/responses.go index 8591b0f..d57bf8e 100644 --- a/internal/api/responses.go +++ b/internal/api/responses.go @@ -13,10 +13,38 @@ import ( "github.com/nyoungo/dsfree2api/internal/upstream" ) +// responsesUsage renders token usage with the Responses API field names +// (input_tokens/output_tokens/total_tokens). Reusing the Chat Completions +// shape here would leave every client's response.usage.input_tokens as null. +func responsesUsage(u openai.Usage) map[string]any { + return map[string]any{ + "input_tokens": u.PromptTokens, + "output_tokens": u.CompletionTokens, + "total_tokens": u.TotalTokens, + } +} + +// responsesPromptOptions 把 Responses 请求翻译成 BuildPrompt 选项。 +// text.format 里的输出格式约束必须带过去 —— 与 Chat Completions 的 +// response_format 等价，漏掉它客户端的 json_schema 就是静默失效。 +func responsesPromptOptions(req *openai.ResponsesRequest, tools []openai.ToolDef) openai.PromptOptions { + var rf *openai.ResponseFormat + if req.Text != nil { + rf = req.Text.Format + } + return openai.PromptOptions{ + Tools: tools, + ToolChoice: req.ToolChoice, + Temperature: req.Temperature, + TopP: req.TopP, + MaxTokens: req.MaxOutputTokens, + ResponseFormat: rf, + } +} + func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) { var req openai.ResponsesRequest - if err := json.NewDecoder(r.Body).Decode(&req); err != nil { - writeError(w, http.StatusBadRequest, "invalid JSON body: "+err.Error(), "invalid_request_error") + if !decodeBody(w, r, &req, writeError) { return } messages, err := openai.ResponsesToMessages(&req) @@ -43,13 +71,7 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) { return } - prompt := openai.BuildPrompt(messages, openai.PromptOptions{ - Tools: tools, - ToolChoice: req.ToolChoice, - Temperature: req.Temperature, - TopP: req.TopP, - MaxTokens: req.MaxOutputTokens, - }) + prompt := openai.BuildPrompt(messages, responsesPromptOptions(&req, tools)) s.met.IncInFlight(1) defer s.met.IncInFlight(-1) @@ -86,7 +108,7 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) { "output": output, "output_text": text, "stop_reason": stopReason, - "usage": out.usage, + "usage": responsesUsage(out.usage), }) } @@ -142,12 +164,22 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req return nil } - _ = write(map[string]any{ - "type": "response.created", - "response": map[string]any{ - "id": id, "model": req.Model, "status": "in_progress", "created_at": time.Now().Unix(), - }, - }) + // The created event is held back until the upstream produces output, so a + // failure with nothing streamed yet can still use a real HTTP status + // instead of a failed event inside an already-200 stream. + var started bool + beginStream := func() error { + if started { + return nil + } + started = true + return write(map[string]any{ + "type": "response.created", + "response": map[string]any{ + "id": id, "model": req.Model, "status": "in_progress", "created_at": time.Now().Unix(), + }, + }) + } // With tools the raw tool-call JSON is converted into a live // function_call item stream: output_item.added once id+name are known, @@ -200,6 +232,9 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req if d.Index < openIdx { return nil // models emit calls in order — drop stale index } + if err := beginStream(); err != nil { + return err + } if d.ID != "" { callIDs[d.Index] = d.ID } @@ -259,6 +294,9 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req firstDelta = time.Now() } }, func(seg string) error { + if err := beginStream(); err != nil { + return err + } text.WriteString(seg) if args == nil { return write(map[string]any{"type": "response.output_text.delta", "delta": seg}) @@ -285,6 +323,12 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req Status: metrics.StatusError, Error: err.Error(), }) // Never report a failed stream as completed. + if !started { + // nothing on the wire yet — report a real status like the + // non-streaming path does + writeUpstreamError(w, err) + return + } _ = write(map[string]any{ "type": "response.failed", "response": map[string]any{ @@ -303,6 +347,10 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req Model: req.Model, ServedBy: servedBy, Site: servedSite, Stream: true, DurationMs: duration.Milliseconds(), Status: metrics.StatusEmpty, Error: "empty response", }) + if !started { + writeError(w, http.StatusBadGateway, emptyResponseMessage, "upstream_empty_response") + return + } _ = write(map[string]any{ "type": "response.failed", "response": map[string]any{ @@ -317,6 +365,11 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req return } + // Success: response.created must precede every event below. + if err := beginStream(); err != nil { + return + } + usage := estimateUsage(prompt, body) outText, calls := splitToolOutput(body, tools) @@ -420,7 +473,7 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req "type": "response.completed", "response": map[string]any{ "id": id, "model": req.Model, "status": "completed", "created_at": time.Now().Unix(), - "output": output, "stop_reason": stopReason, "usage": usage, + "output": output, "stop_reason": stopReason, "usage": responsesUsage(usage), }, }) _, _ = fmt.Fprint(w, "data: [DONE]\n\n") diff --git a/internal/api/responses_format_test.go b/internal/api/responses_format_test.go new file mode 100644 index 0000000..9b9be66 --- /dev/null +++ b/internal/api/responses_format_test.go @@ -0,0 +1,55 @@ +package api + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/nyoungo/dsfree2api/internal/openai" +) + +// 回归：Responses API 的 text.format 必须转成提示词里的输出格式约束。 +// 之前 ResponsesRequest 根本不解析 text 字段，客户端指定的 json_schema +// 被静默丢弃 —— 模型可以随便输出非 JSON，客户端的 schema 校验必炸。 +func TestResponsesTextFormatBecomesPromptConstraint(t *testing.T) { + body := `{"model":"m","input":"hi",` + + `"text":{"format":{"type":"json_schema",` + + `"json_schema":{"name":"addr","schema":{"type":"object",` + + `"properties":{"city":{"type":"string"}},"required":["city"]}}}}}` + var req openai.ResponsesRequest + if err := json.Unmarshal([]byte(body), &req); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if req.Text == nil || req.Text.Format == nil { + t.Fatalf("text.format 没解析出来: %+v", req.Text) + } + + messages, err := openai.ResponsesToMessages(&req) + if err != nil { + t.Fatalf("responses to messages: %v", err) + } + prompt := openai.BuildPrompt(messages, responsesPromptOptions(&req, nil)) + + if !strings.Contains(prompt, "JSON Schema") { + t.Fatalf("提示词里没有 schema 约束: %q", prompt) + } + if !strings.Contains(prompt, "city") { + t.Fatalf("提示词里没有 schema 内容: %q", prompt) + } +} + +// 回归：没有 text.format 的请求照常工作，不能凭空多出约束。 +func TestResponsesWithoutTextFormatHasNoConstraint(t *testing.T) { + var req openai.ResponsesRequest + if err := json.Unmarshal([]byte(`{"model":"m","input":"hi"}`), &req); err != nil { + t.Fatalf("unmarshal: %v", err) + } + messages, err := openai.ResponsesToMessages(&req) + if err != nil { + t.Fatalf("responses to messages: %v", err) + } + prompt := openai.BuildPrompt(messages, responsesPromptOptions(&req, nil)) + if strings.Contains(prompt, "JSON Schema") { + t.Fatalf("没有 text.format 却出现了 schema 约束: %q", prompt) + } +} diff --git a/internal/api/server.go b/internal/api/server.go index 4e94b41..4d47123 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -3,6 +3,7 @@ package api import ( "crypto/subtle" "encoding/json" + "errors" "log/slog" "net/http" "runtime/debug" @@ -183,18 +184,47 @@ func (s *Server) logRequest(next http.Handler) http.Handler { func (s *Server) recoverPanic(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + tw := &headerTracker{ResponseWriter: w} defer func() { if rec := recover(); rec != nil { s.log.Error("panic", "err", rec, "stack", string(debug.Stack())) - writeJSON(w, http.StatusInternalServerError, map[string]any{ - "error": map[string]any{"message": "internal server error", "type": "server_error"}, - }) + // 响应头已经发出去（SSE 已开流）就不能再改状态码：硬写 500 + // 只会把错误 JSON 拼进事件流，客户端解析流会直接失败。 + if !tw.wroteHeader { + writeJSON(tw, http.StatusInternalServerError, map[string]any{ + "error": map[string]any{"message": "internal server error", "type": "server_error"}, + }) + } } }() - next.ServeHTTP(w, r) + next.ServeHTTP(tw, r) }) } +// headerTracker 记录响应头是否已发出，并把 Flush 透传给底层 writer。 +type headerTracker struct { + http.ResponseWriter + wroteHeader bool +} + +func (w *headerTracker) WriteHeader(code int) { + w.wroteHeader = true + w.ResponseWriter.WriteHeader(code) +} + +func (w *headerTracker) Write(b []byte) (int, error) { + if !w.wroteHeader { + w.WriteHeader(http.StatusOK) + } + return w.ResponseWriter.Write(b) +} + +func (w *headerTracker) Flush() { + if f, ok := w.ResponseWriter.(http.Flusher); ok { + f.Flush() + } +} + type statusWriter struct { http.ResponseWriter status int @@ -225,6 +255,31 @@ func writeError(w http.ResponseWriter, code int, msg, typ string) { }) } +// maxRequestBodyBytes caps every public endpoint's body: these endpoints are +// unauthenticated when api_keys is empty, so an uncapped body would let one +// client pin a connection and allocate without bound. +const maxRequestBodyBytes = 8 << 20 + +type errorWriter func(w http.ResponseWriter, code int, msg, typ string) + +// decodeBody caps the request body before JSON-decoding it and reports any +// failure through the caller's error shape (OpenAI vs Anthropic). It returns +// false after having written the error response. +func decodeBody(w http.ResponseWriter, r *http.Request, dst any, fail errorWriter) bool { + r.Body = http.MaxBytesReader(w, r.Body, maxRequestBodyBytes) + err := json.NewDecoder(r.Body).Decode(dst) + if err == nil { + return true + } + var mbe *http.MaxBytesError + if errors.As(err, &mbe) { + fail(w, http.StatusRequestEntityTooLarge, "request body too large", "invalid_request_error") + return false + } + fail(w, http.StatusBadRequest, "invalid JSON body: "+err.Error(), "invalid_request_error") + return false +} + func sortModelInfo(data []map[string]any) { for i := 1; i < len(data); i++ { for j := i; j > 0; j-- { diff --git a/internal/config/config.go b/internal/config/config.go index f5fc313..6cb2c6e 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -1,6 +1,7 @@ package config import ( + "encoding/json" "errors" "fmt" "os" @@ -8,6 +9,7 @@ import ( "strconv" "strings" "sync" + "time" "github.com/BurntSushi/toml" ) @@ -537,13 +539,16 @@ func (c *Config) applyEnv() error { } c.Server.Port = n } - if v := str("API_KEYS"); v != nil { + // 只有"变量存在且非空"才覆盖配置文件：compose 的 `API_KEYS: ${API_KEYS:-}` + // 会让变量恒存在，空值必须视为"不覆盖"，否则会把 config.toml 的 key 清空 + // 并静默关闭鉴权。 + if v := str("API_KEYS"); v != nil && strings.TrimSpace(*v) != "" { c.Security.APIKeys = splitList(*v) } - if v := str("PROXY_URL"); v != nil { + if v := str("PROXY_URL"); v != nil && strings.TrimSpace(*v) != "" { c.Proxy.URL = strings.TrimSpace(*v) } - if v := str("PROXY_FALLBACK_URLS"); v != nil { + if v := str("PROXY_FALLBACK_URLS"); v != nil && strings.TrimSpace(*v) != "" { c.Proxy.FallbackURLs = splitList(*v) } if v := str("PROXY_SLOW_START_SECONDS"); v != nil { @@ -633,7 +638,15 @@ func splitList(v string) []string { return out } +// Validate takes the write lock because it also normalizes fields +// (upstream.continue_rounds), so it must not run concurrently with a mutation. func (c *Config) Validate() error { + c.mu.Lock() + defer c.mu.Unlock() + return c.validateLocked() +} + +func (c *Config) validateLocked() error { if c.Server.Port < 1 || c.Server.Port > 65535 { return fmt.Errorf("server.port out of range: %d", c.Server.Port) } @@ -723,14 +736,24 @@ func (c *Config) Save() error { return c.SaveTo(c.path) } +// SaveTo encodes the config with a read lock held: toml.Encode iterates the +// Sites/Models/... maps, which admin actions mutate under the write lock, and +// an unlocked iteration racing a map write is a fatal (unrecoverable) error. func (c *Config) SaveTo(path string) error { + c.mu.RLock() + defer c.mu.RUnlock() + return c.saveToLocked(path) +} + +func (c *Config) saveToLocked(path string) error { if path == "" { return errors.New("config path is empty") } if _, err := os.Stat(path); err == nil { backup := path + ".bak" if raw, err := os.ReadFile(path); err == nil { - _ = os.WriteFile(backup, raw, 0o644) + // .bak 里同样是 api_keys / admin password，收紧到 0600 + _ = os.WriteFile(backup, raw, 0o600) } } if dir := filepath.Dir(path); dir != "" && dir != "." { @@ -743,7 +766,92 @@ func (c *Config) SaveTo(path string) error { return fmt.Errorf("encode config: %w", err) } header := "# dsfree2api configuration\n# Managed by the web console; a .bak copy is kept next to this file.\n\n" - return os.WriteFile(path, []byte(header+sb.String()), 0o644) + data := []byte(header + sb.String()) + + // 原子替换：先写同目录临时文件再 rename。裸 WriteFile 会先把目标截断， + // 写一半时进程崩溃就留下损坏的 config.toml（下次启动直接 config error）， + // 并发读者也会读到空/半截内容。 + tmp := path + ".tmp" + if err := os.WriteFile(tmp, data, 0o600); err != nil { + return fmt.Errorf("write config temp: %w", err) + } + if err := renameReplace(tmp, path); err != nil { + _ = os.Remove(tmp) + return fmt.Errorf("replace config: %w", err) + } + return nil +} + +// renameReplace 把 src 重命名覆盖到 dst。Windows 上目标文件被并发读者以 +// 不带 FILE_SHARE_DELETE 的方式打开时 MoveFileEx 会报 sharing violation， +// 而读者都是瞬间的 ReadFile，退避重试即可等到空档；其他平台首次即成功。 +func renameReplace(src, dst string) error { + var err error + for i := 0; i < 100; i++ { + if err = os.Rename(src, dst); err == nil { + return nil + } + time.Sleep(10 * time.Millisecond) + } + return err +} + +// ReplaceJSON applies a console payload in one shot: decode into the live +// config only after the payload itself is known to be well-formed, and roll +// back to the pre-request snapshot when reindexing or validation rejects it, +// so a failed PUT never leaves half-applied state behind in memory. +func (c *Config) ReplaceJSON(raw []byte) error { + var probe map[string]json.RawMessage + if err := json.Unmarshal(raw, &probe); err != nil { + return err + } + + c.mu.Lock() + defer c.mu.Unlock() + + snapshot, err := json.Marshal(c) + if err != nil { + return err + } + // json 解码对 map 是"合并"语义，待替换的键必须先清空才能真正替换 + if _, ok := probe["sites"]; ok { + c.Sites = nil + } + if _, ok := probe["models"]; ok { + c.Models = nil + } + if _, ok := probe["proxypool"]; ok { + c.ProxyPool.Entries = nil + c.ProxyPool.Subscriptions = nil + } + curPassword, curTSKey := c.Admin.Password, c.Turnstile.APIKey + if err := json.Unmarshal(raw, c); err != nil { + c.restoreLocked(snapshot) + return err + } + // 控制台 GET 会把这两个密钥抹成空串再原样回传：空值 = 保持不变 + if c.Admin.Password == "" { + c.Admin.Password = curPassword + } + if strings.TrimSpace(c.Turnstile.APIKey) == "" { + c.Turnstile.APIKey = curTSKey + } + if err := c.Reindex(); err != nil { + c.restoreLocked(snapshot) + return err + } + if err := c.validateLocked(); err != nil { + c.restoreLocked(snapshot) + return err + } + return nil +} + +// restoreLocked puts back a JSON snapshot taken under the same lock. +func (c *Config) restoreLocked(snapshot []byte) { + c.Sites, c.Models = nil, nil + c.ProxyPool.Entries, c.ProxyPool.Subscriptions = nil, nil + _ = json.Unmarshal(snapshot, c) } // DataDir resolves (and creates) the runtime data directory. diff --git a/internal/config/regression_test.go b/internal/config/regression_test.go new file mode 100644 index 0000000..7001013 --- /dev/null +++ b/internal/config/regression_test.go @@ -0,0 +1,190 @@ +package config + +import ( + "os" + "path/filepath" + "runtime" + "sync" + "sync/atomic" + "testing" + "time" +) + +// 回归：docker-compose 的 `API_KEYS: ${API_KEYS:-}` 让环境变量恒存在（值为空）， +// 空值必须视为"不覆盖"，否则 config.toml 里的 api_keys 会被清空， +// api/server.go:120 的 len(keys) > 0 判定随之跳过全部鉴权。 +// 对照：TURNSTILE_API_KEY 一直有 != "" 判断。 +func TestEmptyAPIKeysEnvKeepsFileKeys(t *testing.T) { + t.Setenv("API_KEYS", "") + + cfg, err := Load(examplePath) + if err != nil { + t.Fatalf("load: %v", err) + } + if len(cfg.Security.APIKeys) == 0 { + t.Fatalf("空的 API_KEYS 环境变量清空了 config.toml 的 api_keys: %v "+ + "→ api/server.go:120 将跳过全部鉴权", cfg.Security.APIKeys) + } + if cfg.Security.APIKeys[0] != "sk-dsfr-local-change-me" { + t.Fatalf("want file key preserved, got %v", cfg.Security.APIKeys) + } +} + +// 回归：SaveTo 必须持读锁遍历 Sites/Models 等 map —— 与并发的管理端写操作 +// 相撞是 "concurrent map iteration and map write"，属不可 recover 的 fatal +// error，进程直接退出。用 -race 运行可稳定检出。 +func TestSaveToConcurrentWithMutation(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + cfg, err := Load(examplePath) + if err != nil { + t.Fatalf("load: %v", err) + } + cfg.path = path + + stop := make(chan struct{}) + var wg sync.WaitGroup + + // 写者：持锁改 map（与真实 admin action 行为一致） + wg.Add(1) + go func() { + defer wg.Done() + i := 0 + for { + select { + case <-stop: + return + default: + } + cfg.Lock() + cfg.Sites["race-probe"] = &Site{Code: "race-probe", Enabled: false} + if i%4 == 0 { + delete(cfg.Sites, "race-probe") + } + cfg.Unlock() + i++ + } + }() + + // SaveTo：必须在读锁下编码（与真实 s.save() 行为一致） + deadline := time.Now().Add(300 * time.Millisecond) + for time.Now().Before(deadline) { + if err := cfg.SaveTo(path); err != nil { + t.Logf("save: %v", err) + } + } + + close(stop) + wg.Wait() +} + +// 回归：PUT 校验失败必须整体回滚，不能把半套配置留在内存里。 +func TestReplaceJSONRollsBackOnInvalidPayload(t *testing.T) { + cfg, err := Load(examplePath) + if err != nil { + t.Fatalf("load: %v", err) + } + before := cfg.Server.Port + + err = cfg.ReplaceJSON([]byte(`{"server":{"port":99999,"host":"127.0.0.1","log_level":"INFO","log_file":""}}`)) + if err == nil { + t.Fatalf("want validation error for port 99999, got nil") + } + if cfg.Server.Port != before { + t.Fatalf("失败的 PUT 没有回滚: port=%d, want %d", cfg.Server.Port, before) + } +} + +// 回归：控制台 GET 会把两个密钥抹成空串再原样回传，空值必须表示"保持不变"。 +func TestReplaceJSONBlankSecretsKeepCurrent(t *testing.T) { + cfg, err := Load(examplePath) + if err != nil { + t.Fatalf("load: %v", err) + } + cfg.Admin.Password = "admin123" + cfg.Turnstile.APIKey = "solver-secret" + + payload := []byte(`{"admin":{"enabled":true,"host":"127.0.0.1","port":8001,"password":""},` + + `"turnstile":{"enabled":false,"api_key":""},"server":{"host":"127.0.0.1","port":18000,` + + `"log_level":"INFO","log_file":""},"security":{"api_keys":["sk-a"]}}`) + if err := cfg.ReplaceJSON(payload); err != nil { + t.Fatalf("replace: %v", err) + } + if cfg.Admin.Password != "admin123" { + t.Fatalf("空 password 覆盖成了 %q", cfg.Admin.Password) + } + if cfg.Turnstile.APIKey != "solver-secret" { + t.Fatalf("空 api_key 覆盖成了 %q", cfg.Turnstile.APIKey) + } + if cfg.Server.Port != 18000 { + t.Fatalf("port = %d, want 18000", cfg.Server.Port) + } +} + +// 回归：SaveTo 必须原子替换目标文件。裸 os.WriteFile 是"先截断再写"， +// 进程崩在中间、或并发读者正好撞上截断窗口，就留下空/半截 config.toml +// —— 下次启动直接 config error 起不来。顺带把权限从 0644 收紧到 0600 +// （config 里有 api_keys 和 admin password）。 +func TestSaveToIsAtomicAndPrivate(t *testing.T) { + cfg, err := Load(examplePath) + if err != nil { + t.Fatalf("load: %v", err) + } + path := filepath.Join(t.TempDir(), "config.toml") + if err := cfg.SaveTo(path); err != nil { + t.Fatalf("save: %v", err) + } + want, err := os.Stat(path) + if err != nil { + t.Fatalf("stat: %v", err) + } + + stop := make(chan struct{}) + var wg sync.WaitGroup + var bad int32 + for i := 0; i < 4; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + raw, err := os.ReadFile(path) + if err != nil { + continue + } + // 完整文件 ~7KB，撞上截断窗口会读到 0 或半截 + if int64(len(raw)) < want.Size()-8 { + atomic.AddInt32(&bad, 1) + } + } + }() + } + + for i := 0; i < 30; i++ { + cfg.Lock() + cfg.Server.Port = 8000 + i + cfg.Unlock() + if err := cfg.SaveTo(path); err != nil { + t.Fatalf("save #%d: %v", i, err) + } + } + close(stop) + wg.Wait() + + if bad > 0 { + t.Fatalf("并发读者读到了写了一半的 config: %d 次", bad) + } + if _, err := os.Stat(path + ".tmp"); err == nil { + t.Fatalf("临时文件没清理: %s", path+".tmp") + } + st, err := os.Stat(path) + if err != nil { + t.Fatalf("stat: %v", err) + } + if runtime.GOOS != "windows" && st.Mode().Perm() != 0o600 { + t.Fatalf("perm = %o, want 600（config 里有 api_keys / admin password）", st.Mode().Perm()) + } +} diff --git a/internal/httpx/client.go b/internal/httpx/client.go index 1570220..1838a29 100644 --- a/internal/httpx/client.go +++ b/internal/httpx/client.go @@ -6,6 +6,7 @@ package httpx import ( "bytes" "context" + "errors" "fmt" "io" "net/http" @@ -13,6 +14,7 @@ import ( "sort" "sync" "time" + "unicode/utf8" fhttp "github.com/bogdanfinn/fhttp" tls_client "github.com/bogdanfinn/tls-client" @@ -22,6 +24,35 @@ import ( // DefaultUserAgent mirrors config.DefaultUserAgent without an import cycle. const DefaultUserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/149.0.0.0 Safari/537.36" +// DefaultProfile 返回出站连接使用的 TLS/HTTP2 指纹档案。必须与 +// DefaultUserAgent 保持同一代际：JA3 指纹说 Chrome/120、UA 却说 +// Chrome/149，属于自相矛盾的客户端特征，容易被上游当成异常流量。 +func DefaultProfile() profiles.ClientProfile { + return profiles.Chrome_150 +} + +// DefaultMaxBodyBytes 非流式响应体的默认读取上限。上游（或中间人）发一个 +// 无限大的 body 时，io.ReadAll 会把网关内存直接吃光。 +const DefaultMaxBodyBytes int64 = 64 << 20 + +// ErrBodyTooLarge 响应体超过读取上限时返回（原样透出，方便上层判断）。 +var ErrBodyTooLarge = errors.New("response body exceeds the configured limit") + +// readLimited 读满上限即报错，不把超限部分继续吞进内存。 +func readLimited(r io.Reader, max int64) ([]byte, error) { + if max <= 0 { + max = DefaultMaxBodyBytes + } + b, err := io.ReadAll(io.LimitReader(r, max+1)) + if err != nil { + return nil, err + } + if int64(len(b)) > max { + return nil, ErrBodyTooLarge + } + return b, nil +} + type Request struct { Method string URL string @@ -49,6 +80,8 @@ type Session interface { type Options struct { ProxyURL string Timeout time.Duration + // MaxBodyBytes 非流式响应体读取上限，0 = DefaultMaxBodyBytes。 + MaxBodyBytes int64 } type tlsSession struct { @@ -57,6 +90,8 @@ type tlsSession struct { mu sync.Mutex cookies map[string]string closed bool + + maxBody int64 } // NewSession builds a Chrome-impersonating session. @@ -66,7 +101,7 @@ func NewSession(opts Options) (Session, error) { timeout = 60 * time.Second } tlsOpts := []tls_client.HttpClientOption{ - tls_client.WithClientProfile(profiles.Chrome_120), + tls_client.WithClientProfile(DefaultProfile()), tls_client.WithTimeoutSeconds(int(timeout.Seconds()) + 1), tls_client.WithRandomTLSExtensionOrder(), } @@ -80,6 +115,7 @@ func NewSession(opts Options) (Session, error) { return &tlsSession{ inner: inner, cookies: map[string]string{}, + maxBody: opts.MaxBodyBytes, }, nil } @@ -204,7 +240,7 @@ func (s *tlsSession) Do(ctx context.Context, r Request) (*Response, error) { } defer resp.Body.Close() s.capture(resp) - body, err := io.ReadAll(resp.Body) + body, err := readLimited(resp.Body, s.maxBody) if err != nil { return nil, fmt.Errorf("read body: %w", err) } @@ -227,7 +263,7 @@ func (s *tlsSession) Stream(ctx context.Context, r Request) (io.ReadCloser, int, } s.capture(resp) if resp.StatusCode >= 400 { - body, _ := io.ReadAll(resp.Body) + body, _ := readLimited(resp.Body, s.maxBody) resp.Body.Close() return nil, resp.StatusCode, fmt.Errorf("upstream status %d: %s", resp.StatusCode, truncate(string(body), 400)) } @@ -238,6 +274,10 @@ func truncate(s string, n int) string { if len(s) <= n { return s } + // 按字节截断可能切碎多字节字符，回退到字符边界，避免产出非法 UTF-8。 + for n > 0 && !utf8.RuneStart(s[n]) { + n-- + } return s[:n] + "..." } diff --git a/internal/httpx/regression_test.go b/internal/httpx/regression_test.go new file mode 100644 index 0000000..9045fc8 --- /dev/null +++ b/internal/httpx/regression_test.go @@ -0,0 +1,123 @@ +package httpx + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" + "unicode/utf8" +) + +// 回归：按字节截断会切碎多字节字符，产出非法 UTF-8 的错误文本 +// （随后进 JSON/日志时会变成乱码）。 +func TestTruncateKeepsUTF8Valid(t *testing.T) { + src := strings.Repeat("汉", 100) // 300 bytes + got := truncate(src, 10) + if !utf8.ValidString(got) { + t.Fatalf("truncate 切出了非法 UTF-8: %q (bytes=%v)", got, []byte(got[:min(len(got), 14)])) + } + if !strings.HasSuffix(got, "...") { + t.Fatalf("got %q, want an ellipsis suffix", got) + } +} + +// 回归：不超长的字符串原样返回。 +func TestTruncateShortInputUnchanged(t *testing.T) { + if got := truncate("hello", 10); got != "hello" { + t.Fatalf("got %q", got) + } +} + +// 回归：非流式响应体必须有读取上限。裸 io.ReadAll 会让上游（或中间人） +// 用一个超大响应把网关内存打爆（OOM）。 +func TestDoRejectsOversizedBody(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + chunk := make([]byte, 32<<10) + for i := 0; i < 64; i++ { // 2MB，远超下面设的 64KB 上限 + if _, err := w.Write(chunk); err != nil { + return + } + } + })) + defer srv.Close() + + s, err := NewSession(Options{Timeout: 15 * time.Second, MaxBodyBytes: 64 << 10}) + if err != nil { + t.Fatalf("session: %v", err) + } + defer s.Close() + + resp, err := s.Do(context.Background(), Request{Method: "GET", URL: srv.URL}) + if !errors.Is(err, ErrBodyTooLarge) { + t.Fatalf("err = %v (resp=%v), want ErrBodyTooLarge", err, resp) + } +} + +// 回归：上限以内的响应体必须原样读全，不能误伤正常请求。 +func TestDoReadsBodyWithinLimit(t *testing.T) { + want := strings.Repeat("x", 4096) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(want)) + })) + defer srv.Close() + + s, err := NewSession(Options{Timeout: 15 * time.Second, MaxBodyBytes: 64 << 10}) + if err != nil { + t.Fatalf("session: %v", err) + } + defer s.Close() + + resp, err := s.Do(context.Background(), Request{Method: "GET", URL: srv.URL}) + if err != nil { + t.Fatalf("do: %v", err) + } + if string(resp.Body) != want { + t.Fatalf("body len = %d, want %d", len(resp.Body), len(want)) + } +} + +// 回归：TLS 指纹档案的 Chrome 主版本必须与默认 UA 的主版本贴近。 +// 曾经 UA=Chrome/149 而指纹=Chrome_120，差 29 个大版本 —— JA3 指纹与 +// UA 自相矛盾，是可识别的异常客户端特征。 +func TestTLSFingerprintAlignedWithDefaultUA(t *testing.T) { + uaMajor, ok := uaChromeMajor(DefaultUserAgent) + if !ok { + t.Fatalf("DefaultUserAgent 里解析不出 Chrome 主版本: %q", DefaultUserAgent) + } + id := DefaultProfile().GetClientHelloId() + if !strings.EqualFold(id.Client, "Chrome") { + t.Fatalf("指纹档案客户端 = %q, want Chrome", id.Client) + } + profMajor, err := strconv.Atoi(id.Version) + if err != nil { + t.Fatalf("指纹档案版本 = %q, want 纯数字主版本", id.Version) + } + if d := uaMajor - profMajor; d < -1 || d > 1 { + t.Fatalf("UA 主版本 %d 与指纹档案主版本 %d 相差 %d（>1），指纹与 UA 脱节", + uaMajor, profMajor, d) + } +} + +// uaChromeMajor 从 UA 串里取出 "Chrome/" 的主版本号。 +func uaChromeMajor(ua string) (int, bool) { + const marker = "Chrome/" + i := strings.Index(ua, marker) + if i < 0 { + return 0, false + } + rest := ua[i+len(marker):] + j := strings.Index(rest, ".") + if j <= 0 { + return 0, false + } + n, err := strconv.Atoi(rest[:j]) + if err != nil { + return 0, false + } + return n, true +} diff --git a/internal/logbuf/logbuf.go b/internal/logbuf/logbuf.go index df2354b..2989ba5 100644 --- a/internal/logbuf/logbuf.go +++ b/internal/logbuf/logbuf.go @@ -32,19 +32,19 @@ func New(limit int) *Buffer { return &Buffer{subs: map[int]chan Entry{}, limit: limit} } +// Add appends an entry and fans it out. The sends stay under the lock: an +// unlocked send races Subscribe's unsubscribe closure (which closes the +// channel under the same lock) and "send on closed channel" is a panic that +// takes down whatever goroutine was logging. The sends never block, so the +// lock is only held for a bounded number of non-blocking sends. func (b *Buffer) Add(e Entry) { b.mu.Lock() + defer b.mu.Unlock() b.entries = append(b.entries, e) if len(b.entries) > b.limit { b.entries = b.entries[len(b.entries)-b.limit:] } - chans := make([]chan Entry, 0, len(b.subs)) for _, ch := range b.subs { - chans = append(chans, ch) - } - b.mu.Unlock() - - for _, ch := range chans { select { case ch <- e: default: diff --git a/internal/logbuf/logbuf_test.go b/internal/logbuf/logbuf_test.go new file mode 100644 index 0000000..c2f09fd --- /dev/null +++ b/internal/logbuf/logbuf_test.go @@ -0,0 +1,43 @@ +package logbuf + +import ( + "testing" + "time" +) + +// 回归：Add 必须在持锁状态下发送，否则会与 Subscribe 的取消闭包（锁内 +// close(c)，logbuf.go:82-89）竞争 —— 向已关闭的 channel 发送是 panic， +// 发生在后台 goroutine 时整个进程都会挂。 +// +// 订阅者越多，遍历发送的时间越长，越容易命中这个窗口。 +func TestAddConcurrentUnsubscribeDoesNotFatal(t *testing.T) { + b := New(1000) + + const n = 8000 + unsubs := make([]func(), 0, n) + for i := 0; i < n; i++ { + _, unsub := b.Subscribe() + unsubs = append(unsubs, unsub) + } + + stop := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer close(done) + for { + select { + case <-stop: + return + default: + b.Add(Entry{Level: "INFO", Msg: "probe"}) + } + } + }() + + time.Sleep(20 * time.Millisecond) + for _, u := range unsubs { + u() + } + close(stop) + <-done +} diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index e23acc3..73b1951 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -3,6 +3,7 @@ package metrics import ( "encoding/json" + "log/slog" "os" "path/filepath" "sort" @@ -170,6 +171,13 @@ func (r *Recorder) IncInFlight(delta int64) { r.mu.Unlock() } +// InFlight 当前正在处理的请求数（并发数仪表盘用）。 +func (r *Recorder) InFlight() int64 { + r.mu.Lock() + defer r.mu.Unlock() + return r.inflight +} + // Record stores one completed request. func (r *Recorder) Record(rec Record) { if rec.At.IsZero() { @@ -391,7 +399,10 @@ func (r *Recorder) Reset() { r.latency = nil } -func (r *Recorder) Save() { +// Save 持久化计数。失败会写日志并把错误返回给调用方 —— 之前 +// mkdir/写临时文件/rename 的错误全被静默吞掉，磁盘满或权限出问题时 +// 指标会悄悄丢光、日志里毫无痕迹。 +func (r *Recorder) Save() error { r.mu.Lock() p := persisted{ Total: r.total, Errors: r.errs, Empty: r.empty, Denied: r.denied, @@ -402,20 +413,33 @@ func (r *Recorder) Save() { r.mu.Unlock() if path == "" { - return + return nil } + if err := persistMetrics(path, p); err != nil { + slog.Default().Error("metrics save failed", "path", path, "err", err) + return err + } + return nil +} + +// persistMetrics 写 .tmp 后原子替换；任何一步失败都原样返回错误。 +func persistMetrics(path string, p persisted) error { if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - return + return err } raw, err := json.MarshalIndent(p, "", " ") if err != nil { - return + return err } tmp := path + ".tmp" if err := os.WriteFile(tmp, raw, 0o644); err != nil { - return + return err } - _ = os.Rename(tmp, path) + if err := os.Rename(tmp, path); err != nil { + _ = os.Remove(tmp) + return err + } + return nil } func (r *Recorder) load() { diff --git a/internal/metrics/metrics_test.go b/internal/metrics/metrics_test.go index 04e9553..3c2035a 100644 --- a/internal/metrics/metrics_test.go +++ b/internal/metrics/metrics_test.go @@ -1,6 +1,10 @@ package metrics -import "testing" +import ( + "os" + "path/filepath" + "testing" +) func TestRecordTracksSiteTokensAndQuota(t *testing.T) { r := New(t.TempDir()) @@ -24,3 +28,34 @@ func TestRecordTracksSiteTokensAndQuota(t *testing.T) { t.Errorf("unknown site totals = %+v, want zero", empty) } } + +// 回归：落盘失败必须能被调用方看到。之前 Save 把 mkdir/写文件/rename +// 的错误全部吞掉，磁盘满或权限出问题时指标悄悄丢光、日志毫无痕迹。 +func TestSaveSurfacesPersistFailure(t *testing.T) { + dir := t.TempDir() + // 把目标 metrics.json 建成目录 → rename 必然失败 + if err := os.Mkdir(filepath.Join(dir, "metrics.json"), 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + r := New(dir) + if err := r.Save(); err == nil { + t.Fatal("Save() = nil, want rename failure to surface") + } +} + +// 成功路径返回 nil，且文件真实落盘。 +func TestSaveSucceedsNormally(t *testing.T) { + dir := t.TempDir() + r := New(dir) + r.Record(Record{Model: "m", Site: "de", PromptTokens: 1, Status: StatusOK}) + if err := r.Save(); err != nil { + t.Fatalf("Save() = %v, want nil", err) + } + raw, err := os.ReadFile(filepath.Join(dir, "metrics.json")) + if err != nil { + t.Fatalf("read persisted metrics: %v", err) + } + if len(raw) == 0 { + t.Fatal("metrics.json is empty") + } +} diff --git a/internal/openai/prompt.go b/internal/openai/prompt.go index a71d2b7..d8d66ab 100644 --- a/internal/openai/prompt.go +++ b/internal/openai/prompt.go @@ -94,8 +94,18 @@ func samplingPrompt(opts PromptOptions) string { if len(opts.Stop) > 0 { add("- Stop sequences (do not emit them): %s", strings.Join(opts.Stop, " | ")) } - if opts.ResponseFormat != nil && opts.ResponseFormat.Type == "json_object" { - add("- Respond with a single valid JSON object only, no prose, no markdown fences.") + if opts.ResponseFormat != nil { + switch opts.ResponseFormat.Type { + case "json_object": + add("- Respond with a single valid JSON object only, no prose, no markdown fences.") + case "json_schema": + add("- Respond with a single valid JSON object only, no prose, no markdown fences.") + // schema 必须进入约束，否则声明的 json_schema 完全没生效， + // 客户端拿到的对象无法通过它自己的 schema 校验。 + if schema := jsonSchemaConstraint(opts.ResponseFormat.JSONSchema); schema != "" { + add("- The JSON must validate against this JSON Schema: %s", schema) + } + } } if len(lines) == 0 { return "" @@ -103,6 +113,21 @@ func samplingPrompt(opts PromptOptions) string { return "=== GENERATION CONSTRAINTS ===\n" + strings.Join(lines, "\n") + "\n=== END GENERATION CONSTRAINTS ===" } +// jsonSchemaConstraint 取出 response_format.json_schema 里的 schema 主体； +// 结构不是预期形状时退回原始负载，至少让模型看到声明。 +func jsonSchemaConstraint(raw json.RawMessage) string { + if len(raw) == 0 { + return "" + } + var fs struct { + Schema json.RawMessage `json:"schema"` + } + if err := json.Unmarshal(raw, &fs); err == nil && len(fs.Schema) > 0 { + return string(fs.Schema) + } + return string(raw) +} + // ResponsesToMessages converts an OpenAI Responses API request into chat // messages. Besides plain message items it understands function_call and // function_call_output items so multi-turn tool loops round-trip. diff --git a/internal/openai/regression_test.go b/internal/openai/regression_test.go new file mode 100644 index 0000000..b2f608b --- /dev/null +++ b/internal/openai/regression_test.go @@ -0,0 +1,81 @@ +package openai + +import ( + "encoding/json" + "reflect" + "strings" + "testing" +) + +// 回归：OpenAI 的 stop 允许单个字符串或字符串数组；字段类型是 []string 时 +// 单字符串会让整个请求体解码失败（400），而这是合法的 OpenAI 请求。 +func TestStopAcceptsSingleString(t *testing.T) { + var req ChatCompletionRequest + raw := `{"model":"m","messages":[{"role":"user","content":"hi"}],"stop":"END"}` + if err := json.Unmarshal([]byte(raw), &req); err != nil { + t.Fatalf("stop 写成字符串时解码失败（OpenAI 合法请求）: %v", err) + } + if !reflect.DeepEqual([]string(req.Stop), []string{"END"}) { + t.Fatalf("stop = %#v, want [END]", req.Stop) + } +} + +// 回归：数组写法必须照常工作。 +func TestStopStillAcceptsArray(t *testing.T) { + var req ChatCompletionRequest + raw := `{"model":"m","messages":[],"stop":["A","B"]}` + if err := json.Unmarshal([]byte(raw), &req); err != nil { + t.Fatalf("数组写法解码失败: %v", err) + } + if !reflect.DeepEqual([]string(req.Stop), []string{"A", "B"}) { + t.Fatalf("stop = %#v, want [A B]", req.Stop) + } +} + +// 回归：response_format.type = json_schema 时，schema 必须进入生成约束， +// 否则结构化输出没有任何约束，客户端拿到的对象不符合声明的 schema。 +func TestPromptCarriesJSONSchema(t *testing.T) { + const schema = `{"type":"object","properties":{"title":{"type":"string"}},"required":["title"]}` + prompt := BuildPrompt([]ChatMessage{{Role: "user", Content: MessageContent{Text: "hi"}}}, + PromptOptions{ResponseFormat: &ResponseFormat{Type: "json_schema", JSONSchema: json.RawMessage(schema)}}) + + if !strings.Contains(prompt, "json_schema") && !strings.Contains(prompt, "JSON Schema") { + t.Fatalf("提示词没有提到 schema:\n%s", prompt) + } + for _, want := range []string{"title", `"required"`} { + if !strings.Contains(prompt, want) { + t.Fatalf("提示词缺少 schema 内容 %q:\n%s", want, prompt) + } + } +} + +// 回归：json_object 老行为不能回归。 +func TestPromptStillHandlesJSONObject(t *testing.T) { + prompt := BuildPrompt([]ChatMessage{{Role: "user", Content: MessageContent{Text: "hi"}}}, + PromptOptions{ResponseFormat: &ResponseFormat{Type: "json_object"}}) + if !strings.Contains(prompt, "JSON object") { + t.Fatalf("json_object 提示缺失:\n%s", prompt) + } +} + +// 回归：tool_calls 条目缺少/损坏的 function 时不能吞掉错误并产出一个 +// name 为空的调用丢给客户端（客户端会在执行阶段才炸）。 +func TestMalformedToolCallFunctionIsRejected(t *testing.T) { + _, calls, ok := TryParseToolCalls(`{"tool_calls":[{"id":"c1","type":"function"}]}`) + for _, c := range calls { + if c.Function.Name == "" { + t.Fatalf("解析出 name 为空的 tool call: %+v (ok=%v)", c, ok) + } + } +} + +// 回归：正常的 tool_calls 解析不能受影响。 +func TestValidToolCallStillParses(t *testing.T) { + _, calls, ok := TryParseToolCalls(`{"tool_calls":[{"id":"c1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"SH\"}"}}]}`) + if !ok { + t.Fatalf("正常 tool_calls 没被识别") + } + if len(calls) != 1 || calls[0].Function.Name != "get_weather" { + t.Fatalf("calls = %+v", calls) + } +} diff --git a/internal/openai/toolcalls.go b/internal/openai/toolcalls.go index 3b589ca..991f465 100644 --- a/internal/openai/toolcalls.go +++ b/internal/openai/toolcalls.go @@ -63,7 +63,13 @@ func TryParseToolCalls(text string) (string, []ToolCall, bool) { Name string `json:"name"` Arguments json.RawMessage `json:"arguments"` } - _ = json.Unmarshal(tc.Function, &fn) + // function 字段缺失或解析不出名字的条目是坏数据：标记 bad 让整 + // 个候选被丢弃，而不是把一个 name 为空的 tool call 交给客户端 + // （错误被吞掉后要到执行阶段才暴露）。 + if err := json.Unmarshal(tc.Function, &fn); err != nil || fn.Name == "" { + bad = true + break + } args := "{}" if len(fn.Arguments) > 0 { var s string diff --git a/internal/openai/types.go b/internal/openai/types.go index 13543c1..a0d3a6f 100644 --- a/internal/openai/types.go +++ b/internal/openai/types.go @@ -1,6 +1,9 @@ package openai -import "encoding/json" +import ( + "encoding/json" + "fmt" +) type ChatMessage struct { Role string `json:"role"` @@ -67,6 +70,32 @@ type StreamOptions struct { IncludeUsage bool `json:"include_usage,omitempty"` } +// StopSequences 兼容 OpenAI 对 stop 的两种写法：单个字符串或字符串数组。 +// 直接用 []string 会让 `"stop": "END"` 这种合法请求整个解码失败（400）。 +type StopSequences []string + +func (s *StopSequences) UnmarshalJSON(data []byte) error { + if string(data) == "null" { + *s = nil + return nil + } + var one string + if err := json.Unmarshal(data, &one); err == nil { + if one == "" { + *s = nil + } else { + *s = StopSequences{one} + } + return nil + } + var many []string + if err := json.Unmarshal(data, &many); err != nil { + return fmt.Errorf("stop must be a string or an array of strings") + } + *s = many + return nil +} + type ChatCompletionRequest struct { Model string `json:"model"` Messages []ChatMessage `json:"messages"` @@ -78,7 +107,7 @@ type ChatCompletionRequest struct { Tools []ToolDef `json:"tools,omitempty"` ToolChoice json.RawMessage `json:"tool_choice,omitempty"` User string `json:"user,omitempty"` - Stop []string `json:"stop,omitempty"` + Stop StopSequences `json:"stop,omitempty"` ResponseFormat *ResponseFormat `json:"response_format,omitempty"` StreamOptions *StreamOptions `json:"stream_options,omitempty"` N *int `json:"n,omitempty"` @@ -132,6 +161,15 @@ type ResponsesRequest struct { MaxOutputTokens *int `json:"max_output_tokens,omitempty"` Tools []ResponsesTool `json:"tools,omitempty"` ToolChoice json.RawMessage `json:"tool_choice,omitempty"` + // Text 是 Responses API 的输出配置块，text.format 承载输出格式约束 + // （json_schema / json_object / text）。不解析它，客户端指定的 + // json_schema 会被静默丢弃，模型输出的 JSON 无从校验。 + Text *TextFormat `json:"text,omitempty"` +} + +// TextFormat 是 Responses API 的 text 对象，这里只消费 format。 +type TextFormat struct { + Format *ResponseFormat `json:"format,omitempty"` } // ResponsesTool is a function tool in the flat Responses API shape; the nested diff --git a/internal/proxypool/manager.go b/internal/proxypool/manager.go index c738b97..8c62632 100644 --- a/internal/proxypool/manager.go +++ b/internal/proxypool/manager.go @@ -65,6 +65,13 @@ type Manager struct { started bool lastCheck time.Time lastErr string + + // 关闭协调：stopCh 通知循环退出，doneCh 在循环真正退出后关闭， + // loopCancel 取消 Start 传入的 ctx（让进行中的 cycle 立刻收尾）。 + stopCh chan struct{} + doneCh chan struct{} + loopCancel context.CancelFunc + closeOnce sync.Once } func New(cfg *config.Config, dataDir string, log *slog.Logger) *Manager { @@ -80,6 +87,8 @@ func New(cfg *config.Config, dataDir string, log *slog.Logger) *Manager { subs: map[string]*subState{}, sticky: map[string]string{}, fetching: map[string]bool{}, + stopCh: make(chan struct{}), + doneCh: make(chan struct{}), } } @@ -92,11 +101,36 @@ func (m *Manager) Start(ctx context.Context) { return } m.started = true + ctx, m.loopCancel = context.WithCancel(ctx) m.mu.Unlock() go m.loop(ctx) } +// Close 停掉后台循环并同步回收 xray 子进程，阻塞到两件事都完成为止。 +// 进程退出路径必须调用它：仅靠 ctx 取消是异步的，主 goroutine 退出时 +// 循环可能还没跑到 xray.stop()，子进程会变成孤儿（Windows 上父进程退出 +// 不会带走子进程）。 +func (m *Manager) Close() { + m.closeOnce.Do(func() { + m.mu.Lock() + cancel := m.loopCancel + started := m.started + m.mu.Unlock() + if cancel != nil { + cancel() + } + close(m.stopCh) + if !started { + // 循环没启动过，没有 doneCh 可等，直接回收。 + m.xray.stop() + return + } + <-m.doneCh + }) +} + func (m *Manager) loop(ctx context.Context) { + defer close(m.doneCh) ticker := time.NewTicker(5 * time.Second) defer ticker.Stop() for { @@ -104,6 +138,9 @@ func (m *Manager) loop(ctx context.Context) { case <-ctx.Done(): m.xray.stop() return + case <-m.stopCh: + m.xray.stop() + return case <-ticker.C: } m.cycle(ctx) @@ -616,9 +653,20 @@ func (m *Manager) Status() Status { m.cfg.RUnlock() m.mu.Lock() - endpoints := make([]*endpoint, 0, len(m.endpoints)) + // 在锁内把 endpoint 快照成值类型：检查循环/Report 会在 m.mu 下改 + // ep.healthy/lastErr/latencyMs，锁外再读这些字段就是数据竞争。 + entries := make([]EntryStatus, 0, len(m.endpoints)) for _, ep := range m.endpoints { - endpoints = append(endpoints, ep) + es := EntryStatus{ + Name: ep.name, Kind: ep.kind, Scheme: ep.scheme, Target: ep.display, + Source: ep.source, Enabled: ep.enabled, Healthy: ep.healthy, + Checked: ep.checked, LatencyMs: ep.latencyMs, LastErr: ep.lastErr, + ParseError: ep.parseErr, + } + if !ep.checkedAt.IsZero() { + es.CheckedAt = ep.checkedAt.Unix() + } + entries = append(entries, es) } sticky := make(map[string]string, len(m.sticky)) for k, v := range m.sticky { @@ -642,21 +690,10 @@ func (m *Manager) Status() Status { XrayVersion: poolCfg.XrayVersion, XrayAutoDownload: poolCfg.XrayAutoDownload, Xray: m.xray.status(), + Entries: entries, Sticky: sticky, LastError: lastErr, } - for _, ep := range endpoints { - es := EntryStatus{ - Name: ep.name, Kind: ep.kind, Scheme: ep.scheme, Target: ep.display, - Source: ep.source, Enabled: ep.enabled, Healthy: ep.healthy, - Checked: ep.checked, LatencyMs: ep.latencyMs, LastErr: ep.lastErr, - ParseError: ep.parseErr, - } - if !ep.checkedAt.IsZero() { - es.CheckedAt = ep.checkedAt.Unix() - } - st.Entries = append(st.Entries, es) - } sort.Slice(st.Entries, func(i, j int) bool { a, b := st.Entries[i], st.Entries[j] if (a.Source == "") != (b.Source == "") { diff --git a/internal/proxypool/regression_test.go b/internal/proxypool/regression_test.go new file mode 100644 index 0000000..fd76ada --- /dev/null +++ b/internal/proxypool/regression_test.go @@ -0,0 +1,115 @@ +package proxypool + +import ( + "context" + "os" + "os/exec" + "sync" + "testing" + "time" + + "github.com/nyoungo/dsfree2api/internal/config" +) + +// 回归：Status() 必须在持锁状态读取 endpoint 字段。检查循环/Report 会在 +// m.mu 下写 ep.healthy/lastErr/latencyMs，锁外读它们就是数据竞争 +// （go test -race 会直接报 DATA RACE，线上是控制台刷新时的随机脏数据）。 +func TestStatusDoesNotRaceWithEndpointUpdates(t *testing.T) { + cfg, err := config.Load(examplePath) + if err != nil { + t.Fatalf("load config: %v", err) + } + cfg.ProxyPool.Enabled = true + cfg.ProxyPool.Entries = map[string]*config.ProxyEntry{ + "a": {Link: "http://127.0.0.1:9", Enabled: true}, + "b": {Link: "socks5://127.0.0.1:10", Enabled: true}, + } + m := New(cfg, t.TempDir(), nil) + m.Reload() + + stop := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + // 模拟检查循环：只在锁内写 endpoint 字段 + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + m.mu.Lock() + for _, ep := range m.endpoints { + ep.checked = true + ep.healthy = !ep.healthy + ep.latencyMs++ + ep.lastErr = "probe failed" + ep.checkedAt = time.Now() + } + m.mu.Unlock() + } + }() + + for i := 0; i < 300; i++ { + st := m.Status() + if len(st.Entries) != 2 { + close(stop) + wg.Wait() + t.Fatalf("entries = %d, want 2", len(st.Entries)) + } + } + close(stop) + wg.Wait() +} + +// 供 TestCloseKillsXrayChild 拉起的子进程桩：当成 xray 起着，等着被杀。 +func TestXrayChildStub(t *testing.T) { + if os.Getenv("DSFREE_XRAY_CHILD") != "1" { + t.Skip("只给 Close 回归测试当子进程用") + } + time.Sleep(time.Hour) +} + +// 回归：Close 必须同步杀掉 xray 子进程并等后台循环退出。进程退出时 +// Windows/Linux 都不会替父进程带走子进程，漏杀就是孤儿 xray。 +func TestCloseKillsXrayChild(t *testing.T) { + cfg, err := config.Load(examplePath) + if err != nil { + t.Fatalf("load config: %v", err) + } + m := New(cfg, t.TempDir(), nil) + m.Start(context.Background()) + + cmd := exec.Command(os.Args[0], "-test.run=^TestXrayChildStub$") + cmd.Env = append(os.Environ(), "DSFREE_XRAY_CHILD=1") + if err := cmd.Start(); err != nil { + t.Fatalf("start child: %v", err) + } + exited := make(chan struct{}) + go func() { + _ = cmd.Wait() + close(exited) + }() + + m.xray.mu.Lock() + m.xray.proc = cmd + m.xray.exited = exited + m.xray.running = true + m.xray.pid = cmd.Process.Pid + m.xray.mu.Unlock() + + m.Close() + + select { + case <-exited: + case <-time.After(5 * time.Second): + _ = cmd.Process.Kill() + t.Fatal("Close 没有回收 xray 子进程") + } + if st := m.xray.status(); st.Running || st.PID != 0 { + t.Fatalf("xray 状态没清干净: %+v", st) + } + // 幂等：重复 Close 不能 panic 也不能卡死 + m.Close() +} diff --git a/internal/turnstile/regression_test.go b/internal/turnstile/regression_test.go new file mode 100644 index 0000000..ac76db2 --- /dev/null +++ b/internal/turnstile/regression_test.go @@ -0,0 +1,99 @@ +package turnstile + +import ( + "context" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/nyoungo/dsfree2api/internal/config" + "github.com/nyoungo/dsfree2api/internal/httpx" +) + +// 回归：retries=0 时求解循环必须至少跑一次，不能直接返回 ("", nil) —— +// 那会让 refreshLocked 以为求解成功，把未经验证的 cookie 缓存进池子。 +func TestSolveTokenAPIRetriesZeroStillAttemptsOnce(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"token":"tok-1","elapsed":0.1}`)) + })) + defer srv.Close() + + solver := New(config.Turnstile{ + Enabled: true, Provider: config.ProviderAPI, APIStyle: "ezsolver", + APIURL: srv.URL, TimeoutSeconds: 5, Retries: 0, + }, httpx.DefaultUserAgent) + + token, err := solver.solveTokenAPI(context.Background(), "https://deepseek.de/", "0xabc", solver.coreFor("")) + if err != nil { + t.Fatalf("solve: %v", err) + } + if calls.Load() != 1 { + t.Fatalf("solve API 被调用 %d 次，want 1", calls.Load()) + } + if token == "" { + t.Fatal("retries=0 返回空 token 且无错误 —— 上游根本没被调用") + } +} + +// 回归：retries=0 时校验循环也必须至少跑一次，否则会跳过 verify 直接成功。 +func TestVerifyTokenRetriesZeroStillVerifiesOnce(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"ok":true}`)) + })) + defer srv.Close() + + solver := New(config.Turnstile{ + Enabled: true, Provider: config.ProviderAPI, + APIURL: srv.URL, TimeoutSeconds: 5, Retries: 0, + }, httpx.DefaultUserAgent) + + sess, err := httpx.NewSession(httpx.Options{Timeout: 5 * time.Second}) + if err != nil { + t.Fatalf("session: %v", err) + } + defer sess.Close() + + site := &config.Site{ + Code: "t", BaseURL: srv.URL, AJAXURL: srv.URL, + VerifyAction: "verify", Language: "en", + } + if err := solver.verifyToken(context.Background(), sess, site, "tok"); err != nil { + t.Fatalf("verify: %v", err) + } + if calls.Load() != 1 { + t.Fatalf("verify 被调用 %d 次，want 1 —— 空 token 会被判定为校验成功", calls.Load()) + } +} + +// 回归：coreFor 拿到非法代理 URL 时不能吞掉 StdClient 的错误并缓存 nil client， +// 否则 solveTokenAPI 里 client.Do 会 nil panic（warmer 后台协程没有 recover， +// 一次就能拖垮整个进程）。 +func TestCoreForInvalidProxyDoesNotPanic(t *testing.T) { + solver := New(config.Turnstile{ + Enabled: true, Provider: config.ProviderAPI, APIStyle: "ezsolver", + APIURL: "https://solve.example.invalid/", TimeoutSeconds: 5, Retries: 1, + }, httpx.DefaultUserAgent) + + c := solver.coreFor("://bad proxy") + if c == nil { + t.Fatal("coreFor returned nil") + } + defer func() { + if r := recover(); r != nil { + t.Fatalf("coreFor 留下 nil client 导致 panic: %v", r) + } + }() + + token, err := solver.solveTokenAPI(context.Background(), "https://deepseek.de/", "0xabc", c) + if err == nil { + t.Fatalf("非法代理应当报错，却返回 token=%q", token) + } +} diff --git a/internal/turnstile/solver.go b/internal/turnstile/solver.go index 3cc995f..4ff37c2 100644 --- a/internal/turnstile/solver.go +++ b/internal/turnstile/solver.go @@ -66,9 +66,12 @@ type cookieState struct { type core struct { proxy string client *http.Client - mu sync.Mutex - cache map[string]*cookieState - locks map[string]*sync.Mutex + // clientErr 记录 StdClient 构建失败的原因：代理 URL 非法时 client 为 nil， + // 调用方必须显式报错而不是 nil panic。 + clientErr error + mu sync.Mutex + cache map[string]*cookieState + locks map[string]*sync.Mutex } func (c *core) lockFor(site string) *sync.Mutex { @@ -140,18 +143,29 @@ func (s *Solver) coreFor(proxy string) *core { defer s.mu.Unlock() c, ok := s.cores[key] if !ok { - client, _ := httpx.StdClient(proxy, time.Duration(s.cfg.SolveTimeoutValue()+30)*time.Second) + client, clientErr := httpx.StdClient(proxy, time.Duration(s.cfg.SolveTimeoutValue()+30)*time.Second) c = &core{ - proxy: key, - client: client, - cache: map[string]*cookieState{}, - locks: map[string]*sync.Mutex{}, + proxy: key, + client: client, + clientErr: clientErr, + cache: map[string]*cookieState{}, + locks: map[string]*sync.Mutex{}, } s.cores[key] = c } return c } +// solveAttempts 把 retries 归一为至少 1 次尝试。retries 表示"额外重试次数"， +// 0 不能意味着一次都不做：那会让 solveToken 返回 ("", nil)、verifyToken 直接 +// 返回 nil，refreshLocked 便把未经求解/校验的 cookie 当成功缓存。 +func solveAttempts(retries int) int { + if retries < 1 { + return 1 + } + return retries +} + func (s *Solver) Invalidate(siteCode string) { s.mu.Lock() cores := make([]*core, 0, len(s.cores)) @@ -423,8 +437,17 @@ func (s *Solver) solveTokenAPI(ctx context.Context, baseURL, sitekey string, c * client := c.client if isLocalHost(cfg.APIURL) { client = &http.Client{Timeout: time.Duration(cfg.SolveTimeoutValue()+30) * time.Second} + } else if client == nil { + // 代理 URL 非法导致 StdClient 构建失败：显式失败，不能带着 nil + // client 进循环（client.Do 会 nil panic）。 + msg := "turnstile solve client unavailable" + if c.clientErr != nil { + msg = fmt.Sprintf("proxy %q is invalid: %v", c.proxy, c.clientErr) + } + return "", &Error{Msg: msg, Retryable: false} } - for attempt := 0; attempt < cfg.Retries; attempt++ { + attempts := solveAttempts(cfg.Retries) + for attempt := 0; attempt < attempts; attempt++ { req, err := http.NewRequestWithContext(ctx, http.MethodPost, cfg.APIURL, bytes.NewReader(payload)) if err != nil { return "", err @@ -541,7 +564,7 @@ func (s *Solver) verifyToken(ctx context.Context, sess httpx.Session, site *conf "Content-Type": "application/x-www-form-urlencoded", } var last error - for attempt := 0; attempt < cfg.Retries; attempt++ { + for attempt := 0; attempt < solveAttempts(cfg.Retries); attempt++ { resp, err := sess.Do(ctx, httpx.Request{ Method: "POST", URL: site.AJAXURL, diff --git a/internal/upstream/client.go b/internal/upstream/client.go index 72e80f9..c10861b 100644 --- a/internal/upstream/client.go +++ b/internal/upstream/client.go @@ -11,12 +11,14 @@ import ( "html" "io" "log/slog" + "net/http" "os" "regexp" "strconv" "strings" "sync" "time" + "unicode/utf8" "github.com/nyoungo/dsfree2api/internal/config" "github.com/nyoungo/dsfree2api/internal/httpx" @@ -65,7 +67,7 @@ type Client struct { stateMu sync.Mutex chatCfg map[string]chatConfig cfgLocks map[string]*sync.Mutex - gates map[string]chan struct{} + gates map[string]*siteGate // quotaCool marks sites that recently answered with quota errors; they // are tried after the healthy mirrors until the cooldown expires. @@ -96,7 +98,7 @@ func New(cfg *config.Config, ts *turnstile.Solver, log *slog.Logger, pool RouteS pool: pool, chatCfg: map[string]chatConfig{}, cfgLocks: map[string]*sync.Mutex{}, - gates: map[string]chan struct{}{}, + gates: map[string]*siteGate{}, quotaCool: map[string]time.Time{}, } } @@ -156,15 +158,14 @@ func (c *Client) quotaCooled(code string) bool { // the transition, not every probe that confirms it. func (c *Client) setQuotaCool(code string) { c.stateMu.Lock() - until, ok := c.quotaCool[code] - now := time.Now() - if ok && now.Add(quotaCoolTTL).Before(until) { - c.stateMu.Unlock() - return - } - c.quotaCool[code] = now.Add(quotaCoolTTL) + _, wasCooled := c.quotaCool[code] + // 每次都把截止时间顺延，重复探测只延长窗口；只有从"未冷却"进入 + // "冷却"这一次状态翻转才记 INFO，之后的确认探测不再刷日志。 + c.quotaCool[code] = time.Now().Add(quotaCoolTTL) c.stateMu.Unlock() - c.log.Info("site quota cooldown", "site", code, "for", quotaCoolTTL) + if !wasCooled { + c.log.Info("site quota cooldown", "site", code, "for", quotaCoolTTL) + } } // sortQuotaCoolLast moves quota-cooled sites behind the healthy candidates, @@ -738,7 +739,7 @@ func (c *Client) chatOnce( headers := ajaxHeaders(site, referer, ua) headers["Accept"] = "text/event-stream" - rc, _, err := sess.Stream(streamCtx, httpx.Request{ + rc, status, err := sess.Stream(streamCtx, httpx.Request{ Method: "GET", URL: cc.AJAXURL + "?" + params, Header: headers, @@ -747,7 +748,7 @@ func (c *Client) chatOnce( if streamCtx.Err() != nil && ctx.Err() == nil { return errf("upstream stream timeout: %v", err) } - return AsUpstream(err) + return streamHTTPError(status, err) } defer rc.Close() @@ -772,6 +773,18 @@ func (c *Client) chatOnce( }) } +// streamHTTPError classifies a stream setup failure by its HTTP status. 401/403 +// mean the cached session is no longer accepted (expired nonce / Turnstile +// cookie), which must be reported as a session error so the caller refreshes +// cookies instead of blaming the egress endpoint (poolReport(false)) and +// keeping the stale cookie. +func streamHTTPError(status int, err error) error { + if status == http.StatusUnauthorized || status == http.StatusForbidden { + return newSession(fmt.Sprintf("upstream stream http %d: %v", status, err)) + } + return AsUpstream(err) +} + // sseTimingReader (diagnosis only, gated by DSFREE_SSE_TIMING) logs when bytes // arrive from the upstream stream. type sseTimingReader struct { @@ -997,25 +1010,51 @@ func (c *Client) lockFor(key string) *sync.Mutex { // ── concurrency gate ──────────────────────────────────────────── +// siteGate 是单个站点的并发闸门。名额用 busy 计数而不是"首次使用时定容的 +// channel"，这样 [limits].max_concurrent_per_site 被修改后立即生效，无需重启 +// （channel 版本只在第一次 acquire 时定容，之后改配置一直不生效）。 +type siteGate struct { + mu sync.Mutex + busy int + // wake 在每次 release 时被关闭并重建，用来广播唤醒等待者；等待者醒来后 + // 按当时最新的配置重新竞争名额。 + wake chan struct{} +} + +func newSiteGate() *siteGate { return &siteGate{wake: make(chan struct{})} } + func (c *Client) acquire(ctx context.Context, site string) error { - c.cfg.RLock() - max := c.cfg.Limits.MaxConcurrentPerSite - c.cfg.RUnlock() - if max <= 0 { - return nil - } c.stateMu.Lock() g, ok := c.gates[site] if !ok { - g = make(chan struct{}, max) + g = newSiteGate() c.gates[site] = g } c.stateMu.Unlock() - select { - case g <- struct{}{}: - return nil - case <-ctx.Done(): - return fmt.Errorf("site %s busy: %w", site, ctx.Err()) + return c.admit(ctx, site, g) +} + +// admit 在当前配置的上限内占一个名额；超限时阻塞到有空位或 ctx 结束。 +func (c *Client) admit(ctx context.Context, site string, g *siteGate) error { + for { + c.cfg.RLock() + max := c.cfg.Limits.MaxConcurrentPerSite + c.cfg.RUnlock() + + g.mu.Lock() + if max <= 0 || g.busy < max { + g.busy++ + g.mu.Unlock() + return nil + } + wake := g.wake + g.mu.Unlock() + + select { + case <-ctx.Done(): + return fmt.Errorf("site %s busy: %w", site, ctx.Err()) + case <-wake: + } } } @@ -1026,10 +1065,13 @@ func (c *Client) release(site string) { if !ok { return } - select { - case <-g: - default: + g.mu.Lock() + if g.busy > 0 { + g.busy-- } + close(g.wake) + g.wake = make(chan struct{}) + g.mu.Unlock() } // ── helpers ───────────────────────────────────────────────────── @@ -1121,6 +1163,10 @@ func truncate(s string, n int) string { if len(s) <= n { return s } + // 按字节截断可能切碎多字节字符，回退到字符边界，避免产出非法 UTF-8。 + for n > 0 && !utf8.RuneStart(s[n]) { + n-- + } return s[:n] + "..." } diff --git a/internal/upstream/regression_test.go b/internal/upstream/regression_test.go new file mode 100644 index 0000000..15f13a9 --- /dev/null +++ b/internal/upstream/regression_test.go @@ -0,0 +1,255 @@ +package upstream + +import ( + "bytes" + "context" + "errors" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + "unicode/utf8" + + "github.com/nyoungo/dsfree2api/internal/httpx" +) + +// 回归：SSE 规范要求同一事件的多个 data 行用 "\n" 拼接后再解析。逐行单独 +// json.Unmarshal 会把跨行 JSON 整条丢弃，导致该事件静默消失。 +func TestTranslateEventJoinsMultiLineData(t *testing.T) { + evs, err := translateEvent("message", []string{`{"delta":`, `"hello"}`}) + if err != nil { + t.Fatalf("translate: %v", err) + } + if len(evs) != 1 || evs[0].Kind != KindDelta || evs[0].Value != "hello" { + t.Fatalf("events = %+v, want one delta %q", evs, "hello") + } +} + +// 回归：readSSE → translateEvent 全链路上，跨行 data 的事件不能丢失， +// 单行事件与 [DONE] 必须照常工作。 +func TestReadSSEMultiLineDataEvent(t *testing.T) { + const body = "event: message\n" + + "data: {\"delta\":\n" + + "data: \"hello\"}\n" + + "\n" + + "data: [DONE]\n" + + "\n" + var got []Event + if err := readSSE(strings.NewReader(body), func(event string, lines []string) error { + evs, err := translateEvent(event, lines) + got = append(got, evs...) + return err + }); err != nil { + t.Fatalf("readSSE: %v", err) + } + if len(got) != 2 { + t.Fatalf("events = %+v, want delta + done", got) + } + if got[0].Kind != KindDelta || got[0].Value != "hello" { + t.Fatalf("first event = %+v, want delta %q", got[0], "hello") + } + if got[1].Kind != KindDone { + t.Fatalf("second event = %+v, want done", got[1]) + } +} + +// 前提：httpx.Stream 对 >=400 把状态码连同错误一起返回，调用方拿得到 401。 +// （chatOnce 里 `rc, _, err := sess.Stream(...)` 把它丢掉了。） +func TestStreamErrorCarriesHTTPStatus(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "expired", http.StatusUnauthorized) + })) + defer srv.Close() + + sess, err := httpx.NewSession(httpx.Options{Timeout: 5 * time.Second}) + if err != nil { + t.Fatalf("session: %v", err) + } + defer sess.Close() + + rc, status, err := sess.Stream(t.Context(), httpx.Request{Method: http.MethodGet, URL: srv.URL}) + if err == nil { + if rc != nil { + rc.Close() + } + t.Fatalf("status %d should have returned an error", status) + } + if status != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", status, http.StatusUnauthorized) + } +} + +// 回归：流建立阶段的 401/403 必须归类为 session 错误，否则 failover 循环会 +// 走 poolReport(proxy, false) 把出口代理当故障标记，同时不会刷新 cookie， +// 带着同一份过期 cookie 重试到超时。 +func TestStreamHTTPErrorClassifiesUnauthorizedAsSession(t *testing.T) { + for _, code := range []int{http.StatusUnauthorized, http.StatusForbidden} { + err := streamHTTPError(code, errors.New("upstream status")) + if !isSession(err) { + t.Fatalf("streamHTTPError(%d) = %#v, want session error", code, err) + } + } + if isSession(streamHTTPError(http.StatusBadGateway, errors.New("upstream status"))) { + t.Fatalf("502 must stay a plain upstream error") + } +} + +// 回归：并发闸门的容量必须跟随 [limits].max_concurrent_per_site 变更。旧实现 +// 用"首次 acquire 时定容的 channel"，改配置（控制台或 PUT /api/config）后要 +// 重启才生效。 +func TestConcurrencyGateFollowsConfigChanges(t *testing.T) { + cfg := testConfig(t) + cfg.Limits.MaxConcurrentPerSite = 4 + c := New(cfg, nil, nil, nil) + ctx := context.Background() + + // 先按 4 建一次闸门 + if err := c.acquire(ctx, "site-a"); err != nil { + t.Fatalf("first acquire: %v", err) + } + c.release("site-a") + + cfg.Limits.MaxConcurrentPerSite = 1 + if err := c.acquire(ctx, "site-a"); err != nil { + t.Fatalf("新上限下的第一个名额: %v", err) + } + short, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + defer cancel() + if err := c.acquire(short, "site-a"); err == nil { + c.release("site-a") + t.Fatalf("上限已改成 1，第二个并发请求却拿到了名额") + } + c.release("site-a") + if err := c.acquire(ctx, "site-a"); err != nil { + t.Fatalf("release 之后应能拿到名额: %v", err) + } + c.release("site-a") +} + +// 回归：上限必须真的挡住并发（防止把闸门改坏）。 +func TestConcurrencyGateBlocksBeyondLimit(t *testing.T) { + cfg := testConfig(t) + cfg.Limits.MaxConcurrentPerSite = 2 + c := New(cfg, nil, nil, nil) + ctx := context.Background() + + if err := c.acquire(ctx, "site-b"); err != nil { + t.Fatalf("slot 1: %v", err) + } + if err := c.acquire(ctx, "site-b"); err != nil { + t.Fatalf("slot 2: %v", err) + } + short, cancel := context.WithTimeout(ctx, 200*time.Millisecond) + defer cancel() + if err := c.acquire(short, "site-b"); err == nil { + t.Fatalf("第三个名额不该发出去") + } + c.release("site-b") + c.release("site-b") + if err := c.acquire(ctx, "site-b"); err != nil { + t.Fatalf("release 后应恢复: %v", err) + } + c.release("site-b") +} + +// 回归：0 = 不限并发（与 config.example.toml 的注释一致）。 +func TestConcurrencyGateZeroMeansUnlimited(t *testing.T) { + cfg := testConfig(t) + cfg.Limits.MaxConcurrentPerSite = 0 + c := New(cfg, nil, nil, nil) + ctx := context.Background() + + for i := 0; i < 8; i++ { + if err := c.acquire(ctx, "site-c"); err != nil { + t.Fatalf("acquire #%d: %v", i, err) + } + } + for i := 0; i < 8; i++ { + c.release("site-c") + } +} + +// 回归：并发增删 + 中途改上限，不能漏发名额也不能多发（-race 下验证记账）。 +func TestConcurrencyGateUnderConcurrentLoad(t *testing.T) { + cfg := testConfig(t) + cfg.Limits.MaxConcurrentPerSite = 3 + c := New(cfg, nil, nil, nil) + + var wg sync.WaitGroup + stop := make(chan struct{}) + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 40; i++ { + cfg.Lock() + cfg.Limits.MaxConcurrentPerSite = []int{1, 3, 5, 0}[i%4] + cfg.Unlock() + time.Sleep(2 * time.Millisecond) + } + close(stop) + }() + for w := 0; w < 4; w++ { + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + if err := c.acquire(ctx, "site-d"); err == nil { + c.release("site-d") + } + cancel() + } + }() + } + wg.Wait() +} + +// 回归：错误信息里的字节截断不能切碎多字节字符。 +func TestUpstreamTruncateKeepsUTF8Valid(t *testing.T) { + got := truncate(strings.Repeat("汉", 100), 10) + if !utf8.ValidString(got) { + t.Fatalf("truncate 切出了非法 UTF-8: %q", got) + } +} + +// 回归：shortJSON 截断同样要落在字符边界上。 +func TestShortJSONKeepsUTF8Valid(t *testing.T) { + got := shortJSON(map[string]any{ + "aaa": strings.Repeat("a", 200), + "error": strings.Repeat("汉", 200), + }) + if !utf8.ValidString(got) { + t.Fatalf("shortJSON 切出了非法 UTF-8: %q", got) + } +} + +// 回归：quota 冷却在窗口内被重复触发时只记一次 INFO（注释说的 +// "the INFO line marks the transition, not every probe that confirms it"）， +// 但冷却截止时间每次都要顺延。 +func TestQuotaCoolLogsOncePerTransition(t *testing.T) { + cfg := testConfig(t) + var buf bytes.Buffer + log := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelInfo})) + c := New(cfg, nil, log, nil) + + c.setQuotaCool("site-a") + first := c.quotaCool["site-a"] + time.Sleep(5 * time.Millisecond) + c.setQuotaCool("site-a") + second := c.quotaCool["site-a"] + + if !second.After(first) { + t.Fatalf("重复触发没有顺延冷却时间: %v -> %v", first, second) + } + if n := strings.Count(buf.String(), "site quota cooldown"); n != 1 { + t.Fatalf("冷却 INFO 条数 = %d, want 1\nlog: %s", n, buf.String()) + } +} diff --git a/internal/upstream/sse.go b/internal/upstream/sse.go index fc5aba6..a5d4bae 100644 --- a/internal/upstream/sse.go +++ b/internal/upstream/sse.go @@ -65,14 +65,21 @@ func readSSE(r io.Reader, handle func(event string, dataLines []string) error) e // translateEvent converts one upstream SSE event into zero or more events for // the downstream client. It returns an error for upstream failure payloads. func translateEvent(eventType string, dataLines []string) ([]Event, error) { - var payload map[string]any for _, line := range dataLines { if line == "[DONE]" { return []Event{{Kind: KindDone}}, nil } - var p map[string]any - if err := json.Unmarshal([]byte(line), &p); err == nil { - payload = p + } + // SSE 规范：同一事件的多个 data 行用 "\n" 拼接才是完整负载，跨行 JSON + // 必须整条解析；逐行解析会把它丢掉。 + var payload map[string]any + if err := json.Unmarshal([]byte(strings.Join(dataLines, "\n")), &payload); err != nil { + // 兼容上游偶尔夹带的非 JSON 行：退化为逐行解析，取最后一个能解析的。 + for _, line := range dataLines { + var p map[string]any + if err := json.Unmarshal([]byte(line), &p); err == nil { + payload = p + } } } if payload == nil { @@ -176,7 +183,7 @@ func shortJSON(payload map[string]any) string { } s := string(raw) if len(s) > 400 { - return s[:400] + "..." + return truncate(s, 400) } return s }