|
- package middleware
-
- import (
- "fmt"
- "os"
- "path/filepath"
- "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)
- }
- }
-
- // LoadCaptureEnabledUsers 从数据库加载 capture_relay=true 的用户到内存缓存
- 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
- }
-
- // GetCaptureEnabledUsersCount 仅用于测试
- func GetCaptureEnabledUsersCount() int {
- count := 0
- captureEnabledUsers.Range(func(key, value interface{}) bool {
- count++
- return true
- })
- return count
- }
-
- // --- captureResponseWriter ---
-
- type captureResponseWriter struct {
- gin.ResponseWriter
- 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
- }
- return n, err
- }
-
- 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
- }
- 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{})
-
- 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 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, requestID+".log"),
- os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644,
- )
- }
-
- 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")
- }
-
- // RelayCaptureMiddleware 抓包中间件
- // 对开启 capture_relay 的用户,将请求体和响应体写入本地文件
- 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
- }
-
- // 同步写 REQUEST 块
- writeRequestBlock(c, f, userID)
- fmt.Fprintf(f, "\n=== RESPONSE ===\n")
-
- // 启动 writer goroutine
- chIn, doneCh := startCaptureWriter(f, requestID, start)
-
- c.Writer = &captureResponseWriter{
- ResponseWriter: c.Writer,
- ch: chIn,
- }
-
- c.Next()
-
- // 关闭 chIn,等待 writer 完成(带超时保护)
- close(chIn)
- select {
- case <-doneCh:
- case <-time.After(30 * time.Second):
- common.SysError("relay capture: timeout waiting for writer goroutine, request_id=" + requestID)
- }
- }
- }
|