Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

178 строки
4.0 KiB

  1. package middleware
  2. import (
  3. "fmt"
  4. "os"
  5. "path/filepath"
  6. "sync"
  7. "time"
  8. "github.com/QuantumNous/new-api/common"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/gin-gonic/gin"
  11. )
  12. var captureEnabledUsers sync.Map // key: int64, value: struct{}
  13. func isCaptureEnabled(userID int64) bool {
  14. _, ok := captureEnabledUsers.Load(userID)
  15. return ok
  16. }
  17. func SetCaptureEnabled(userID int64, enabled bool) {
  18. if enabled {
  19. captureEnabledUsers.Store(userID, struct{}{})
  20. } else {
  21. captureEnabledUsers.Delete(userID)
  22. }
  23. }
  24. // LoadCaptureEnabledUsers 从数据库加载 capture_relay=true 的用户到内存缓存
  25. func LoadCaptureEnabledUsers() error {
  26. var ids []int64
  27. if err := model.DB.Model(&model.User{}).
  28. Where("capture_relay = ?", true).
  29. Pluck("id", &ids).Error; err != nil {
  30. return err
  31. }
  32. for _, id := range ids {
  33. captureEnabledUsers.Store(id, struct{}{})
  34. }
  35. return nil
  36. }
  37. // GetCaptureEnabledUsersCount 仅用于测试
  38. func GetCaptureEnabledUsersCount() int {
  39. count := 0
  40. captureEnabledUsers.Range(func(key, value interface{}) bool {
  41. count++
  42. return true
  43. })
  44. return count
  45. }
  46. // --- captureResponseWriter ---
  47. type captureResponseWriter struct {
  48. gin.ResponseWriter
  49. ch chan<- []byte
  50. }
  51. func (cw *captureResponseWriter) Write(b []byte) (int, error) {
  52. n, err := cw.ResponseWriter.Write(b)
  53. if n > 0 {
  54. buf := make([]byte, n)
  55. copy(buf, b[:n])
  56. cw.ch <- buf
  57. }
  58. return n, err
  59. }
  60. func (cw *captureResponseWriter) WriteString(s string) (int, error) {
  61. n, err := cw.ResponseWriter.WriteString(s)
  62. if n > 0 {
  63. buf := make([]byte, n)
  64. copy(buf, s[:n])
  65. cw.ch <- buf
  66. }
  67. return n, err
  68. }
  69. // --- startCaptureWriter ---
  70. // startCaptureWriter 启动单个 writer goroutine,从 chIn 顺序写入文件
  71. func startCaptureWriter(f *os.File, requestID string, start time.Time) (in chan<- []byte, done <-chan struct{}) {
  72. chIn := make(chan []byte, 128)
  73. doneCh := make(chan struct{})
  74. go func() {
  75. defer close(doneCh)
  76. defer f.Close()
  77. for chunk := range chIn {
  78. f.Write(chunk)
  79. }
  80. fmt.Fprintf(f, "\n=== END duration_ms=%d ===\n", time.Since(start).Milliseconds())
  81. }()
  82. return chIn, doneCh
  83. }
  84. // --- 辅助函数 ---
  85. func createCaptureFile(requestID string) (*os.File, error) {
  86. dir := filepath.Join("./data", "relay-capture", time.Now().Format("2006-01-02"))
  87. if err := os.MkdirAll(dir, 0755); err != nil {
  88. return nil, err
  89. }
  90. return os.OpenFile(
  91. filepath.Join(dir, requestID+".log"),
  92. os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644,
  93. )
  94. }
  95. func writeRequestBlock(c *gin.Context, f *os.File, userID int64) {
  96. fmt.Fprintf(f, "=== REQUEST %s ===\n", time.Now().UTC().Format(time.RFC3339Nano))
  97. fmt.Fprintf(f, "%s %s\n", c.Request.Method, c.Request.URL.Path)
  98. fmt.Fprintf(f, "user_id: %d\n", userID)
  99. for key, values := range c.Request.Header {
  100. for _, v := range values {
  101. fmt.Fprintf(f, "%s: %s\n", key, v)
  102. }
  103. }
  104. fmt.Fprintf(f, "\n")
  105. // 读取请求体
  106. if c.Request.Body != nil {
  107. if storage, err := common.GetBodyStorage(c); err == nil {
  108. if body, err := storage.Bytes(); err == nil && len(body) > 0 {
  109. f.Write(body)
  110. }
  111. }
  112. }
  113. fmt.Fprintf(f, "\n")
  114. }
  115. // RelayCaptureMiddleware 抓包中间件
  116. // 对开启 capture_relay 的用户,将请求体和响应体写入本地文件
  117. func RelayCaptureMiddleware() gin.HandlerFunc {
  118. return func(c *gin.Context) {
  119. userID := int64(c.GetInt("id"))
  120. if !isCaptureEnabled(userID) {
  121. c.Next()
  122. return
  123. }
  124. start := time.Now()
  125. requestID := c.GetString(common.RequestIdKey)
  126. f, err := createCaptureFile(requestID)
  127. if err != nil {
  128. c.Next()
  129. return
  130. }
  131. // 同步写 REQUEST 块
  132. writeRequestBlock(c, f, userID)
  133. fmt.Fprintf(f, "\n=== RESPONSE ===\n")
  134. // 启动 writer goroutine
  135. chIn, doneCh := startCaptureWriter(f, requestID, start)
  136. c.Writer = &captureResponseWriter{
  137. ResponseWriter: c.Writer,
  138. ch: chIn,
  139. }
  140. c.Next()
  141. // 关闭 chIn,等待 writer 完成(带超时保护)
  142. close(chIn)
  143. select {
  144. case <-doneCh:
  145. case <-time.After(30 * time.Second):
  146. common.SysError("relay capture: timeout waiting for writer goroutine, request_id=" + requestID)
  147. }
  148. }
  149. }