76 lines
1.7 KiB
Go
76 lines
1.7 KiB
Go
|
package middleware
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"bytes"
|
|||
|
|
"io"
|
|||
|
|
"strings"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"gitea.com/bitwsd/document_ai/pkg/log"
|
|||
|
|
"github.com/gin-gonic/gin"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
maxBodySize = 1024 * 500 // 500KB 限制
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// 自定义 ResponseWriter 来捕获响应
|
|||
|
|
type bodyWriter struct {
|
|||
|
|
gin.ResponseWriter
|
|||
|
|
body *bytes.Buffer
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (w *bodyWriter) Write(b []byte) (int, error) {
|
|||
|
|
w.body.Write(b)
|
|||
|
|
return w.ResponseWriter.Write(b)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func AccessLog() gin.HandlerFunc {
|
|||
|
|
return func(c *gin.Context) {
|
|||
|
|
start := time.Now()
|
|||
|
|
path := c.Request.URL.Path
|
|||
|
|
raw := c.Request.URL.RawQuery
|
|||
|
|
|
|||
|
|
// 处理请求体
|
|||
|
|
var reqBody string
|
|||
|
|
if c.Request.Body != nil && (c.Request.Method == "POST" || c.Request.Method == "PUT") {
|
|||
|
|
// 读取并限制请求体大小
|
|||
|
|
bodyBytes, _ := io.ReadAll(io.LimitReader(c.Request.Body, maxBodySize))
|
|||
|
|
c.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) // 重新设置 body
|
|||
|
|
reqBody = string(bodyBytes)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 设置自定义 ResponseWriter
|
|||
|
|
var responseBody string
|
|||
|
|
if !strings.Contains(c.GetHeader("Accept"), "text/event-stream") {
|
|||
|
|
bw := &bodyWriter{body: &bytes.Buffer{}, ResponseWriter: c.Writer}
|
|||
|
|
c.Writer = bw
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
c.Next()
|
|||
|
|
|
|||
|
|
// 获取响应体（非 SSE）
|
|||
|
|
if writer, ok := c.Writer.(*bodyWriter); ok {
|
|||
|
|
responseBody = writer.body.String()
|
|||
|
|
if len(responseBody) > maxBodySize {
|
|||
|
|
responseBody = responseBody[:maxBodySize] + "... (truncated)"
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 记录访问日志
|
|||
|
|
log.Access(c.Request.Context(),
|
|||
|
|
"request_id", c.GetString("request_id"),
|
|||
|
|
"method", c.Request.Method,
|
|||
|
|
"path", path,
|
|||
|
|
"query", raw,
|
|||
|
|
"ip", c.ClientIP(),
|
|||
|
|
"user_agent", c.Request.UserAgent(),
|
|||
|
|
"status", c.Writer.Status(),
|
|||
|
|
"duration", time.Since(start),
|
|||
|
|
"request_body", reqBody,
|
|||
|
|
"response_body", responseBody,
|
|||
|
|
)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|