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