🎨 添加中间件 && 完善预设注入功能 && 新增流式传输

server/middleware/operation.go

@@ -0,0 +1,119 @@
+package middleware
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"time"
+
+	"git.echol.cn/loser/ai_proxy/server/global"
+	"git.echol.cn/loser/ai_proxy/server/model/system"
+	"git.echol.cn/loser/ai_proxy/server/utils"
+	"github.com/gin-gonic/gin"
+	"go.uber.org/zap"
+)
+
+var respPool sync.Pool
+var bufferSize = 1024
+
+func init() {
+	respPool.New = func() interface{} {
+		return make([]byte, bufferSize)
+	}
+}
+
+// Operation 操作日志记录中间件
+func Operation() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		var body []byte
+		var userId int
+
+		if c.Request.Method != http.MethodGet {
+			var err error
+			body, err = io.ReadAll(c.Request.Body)
+			if err != nil {
+				global.GVA_LOG.Error("read body from request error:", zap.Error(err))
+			} else {
+				c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
+			}
+		}
+
+		userId = int(utils.GetUserID(c))
+
+		writer := responseBodyWriter{
+			ResponseWriter: c.Writer,
+			body:           &bytes.Buffer{},
+		}
+		c.Writer = writer
+		now := time.Now()
+
+		c.Next()
+
+		latency := time.Since(now)
+
+		// 只记录需要的路径
+		if NeedRecordPath(c.Request.URL.Path) && c.Request.Method != http.MethodGet {
+			record := system.SysOperationRecord{
+				Ip:      c.ClientIP(),
+				Method:  c.Request.Method,
+				Path:    c.Request.URL.Path,
+				Agent:   c.Request.UserAgent(),
+				Body:    string(body),
+				UserID:  userId,
+				Status:  c.Writer.Status(),
+				Latency: latency,
+				Resp:    writer.body.String(),
+			}
+
+			values, _ := url.ParseQuery(c.Request.URL.RawQuery)
+			record.Query = values.Encode()
+
+			// 异步记录操作日志
+			go func() {
+				if err := operationRecordService.CreateSysOperationRecord(record); err != nil {
+					global.GVA_LOG.Error("create operation record error:", zap.Error(err))
+				}
+			}()
+		}
+	}
+}
+
+// FilteredPaths 需要过滤的路径
+var FilteredPaths = []string{
+	"/health",
+	"/swagger",
+	"/api/captcha",
+	"/api/base/login",
+}
+
+// ShouldRecord 判断是否应该记录该路径
+func ShouldRecord(path string) bool {
+	for _, filtered := range FilteredPaths {
+		if strings.HasPrefix(path, filtered) {
+			return false
+		}
+	}
+	return true
+}
+
+// MaskSensitiveData 脱敏敏感数据
+func MaskSensitiveData(data string) string {
+	var result map[string]interface{}
+	if err := json.Unmarshal([]byte(data), &result); err != nil {
+		return data
+	}
+
+	sensitiveFields := []string{"password", "token", "secret", "key"}
+	for _, field := range sensitiveFields {
+		if _, ok := result[field]; ok {
+			result[field] = "******"
+		}
+	}
+
+	masked, _ := json.Marshal(result)
+	return string(masked)
+}