refactor: prepare relay capture json helpers

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:

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)
 }