98 lines
2.1 KiB
Go
98 lines
2.1 KiB
Go
|
package middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"net/http"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/sirupsen/logrus"
|
||
|
|
|
||
|
|
"ca-mini/internal/config"
|
||
|
|
"ca-mini/internal/logger"
|
||
|
|
)
|
||
|
|
|
||
|
|
type RequestInfo struct {
|
||
|
|
Method string `json:"method"`
|
||
|
|
Path string `json:"path"`
|
||
|
|
Header http.Header `json:"header"`
|
||
|
|
Body string `json:"body"`
|
||
|
|
}
|
||
|
|
|
||
|
|
type ResponseInfo struct {
|
||
|
|
Code int `json:"code"`
|
||
|
|
Header http.Header `json:"header"`
|
||
|
|
Body string `json:"body"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// responseWriter 包装 http.ResponseWriter 以捕获状态码和响应体
|
||
|
|
type responseWriter struct {
|
||
|
|
http.ResponseWriter
|
||
|
|
statusCode int
|
||
|
|
body *bytes.Buffer
|
||
|
|
}
|
||
|
|
|
||
|
|
// WriteHeader 捕获状态码
|
||
|
|
func (rw *responseWriter) WriteHeader(code int) {
|
||
|
|
rw.statusCode = code
|
||
|
|
rw.ResponseWriter.WriteHeader(code)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Write 捕获响应体
|
||
|
|
func (rw *responseWriter) Write(b []byte) (int, error) {
|
||
|
|
rw.body.Write(b)
|
||
|
|
return rw.ResponseWriter.Write(b)
|
||
|
|
}
|
||
|
|
|
||
|
|
// LoggerMiddleware 日志中间件
|
||
|
|
func LoggerMiddleware(next http.Handler) http.Handler {
|
||
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
cfg, err := config.Load()
|
||
|
|
if err != nil || cfg.Logging.Level == "debug" {
|
||
|
|
next.ServeHTTP(w, r)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// 开始计时
|
||
|
|
start := time.Now()
|
||
|
|
|
||
|
|
// 读取请求体并重置 r.Body
|
||
|
|
var requestBody bytes.Buffer
|
||
|
|
_, err = io.Copy(&requestBody, r.Body)
|
||
|
|
if err != nil {
|
||
|
|
logger.Error("Failed to read request body: %v", err)
|
||
|
|
}
|
||
|
|
r.Body = io.NopCloser(&requestBody)
|
||
|
|
|
||
|
|
// 创建一个响应包装器来捕获响应状态码和响应体
|
||
|
|
lrw := &responseWriter{
|
||
|
|
ResponseWriter: w,
|
||
|
|
statusCode: http.StatusOK,
|
||
|
|
body: &bytes.Buffer{},
|
||
|
|
}
|
||
|
|
|
||
|
|
// 处理请求
|
||
|
|
next.ServeHTTP(lrw, r)
|
||
|
|
|
||
|
|
// 统计请求处理时间
|
||
|
|
duration := time.Since(start).Milliseconds()
|
||
|
|
|
||
|
|
// 记录日志
|
||
|
|
logger.Logger.WithFields(logrus.Fields{
|
||
|
|
"request": RequestInfo{
|
||
|
|
Method: r.Method,
|
||
|
|
Path: r.URL.Path,
|
||
|
|
Header: r.Header,
|
||
|
|
Body: requestBody.String(),
|
||
|
|
},
|
||
|
|
"response": ResponseInfo{
|
||
|
|
Code: lrw.statusCode,
|
||
|
|
Header: lrw.Header(),
|
||
|
|
Body: lrw.body.String(),
|
||
|
|
},
|
||
|
|
"duration": fmt.Sprintf("%dms", duration),
|
||
|
|
}).Debug("Request processed")
|
||
|
|
})
|
||
|
|
}
|