diff --git a/middleware/relay_capture.go b/middleware/relay_capture.go index 480958b..71c0bf3 100644 --- a/middleware/relay_capture.go +++ b/middleware/relay_capture.go @@ -111,21 +111,12 @@ func (collector *responseCaptureCollector) Chunks() []string { type captureResponseWriter struct { gin.ResponseWriter collector *responseCaptureCollector - ch chan<- []byte } func (cw *captureResponseWriter) Write(b []byte) (int, error) { n, err := cw.ResponseWriter.Write(b) if n > 0 { - chunk := b[:n] - if cw.collector != nil { - cw.collector.Add(chunk) - } - if cw.ch != nil { - buf := make([]byte, n) - copy(buf, chunk) - cw.ch <- buf - } + cw.collector.Add(b[:n]) } return n, err } @@ -133,35 +124,11 @@ func (cw *captureResponseWriter) Write(b []byte) (int, error) { func (cw *captureResponseWriter) WriteString(s string) (int, error) { n, err := cw.ResponseWriter.WriteString(s) if n > 0 { - chunk := s[:n] - if cw.collector != nil { - cw.collector.Add([]byte(chunk)) - } - if cw.ch != nil { - buf := make([]byte, n) - copy(buf, chunk) - cw.ch <- buf - } + cw.collector.Add([]byte(s[:n])) } return n, err } -func startCaptureWriter(f *os.File, requestID string, start time.Time) (in chan<- []byte, done <-chan struct{}) { - chIn := make(chan []byte, 128) - doneCh := make(chan struct{}) - - go func() { - defer close(doneCh) - defer f.Close() - for chunk := range chIn { - _, _ = f.Write(chunk) - } - _, _ = fmt.Fprintf(f, "\n=== END duration_ms=%d ===\n", time.Since(start).Milliseconds()) - }() - - return chIn, doneCh -} - func normalizeCaptureRequestID(requestID string) string { if strings.TrimSpace(requestID) != "" { return requestID @@ -241,26 +208,19 @@ func buildCaptureResponseRecord(c *gin.Context, collector *responseCaptureCollec } } -func writeRequestBlock(c *gin.Context, f *os.File, userID int64) { - _, _ = fmt.Fprintf(f, "=== REQUEST %s ===\n", time.Now().UTC().Format(time.RFC3339Nano)) - _, _ = fmt.Fprintf(f, "%s %s\n", c.Request.Method, c.Request.URL.Path) - _, _ = fmt.Fprintf(f, "user_id: %d\n", userID) - - for key, values := range c.Request.Header { - for _, v := range values { - _, _ = fmt.Fprintf(f, "%s: %s\n", key, v) - } +func writeCaptureRecord(requestID string, record relayCaptureRecord) error { + f, err := createCaptureFile(requestID) + if err != nil { + return err } - _, _ = fmt.Fprintf(f, "\n") + defer f.Close() - if c.Request.Body != nil { - if storage, err := common.GetBodyStorage(c); err == nil { - if body, err := storage.Bytes(); err == nil && len(body) > 0 { - _, _ = f.Write(body) - } - } + data, err := common.Marshal(record) + if err != nil { + return err } - _, _ = fmt.Fprintf(f, "\n") + _, err = f.Write(data) + return err } func RelayCaptureMiddleware() gin.HandlerFunc { @@ -272,31 +232,30 @@ func RelayCaptureMiddleware() gin.HandlerFunc { } start := time.Now() - requestID := c.GetString(common.RequestIdKey) - - f, err := createCaptureFile(requestID) - if err != nil { - c.Next() - return - } - - writeRequestBlock(c, f, userID) - _, _ = fmt.Fprintf(f, "\n=== RESPONSE ===\n") - - chIn, doneCh := startCaptureWriter(f, requestID, start) + requestID := normalizeCaptureRequestID(c.GetString(common.RequestIdKey)) + captureErrors := make([]string, 0) + requestRecord := buildCaptureRequestRecord(c, &captureErrors) + collector := newResponseCaptureCollector() c.Writer = &captureResponseWriter{ ResponseWriter: c.Writer, - ch: chIn, + collector: collector, } c.Next() - close(chIn) - select { - case <-doneCh: - case <-time.After(30 * time.Second): - common.SysError("relay capture: timeout waiting for writer goroutine, request_id=" + requestID) + record := relayCaptureRecord{ + RequestID: requestID, + CapturedAt: start.UTC().Format(time.RFC3339Nano), + DurationMs: time.Since(start).Milliseconds(), + UserID: userID, + Request: requestRecord, + Response: buildCaptureResponseRecord(c, collector), + CaptureErrors: captureErrors, + } + + if err := writeCaptureRecord(requestID, record); err != nil { + common.SysError("relay capture: failed to write json capture, request_id=" + requestID + ", error=" + err.Error()) } } } diff --git a/middleware/relay_capture_integration_test.go b/middleware/relay_capture_integration_test.go index 52dadb7..bf6dc33 100644 --- a/middleware/relay_capture_integration_test.go +++ b/middleware/relay_capture_integration_test.go @@ -7,6 +7,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "strings" "sync/atomic" "testing" "time" @@ -22,21 +23,73 @@ func setupCaptureTestRouter() *gin.Engine { return gin.New() } -// captureTestCaptureDir 集成测试共享的捕获文件目录 const captureTestDir = "./data/relay-capture" -// --- 非流式请求测试 --- +type relayCaptureTestRecord struct { + RequestID string `json:"request_id"` + CapturedAt string `json:"captured_at"` + DurationMs int64 `json:"duration_ms"` + UserID int64 `json:"user_id"` + Request struct { + Method string `json:"method"` + Path string `json:"path"` + Query string `json:"query"` + Headers map[string][]string `json:"headers"` + Body string `json:"body"` + } `json:"request"` + Response struct { + StatusCode int `json:"status_code"` + Headers map[string][]string `json:"headers"` + Body string `json:"body"` + BodyChunks []string `json:"body_chunks"` + } `json:"response"` + CaptureErrors []string `json:"capture_errors"` +} + +func readCaptureRecord(t *testing.T, requestID string) relayCaptureTestRecord { + t.Helper() + today := time.Now().Format("2006-01-02") + path := filepath.Join(captureTestDir, today, requestID+".json") + content, err := os.ReadFile(path) + require.NoError(t, err) + require.NotContains(t, string(content), "=== REQUEST") + require.NotContains(t, string(content), "=== RESPONSE") + require.NotContains(t, string(content), "=== END") + + var record relayCaptureTestRecord + require.NoError(t, common.Unmarshal(content, &record)) + return record +} + +func readLatestCaptureRecord(t *testing.T) relayCaptureTestRecord { + t.Helper() + today := time.Now().Format("2006-01-02") + dateDir := filepath.Join(captureTestDir, today) + files, err := os.ReadDir(dateDir) + require.NoError(t, err) + require.NotEmpty(t, files, "should have capture json files") + latest := filepath.Join(dateDir, files[len(files)-1].Name()) + require.True(t, strings.HasSuffix(latest, ".json")) + + content, err := os.ReadFile(latest) + require.NoError(t, err) + var record relayCaptureTestRecord + require.NoError(t, common.Unmarshal(content, &record)) + return record +} func TestRelayCapture_NonStreaming_EnabledUser(t *testing.T) { clearCaptureCache() + t.Cleanup(func() { os.RemoveAll(captureTestDir) }) SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) + requestID := fmt.Sprintf("test-non-stream-%d", time.Now().UnixNano()) router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { - c.Set("id", 42) // must be int, not int64 — matches production TokenAuth behavior - c.Set(common.RequestIdKey, "test-non-stream-"+time.Now().Format("150405")) + c.Set("id", 42) + c.Set(common.RequestIdKey, requestID) c.Next() }) router.Use(RelayCaptureMiddleware()) @@ -58,31 +111,24 @@ func TestRelayCapture_NonStreaming_EnabledUser(t *testing.T) { assert.Equal(t, 200, rec.Code) - today := time.Now().Format("2006-01-02") - dateDir := filepath.Join(captureTestDir, today) - files, err := os.ReadDir(dateDir) - require.NoError(t, err) - require.True(t, len(files) > 0, "should have capture log files") - - logPath := filepath.Join(dateDir, files[len(files)-1].Name()) - content, err := os.ReadFile(logPath) - require.NoError(t, err) - contentStr := string(content) - - t.Logf("Capture file content:\n%s", contentStr) - - assert.Contains(t, contentStr, "=== REQUEST") - assert.Contains(t, contentStr, "POST /v1/chat/completions") - assert.Contains(t, contentStr, "user_id: 42") - assert.Contains(t, contentStr, "Authorization: Bearer sk-test") - assert.Contains(t, contentStr, "=== RESPONSE") - assert.Contains(t, contentStr, "chatcmpl-test") - assert.Contains(t, contentStr, "=== END") - assert.Contains(t, contentStr, "duration_ms=") + record := readCaptureRecord(t, requestID) + + assert.Equal(t, requestID, record.RequestID) + assert.NotEmpty(t, record.CapturedAt) + assert.GreaterOrEqual(t, record.DurationMs, int64(0)) + assert.Equal(t, int64(42), record.UserID) + assert.Equal(t, "POST", record.Request.Method) + assert.Equal(t, "/v1/chat/completions", record.Request.Path) + assert.Equal(t, body, record.Request.Body) + assert.Equal(t, []string{"[REDACTED]"}, record.Request.Headers["Authorization"]) + assert.Equal(t, []string{"application/json"}, record.Request.Headers["Content-Type"]) + assert.Equal(t, 200, record.Response.StatusCode) + assert.Contains(t, record.Response.Headers["Content-Type"][0], "application/json") + assert.Contains(t, record.Response.Body, "chatcmpl-test") + assert.NotEmpty(t, record.Response.BodyChunks) + assert.Empty(t, record.CaptureErrors) } -// --- 流式 SSE 请求测试 --- - func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { clearCaptureCache() t.Cleanup(func() { os.RemoveAll(captureTestDir) }) @@ -90,10 +136,11 @@ func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) + requestID := fmt.Sprintf("test-stream-%d", time.Now().UnixNano()) router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 42) - c.Set(common.RequestIdKey, "test-stream-"+time.Now().Format("150405")) + c.Set(common.RequestIdKey, requestID) c.Next() }) router.Use(RelayCaptureMiddleware()) @@ -111,7 +158,7 @@ func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { } for _, chunk := range chunks { - c.Writer.Write([]byte(chunk)) + _, _ = c.Writer.Write([]byte(chunk)) flusher.Flush() time.Sleep(5 * time.Millisecond) } @@ -127,39 +174,70 @@ func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { assert.Equal(t, 200, rec.Code) - today := time.Now().Format("2006-01-02") - dateDir := filepath.Join(captureTestDir, today) - files, err := os.ReadDir(dateDir) - require.NoError(t, err) - require.True(t, len(files) > 0, "should have capture log files") + record := readCaptureRecord(t, requestID) + + assert.Equal(t, requestID, record.RequestID) + assert.Equal(t, "POST", record.Request.Method) + assert.Equal(t, "/v1/chat/completions", record.Request.Path) + assert.Equal(t, body, record.Request.Body) + assert.Equal(t, 200, record.Response.StatusCode) + assert.Contains(t, record.Response.Body, "Hello ") + assert.Contains(t, record.Response.Body, "world!") + assert.Contains(t, record.Response.Body, "[DONE]") + require.Len(t, record.Response.BodyChunks, 3) + assert.Contains(t, record.Response.BodyChunks[0], "Hello ") + assert.Contains(t, record.Response.BodyChunks[1], "world!") + assert.Contains(t, record.Response.BodyChunks[2], "[DONE]") + assert.Empty(t, record.CaptureErrors) +} + +func TestRelayCapture_RedactsSensitiveRequestHeaders(t *testing.T) { + clearCaptureCache() + t.Cleanup(func() { os.RemoveAll(captureTestDir) }) + + SetCaptureEnabled(42, true) + defer SetCaptureEnabled(42, false) + + requestID := fmt.Sprintf("test-redaction-%d", time.Now().UnixNano()) + router := setupCaptureTestRouter() + router.Use(func(c *gin.Context) { + c.Set("id", 42) + c.Set(common.RequestIdKey, requestID) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + c.JSON(200, gin.H{"result": "ok"}) + }) + + req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4"}`))) + req.Header.Set("Authorization", "Bearer sk-test") + req.Header.Set("X-Api-Key", "x-secret") + req.Header.Set("Cookie", "session=secret") + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) - logPath := filepath.Join(dateDir, files[len(files)-1].Name()) - content, err := os.ReadFile(logPath) + assert.Equal(t, 200, rec.Code) + + today := time.Now().Format("2006-01-02") + path := filepath.Join(captureTestDir, today, requestID+".json") + content, err := os.ReadFile(path) require.NoError(t, err) contentStr := string(content) - - t.Logf("Stream capture file content:\n%s", contentStr) - - assert.Contains(t, contentStr, "=== REQUEST") - assert.Contains(t, contentStr, "POST /v1/chat/completions") - assert.Contains(t, contentStr, "user_id: 42") - assert.Contains(t, contentStr, "=== RESPONSE") - assert.Contains(t, contentStr, "Hello ") - assert.Contains(t, contentStr, "world!") - assert.Contains(t, contentStr, "[DONE]") - assert.Contains(t, contentStr, "=== END") - assert.Contains(t, contentStr, "duration_ms=") + assert.Contains(t, contentStr, "[REDACTED]") + assert.NotContains(t, contentStr, "Bearer sk-test") + assert.NotContains(t, contentStr, "x-secret") + assert.NotContains(t, contentStr, "session=secret") } -// --- 未开启抓包用户测试 --- - func TestRelayCapture_DisabledUser(t *testing.T) { clearCaptureCache() router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 99) - c.Set(common.RequestIdKey, "test-disabled-"+time.Now().Format("150405")) + c.Set(common.RequestIdKey, fmt.Sprintf("test-disabled-%d", time.Now().UnixNano())) c.Next() }) router.Use(RelayCaptureMiddleware()) @@ -185,8 +263,6 @@ func TestRelayCapture_DisabledUser(t *testing.T) { } } -// --- 大量数据流测试 --- - func TestRelayCapture_LargeStream(t *testing.T) { clearCaptureCache() t.Cleanup(func() { os.RemoveAll(captureTestDir) }) @@ -194,10 +270,11 @@ func TestRelayCapture_LargeStream(t *testing.T) { SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) + requestID := fmt.Sprintf("test-large-%d", time.Now().UnixNano()) router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 42) - c.Set(common.RequestIdKey, "test-large-"+time.Now().Format("150405")) + c.Set(common.RequestIdKey, requestID) c.Next() }) router.Use(RelayCaptureMiddleware()) @@ -207,10 +284,10 @@ func TestRelayCapture_LargeStream(t *testing.T) { for i := 0; i < 50; i++ { chunk := fmt.Sprintf("data: {\"id\":\"%d\",\"choices\":[{\"delta\":{\"content\":\"chunk%d \"}}]}\n\n", i, i) - c.Writer.Write([]byte(chunk)) + _, _ = c.Writer.Write([]byte(chunk)) flusher.Flush() } - c.Writer.Write([]byte("data: [DONE]\n\n")) + _, _ = c.Writer.Write([]byte("data: [DONE]\n\n")) flusher.Flush() }) @@ -222,24 +299,13 @@ func TestRelayCapture_LargeStream(t *testing.T) { assert.Equal(t, 200, rec.Code) - today := time.Now().Format("2006-01-02") - dateDir := filepath.Join(captureTestDir, today) - files, err := os.ReadDir(dateDir) - require.NoError(t, err) - require.True(t, len(files) > 0) - - logPath := filepath.Join(dateDir, files[len(files)-1].Name()) - content, err := os.ReadFile(logPath) - require.NoError(t, err) - contentStr := string(content) - - assert.Contains(t, contentStr, "chunk0") - assert.Contains(t, contentStr, "chunk49") - assert.Contains(t, contentStr, "[DONE]") + record := readCaptureRecord(t, requestID) + assert.Contains(t, record.Response.Body, "chunk0") + assert.Contains(t, record.Response.Body, "chunk49") + assert.Contains(t, record.Response.Body, "[DONE]") + assert.Len(t, record.Response.BodyChunks, 51) } -// --- GET 请求测试 --- - func TestRelayCapture_GetRequest(t *testing.T) { clearCaptureCache() t.Cleanup(func() { os.RemoveAll(captureTestDir) }) @@ -247,10 +313,11 @@ func TestRelayCapture_GetRequest(t *testing.T) { SetCaptureEnabled(42, true) defer SetCaptureEnabled(42, false) + requestID := fmt.Sprintf("test-get-%d", time.Now().UnixNano()) router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 42) - c.Set(common.RequestIdKey, "test-get-"+time.Now().Format("150405")) + c.Set(common.RequestIdKey, requestID) c.Next() }) router.Use(RelayCaptureMiddleware()) @@ -268,32 +335,20 @@ func TestRelayCapture_GetRequest(t *testing.T) { assert.Equal(t, 200, rec.Code) - today := time.Now().Format("2006-01-02") - dateDir := filepath.Join(captureTestDir, today) - files, err := os.ReadDir(dateDir) - require.NoError(t, err) - require.True(t, len(files) > 0) - - logPath := filepath.Join(dateDir, files[len(files)-1].Name()) - content, err := os.ReadFile(logPath) - require.NoError(t, err) - contentStr := string(content) - - assert.Contains(t, contentStr, "=== REQUEST") - assert.Contains(t, contentStr, "GET /v1/models") - assert.Contains(t, contentStr, "=== RESPONSE") - assert.Contains(t, contentStr, "gpt-4") - assert.Contains(t, contentStr, "=== END") + record := readCaptureRecord(t, requestID) + assert.Equal(t, "GET", record.Request.Method) + assert.Equal(t, "/v1/models", record.Request.Path) + assert.Equal(t, "", record.Request.Body) + assert.Equal(t, 200, record.Response.StatusCode) + assert.Contains(t, record.Response.Body, "gpt-4") } -// --- 零 ID 用户测试 --- - func TestRelayCapture_ZeroUserID(t *testing.T) { clearCaptureCache() router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { - c.Set(common.RequestIdKey, "test-zero-"+time.Now().Format("150405")) + c.Set(common.RequestIdKey, fmt.Sprintf("test-zero-%d", time.Now().UnixNano())) c.Next() }) router.Use(RelayCaptureMiddleware()) @@ -308,7 +363,42 @@ func TestRelayCapture_ZeroUserID(t *testing.T) { assert.Equal(t, 200, rec.Code) } -// --- 验证响应不被篡改 --- +func TestRelayCapture_EmptyRequestIDUsesFallbackFilename(t *testing.T) { + clearCaptureCache() + t.Cleanup(func() { os.RemoveAll(captureTestDir) }) + + SetCaptureEnabled(42, true) + defer SetCaptureEnabled(42, false) + + router := setupCaptureTestRouter() + router.Use(func(c *gin.Context) { + c.Set("id", 42) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.GET("/v1/models", func(c *gin.Context) { + c.JSON(200, gin.H{"data": []gin.H{}}) + }) + + req, _ := http.NewRequest("GET", "/v1/models", nil) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, 200, rec.Code) + + today := time.Now().Format("2006-01-02") + dateDir := filepath.Join(captureTestDir, today) + files, err := os.ReadDir(dateDir) + require.NoError(t, err) + require.Len(t, files, 1) + assert.True(t, strings.HasPrefix(files[0].Name(), "capture-")) + assert.True(t, strings.HasSuffix(files[0].Name(), ".json")) + + record := readLatestCaptureRecord(t) + assert.True(t, strings.HasPrefix(record.RequestID, "capture-")) + assert.Equal(t, "GET", record.Request.Method) + assert.Equal(t, "/v1/models", record.Request.Path) +} func TestRelayCapture_ResponseUnchanged(t *testing.T) { clearCaptureCache() @@ -320,7 +410,7 @@ func TestRelayCapture_ResponseUnchanged(t *testing.T) { router := setupCaptureTestRouter() router.Use(func(c *gin.Context) { c.Set("id", 42) - c.Set(common.RequestIdKey, "test-unchanged-"+time.Now().Format("150405")) + c.Set(common.RequestIdKey, fmt.Sprintf("test-unchanged-%d", time.Now().UnixNano())) c.Next() }) router.Use(RelayCaptureMiddleware()) @@ -342,8 +432,6 @@ func TestRelayCapture_ResponseUnchanged(t *testing.T) { assert.Contains(t, rec.Body.String(), "chatcmpl-123") } -// --- 并发安全测试 --- - func TestRelayCapture_ConcurrentRequests(t *testing.T) { clearCaptureCache() t.Cleanup(func() { os.RemoveAll(captureTestDir) }) @@ -352,7 +440,7 @@ func TestRelayCapture_ConcurrentRequests(t *testing.T) { defer SetCaptureEnabled(42, false) router := setupCaptureTestRouter() - var reqCounter int64 = 0 + var reqCounter int64 router.Use(func(c *gin.Context) { c.Set("id", 42) n := atomic.AddInt64(&reqCounter, 1) @@ -389,14 +477,17 @@ func TestRelayCapture_ConcurrentRequests(t *testing.T) { dateDir := filepath.Join(captureTestDir, today) files, err := os.ReadDir(dateDir) require.NoError(t, err) - assert.Equal(t, 10, len(files), "should have 10 capture log files") + assert.Equal(t, 10, len(files), "should have 10 capture json files") for _, f := range files { + require.True(t, strings.HasSuffix(f.Name(), ".json")) content, err := os.ReadFile(filepath.Join(dateDir, f.Name())) require.NoError(t, err) - contentStr := string(content) - assert.Contains(t, contentStr, "=== REQUEST") - assert.Contains(t, contentStr, "=== RESPONSE") - assert.Contains(t, contentStr, "=== END") + var record relayCaptureTestRecord + require.NoError(t, common.Unmarshal(content, &record)) + assert.Equal(t, "POST", record.Request.Method) + assert.Equal(t, "/v1/chat/completions", record.Request.Path) + assert.Equal(t, 200, record.Response.StatusCode) + assert.Contains(t, record.Response.Body, "ok") } }