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
	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
		}
	}
	return n, err
}

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
		}
	}
	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
	}
	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 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)
		}
	}
	_, _ = 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)
			}
		}
	}
	_, _ = fmt.Fprintf(f, "\n")
}

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

		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)

		c.Writer = &captureResponseWriter{
			ResponseWriter: c.Writer,
			ch:             chIn,
		}

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