diff --git a/middleware/relay_capture.go b/middleware/relay_capture.go index 1f07090..480958b 100644 --- a/middleware/relay_capture.go +++ b/middleware/relay_capture.go @@ -1,9 +1,12 @@ package middleware import ( + "bytes" "fmt" + "net/http" "os" "path/filepath" + "strings" "sync" "time" @@ -27,7 +30,6 @@ func SetCaptureEnabled(userID int64, enabled bool) { } } -// LoadCaptureEnabledUsers 从数据库加载 capture_relay=true 的用户到内存缓存 func LoadCaptureEnabledUsers() error { var ids []int64 if err := model.DB.Model(&model.User{}). @@ -41,7 +43,6 @@ func LoadCaptureEnabledUsers() error { return nil } -// GetCaptureEnabledUsersCount 仅用于测试 func GetCaptureEnabledUsersCount() int { count := 0 captureEnabledUsers.Range(func(key, value interface{}) bool { @@ -51,19 +52,80 @@ func GetCaptureEnabledUsersCount() int { return count } -// --- captureResponseWriter --- +type relayCaptureRecord struct { + RequestID string `json:"request_id"` + CapturedAt string `json:"captured_at"` + DurationMs int64 `json:"duration_ms"` + UserID int64 `json:"user_id"` + Request relayCaptureRequestRecord `json:"request"` + Response relayCaptureResponseRecord `json:"response"` + CaptureErrors []string `json:"capture_errors"` +} + +type relayCaptureRequestRecord struct { + Method string `json:"method"` + Path string `json:"path"` + Query string `json:"query"` + Headers map[string][]string `json:"headers"` + Body string `json:"body"` +} + +type relayCaptureResponseRecord struct { + StatusCode int `json:"status_code"` + Headers map[string][]string `json:"headers"` + Body string `json:"body"` + BodyChunks []string `json:"body_chunks"` +} + +type responseCaptureCollector struct { + body bytes.Buffer + chunks []string +} + +func newResponseCaptureCollector() *responseCaptureCollector { + return &responseCaptureCollector{ + chunks: make([]string, 0), + } +} + +func (collector *responseCaptureCollector) Add(chunk []byte) { + if len(chunk) == 0 { + return + } + copied := make([]byte, len(chunk)) + copy(copied, chunk) + collector.body.Write(copied) + collector.chunks = append(collector.chunks, string(copied)) +} + +func (collector *responseCaptureCollector) Body() string { + return collector.body.String() +} + +func (collector *responseCaptureCollector) Chunks() []string { + chunks := make([]string, len(collector.chunks)) + copy(chunks, collector.chunks) + return chunks +} type captureResponseWriter struct { gin.ResponseWriter - ch chan<- []byte + collector *responseCaptureCollector + ch chan<- []byte } func (cw *captureResponseWriter) Write(b []byte) (int, error) { n, err := cw.ResponseWriter.Write(b) if n > 0 { - buf := make([]byte, n) - copy(buf, b[:n]) - cw.ch <- buf + 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 + } } return n, err } @@ -71,16 +133,19 @@ 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 { - buf := make([]byte, n) - copy(buf, s[:n]) - cw.ch <- buf + 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 + } } return n, err } -// --- startCaptureWriter --- - -// startCaptureWriter 启动单个 writer goroutine,从 chIn 顺序写入文件 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{}) @@ -89,15 +154,20 @@ func startCaptureWriter(f *os.File, requestID string, start time.Time) (in chan< defer close(doneCh) defer f.Close() for chunk := range chIn { - f.Write(chunk) + _, _ = f.Write(chunk) } - fmt.Fprintf(f, "\n=== END duration_ms=%d ===\n", time.Since(start).Milliseconds()) + _, _ = 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 + } + return fmt.Sprintf("capture-%d", time.Now().UnixNano()) +} func createCaptureFile(requestID string) (*os.File, error) { dir := filepath.Join("./data", "relay-capture", time.Now().Format("2006-01-02")) @@ -105,36 +175,94 @@ func createCaptureFile(requestID string) (*os.File, error) { return nil, err } return os.OpenFile( - filepath.Join(dir, requestID+".log"), + filepath.Join(dir, normalizeCaptureRequestID(requestID)+".json"), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644, ) } +func isSensitiveCaptureHeader(key string) bool { + switch strings.ToLower(key) { + case "authorization", "proxy-authorization", "x-api-key", "api-key", "cookie": + return true + default: + return false + } +} + +func copyHeadersForCapture(headers http.Header, redactSensitive bool) map[string][]string { + copied := make(map[string][]string, len(headers)) + for key, values := range headers { + copiedValues := make([]string, len(values)) + if redactSensitive && isSensitiveCaptureHeader(key) { + for i := range copiedValues { + copiedValues[i] = "[REDACTED]" + } + } else { + copy(copiedValues, values) + } + copied[key] = copiedValues + } + return copied +} + +func buildCaptureRequestRecord(c *gin.Context, captureErrors *[]string) relayCaptureRequestRecord { + record := relayCaptureRequestRecord{ + Method: c.Request.Method, + Path: c.Request.URL.Path, + Query: c.Request.URL.RawQuery, + Headers: copyHeadersForCapture(c.Request.Header, true), + Body: "", + } + + if c.Request.Body == nil { + return record + } + + storage, err := common.GetBodyStorage(c) + if err != nil { + *captureErrors = append(*captureErrors, "failed to read request body: "+err.Error()) + return record + } + body, err := storage.Bytes() + if err != nil { + *captureErrors = append(*captureErrors, "failed to read request body bytes: "+err.Error()) + return record + } + record.Body = string(body) + return record +} + +func buildCaptureResponseRecord(c *gin.Context, collector *responseCaptureCollector) relayCaptureResponseRecord { + return relayCaptureResponseRecord{ + StatusCode: c.Writer.Status(), + Headers: copyHeadersForCapture(c.Writer.Header(), false), + Body: collector.Body(), + BodyChunks: collector.Chunks(), + } +} + 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) + _, _ = 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) + _, _ = fmt.Fprintf(f, "%s: %s\n", key, v) } } - fmt.Fprintf(f, "\n") + _, _ = fmt.Fprintf(f, "\n") - // 读取请求体 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) + _, _ = f.Write(body) } } } - fmt.Fprintf(f, "\n") + _, _ = fmt.Fprintf(f, "\n") } -// RelayCaptureMiddleware 抓包中间件 -// 对开启 capture_relay 的用户,将请求体和响应体写入本地文件 func RelayCaptureMiddleware() gin.HandlerFunc { return func(c *gin.Context) { userID := int64(c.GetInt("id")) @@ -152,11 +280,9 @@ func RelayCaptureMiddleware() gin.HandlerFunc { return } - // 同步写 REQUEST 块 writeRequestBlock(c, f, userID) - fmt.Fprintf(f, "\n=== RESPONSE ===\n") + _, _ = fmt.Fprintf(f, "\n=== RESPONSE ===\n") - // 启动 writer goroutine chIn, doneCh := startCaptureWriter(f, requestID, start) c.Writer = &captureResponseWriter{ @@ -166,7 +292,6 @@ func RelayCaptureMiddleware() gin.HandlerFunc { c.Next() - // 关闭 chIn,等待 writer 完成(带超时保护) close(chIn) select { case <-doneCh: diff --git a/middleware/relay_capture_test.go b/middleware/relay_capture_test.go index 8fad565..62d01d0 100644 --- a/middleware/relay_capture_test.go +++ b/middleware/relay_capture_test.go @@ -1,13 +1,13 @@ package middleware import ( - "fmt" - "io" + "bytes" + "net/http" "net/http/httptest" "os" "path/filepath" + "strings" "testing" - "time" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" @@ -67,17 +67,15 @@ func TestGetCaptureEnabledUsersCount(t *testing.T) { SetCaptureEnabled(2, false) } -// --- captureResponseWriter 测试 --- - func TestCaptureResponseWriter_Write(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) - ch := make(chan []byte, 16) + collector := newResponseCaptureCollector() cw := &captureResponseWriter{ ResponseWriter: c.Writer, - ch: ch, + collector: collector, } n, err := cw.Write([]byte("hello")) @@ -85,23 +83,19 @@ func TestCaptureResponseWriter_Write(t *testing.T) { assert.Equal(t, 5, n) assert.Equal(t, "hello", rec.Body.String()) - select { - case data := <-ch: - assert.Equal(t, []byte("hello"), data) - default: - t.Fatal("expected data in channel") - } + assert.Equal(t, "hello", collector.Body()) + assert.Equal(t, []string{"hello"}, collector.Chunks()) } func TestCaptureResponseWriter_WriteString(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) - ch := make(chan []byte, 16) + collector := newResponseCaptureCollector() cw := &captureResponseWriter{ ResponseWriter: c.Writer, - ch: ch, + collector: collector, } n, err := cw.WriteString("test-string") @@ -109,86 +103,135 @@ func TestCaptureResponseWriter_WriteString(t *testing.T) { assert.Equal(t, len("test-string"), n) assert.Equal(t, "test-string", rec.Body.String()) - select { - case data := <-ch: - assert.Equal(t, []byte("test-string"), data) - default: - t.Fatal("expected data in channel") - } + assert.Equal(t, "test-string", collector.Body()) + assert.Equal(t, []string{"test-string"}, collector.Chunks()) } -// --- startCaptureWriter 测试 --- - -func TestStartCaptureWriter_WritesAllChunks(t *testing.T) { - f, err := os.CreateTemp("", "test-capture-*.log") - require.NoError(t, err) - filePath := f.Name() - defer os.Remove(filePath) +type partialResponseWriter struct { + gin.ResponseWriter +} - start := time.Now() - chIn, doneCh := startCaptureWriter(f, "test-req", start) +func (w partialResponseWriter) Write(b []byte) (int, error) { + _, err := w.ResponseWriter.Write(b[:2]) + return 2, err +} - // 发送 130 个 chunk(大于 channel buffer 128) - for i := 0; i < 130; i++ { - chIn <- []byte(fmt.Sprintf("data: %d\n", i)) - } - close(chIn) +func TestCaptureResponseWriter_RecordsOnlyWrittenBytes(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + collector := newResponseCaptureCollector() - select { - case <-doneCh: - case <-time.After(5 * time.Second): - t.Fatal("timeout waiting for doneCh") + cw := &captureResponseWriter{ + ResponseWriter: partialResponseWriter{ResponseWriter: c.Writer}, + collector: collector, } - content, err := os.ReadFile(filePath) - require.NoError(t, err) - contentStr := string(content) - assert.Contains(t, contentStr, "data: 0") - assert.Contains(t, contentStr, "data: 19") - assert.Contains(t, contentStr, "=== END") - assert.Contains(t, contentStr, "duration_ms=") + n, err := cw.Write([]byte("hello")) + assert.NoError(t, err) + assert.Equal(t, 2, n) + assert.Equal(t, "he", rec.Body.String()) + assert.Equal(t, "he", collector.Body()) + assert.Equal(t, []string{"he"}, collector.Chunks()) } -// --- 辅助函数测试 --- - func TestCreateCaptureFile(t *testing.T) { f, err := createCaptureFile("test-req-123") require.NoError(t, err) defer os.RemoveAll(filepath.Dir(f.Name())) + defer f.Close() assert.Contains(t, f.Name(), "relay-capture") - assert.Contains(t, f.Name(), "test-req-123.log") + assert.Contains(t, f.Name(), "test-req-123.json") + assert.NotContains(t, f.Name(), ".log") dir := filepath.Dir(f.Name()) info, err := os.Stat(dir) require.NoError(t, err) assert.True(t, info.IsDir()) +} + +func TestNormalizeCaptureRequestID_UsesExistingID(t *testing.T) { + requestID := normalizeCaptureRequestID("req-123") + assert.Equal(t, "req-123", requestID) +} - f.Close() +func TestNormalizeCaptureRequestID_FallbackForEmptyID(t *testing.T) { + requestID := normalizeCaptureRequestID("") + assert.True(t, strings.HasPrefix(requestID, "capture-")) + assert.False(t, strings.HasSuffix(requestID, ".json")) } -func TestWriteRequestBlock(t *testing.T) { +func TestCopyHeadersForCapture_RedactsSensitiveRequestHeaders(t *testing.T) { + headers := http.Header{ + "Authorization": []string{"Bearer sk-test"}, + "authorization": []string{"Bearer lower"}, + "X-Api-Key": []string{"x-key"}, + "Api-Key": []string{"api-key"}, + "Cookie": []string{"session=secret"}, + "Proxy-Authorization": []string{"Basic secret"}, + "Content-Type": []string{"application/json"}, + "X-Trace-ID": []string{"trace-1"}, + } + + copied := copyHeadersForCapture(headers, true) + + assert.Equal(t, []string{"[REDACTED]"}, copied["Authorization"]) + assert.Equal(t, []string{"[REDACTED]"}, copied["authorization"]) + assert.Equal(t, []string{"[REDACTED]"}, copied["X-Api-Key"]) + assert.Equal(t, []string{"[REDACTED]"}, copied["Api-Key"]) + assert.Equal(t, []string{"[REDACTED]"}, copied["Cookie"]) + assert.Equal(t, []string{"[REDACTED]"}, copied["Proxy-Authorization"]) + assert.Equal(t, []string{"application/json"}, copied["Content-Type"]) + assert.Equal(t, []string{"trace-1"}, copied["X-Trace-ID"]) + + assert.Equal(t, []string{"Bearer sk-test"}, headers["Authorization"]) +} + +func TestCopyHeadersForCapture_ResponseHeadersNotRedacted(t *testing.T) { + headers := http.Header{} + headers.Add("Authorization", "response-token") + headers.Add("Content-Type", "application/json") + + copied := copyHeadersForCapture(headers, false) + + assert.Equal(t, []string{"response-token"}, copied["Authorization"]) + assert.Equal(t, []string{"application/json"}, copied["Content-Type"]) +} + +func TestBuildCaptureRequestRecord_PostBody(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) - - c.Set("id", int64(42)) - c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil) + body := `{"model":"gpt-4"}` + c.Request = httptest.NewRequest("POST", "/v1/chat/completions?stream=true", bytes.NewReader([]byte(body))) c.Request.Header.Set("Authorization", "Bearer sk-test") c.Request.Header.Set("Content-Type", "application/json") - f, err := os.CreateTemp("", "test-req-block-*.log") - require.NoError(t, err) - filePath := f.Name() - defer os.Remove(filePath) + var captureErrors []string + record := buildCaptureRequestRecord(c, &captureErrors) - writeRequestBlock(c, f, 42) + assert.Equal(t, "POST", record.Method) + assert.Equal(t, "/v1/chat/completions", record.Path) + assert.Equal(t, "stream=true", record.Query) + assert.Equal(t, body, record.Body) + assert.Equal(t, []string{"[REDACTED]"}, record.Headers["Authorization"]) + assert.Equal(t, []string{"application/json"}, record.Headers["Content-Type"]) + assert.Empty(t, captureErrors) +} - f.Seek(0, io.SeekStart) - content, err := io.ReadAll(f) - require.NoError(t, err) - contentStr := string(content) - assert.Contains(t, contentStr, "=== REQUEST") - assert.Contains(t, contentStr, "POST /v1/chat/completions") - assert.Contains(t, contentStr, "user_id: 42") +func TestBuildCaptureRequestRecord_GetNoBody(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/v1/models", nil) + + var captureErrors []string + record := buildCaptureRequestRecord(c, &captureErrors) + + assert.Equal(t, "GET", record.Method) + assert.Equal(t, "/v1/models", record.Path) + assert.Equal(t, "", record.Query) + assert.Equal(t, "", record.Body) + assert.Empty(t, captureErrors) }