new-api/middleware/relay_capture.go

package middleware

import (
	"bytes"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/gin-gonic/gin"
)

var captureEnabledUsers sync.Map // key: int64, value: struct{}

func isCaptureEnabled(userID int64) bool {
	_, ok := captureEnabledUsers.Load(userID)
	return ok
}

func SetCaptureEnabled(userID int64, enabled bool) {
	if enabled {
		captureEnabledUsers.Store(userID, struct{}{})
	} else {
		captureEnabledUsers.Delete(userID)
	}
}

func LoadCaptureEnabledUsers() error {
	var ids []int64
	if err := model.DB.Model(&model.User{}).
		Where("capture_relay = ?", true).
		Pluck("id", &ids).Error; err != nil {
		return err
	}
	for _, id := range ids {
		captureEnabledUsers.Store(id, struct{}{})
	}
	return nil
}

func GetCaptureEnabledUsersCount() int {
	count := 0
	captureEnabledUsers.Range(func(key, value interface{}) bool {
		count++
		return true
	})
	return count
}

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
	collector *responseCaptureCollector
}

func (cw *captureResponseWriter) Write(b []byte) (int, error) {
	n, err := cw.ResponseWriter.Write(b)
	if n > 0 {
		cw.collector.Add(b[:n])
	}
	return n, err
}

func (cw *captureResponseWriter) WriteString(s string) (int, error) {
	n, err := cw.ResponseWriter.WriteString(s)
	if n > 0 {
		cw.collector.Add([]byte(s[:n]))
	}
	return n, err
}

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"))
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}
	return os.OpenFile(
		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 writeCaptureRecord(requestID string, record relayCaptureRecord) error {
	f, err := createCaptureFile(requestID)
	if err != nil {
		return err
	}
	defer f.Close()

	data, err := common.Marshal(record)
	if err != nil {
		return err
	}
	_, err = f.Write(data)
	return err
}

func RelayCaptureMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		userID := int64(c.GetInt("id"))
		if !isCaptureEnabled(userID) {
			c.Next()
			return
		}

		start := time.Now()
		requestID := normalizeCaptureRequestID(c.GetString(common.RequestIdKey))
		captureErrors := make([]string, 0)

		requestRecord := buildCaptureRequestRecord(c, &captureErrors)
		collector := newResponseCaptureCollector()
		c.Writer = &captureResponseWriter{
			ResponseWriter: c.Writer,
			collector:      collector,
		}

		c.Next()

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