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