103 lines
3.1 KiB
Go
103 lines
3.1 KiB
Go
package middleware
|
||
|
||
import (
|
||
"bytes"
|
||
"io"
|
||
"strings"
|
||
"time"
|
||
|
||
"modelRT/constants"
|
||
"modelRT/logger"
|
||
"modelRT/util"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
// StartTrace define func of set trace info from request header
|
||
func StartTrace() gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
traceID := c.Request.Header.Get(constants.HeaderTraceID)
|
||
parentSpanID := c.Request.Header.Get(constants.HeaderSpanID)
|
||
spanID := util.GenerateSpanID(c.Request.RemoteAddr)
|
||
// if traceId is empty, it means it is the origin of the link. Set it to the spanId of this time. The originating spanId is the root spanId.
|
||
if traceID == "" {
|
||
// traceId identifies the entire request link, and spanId identifies the different services in the link.
|
||
traceID = spanID
|
||
}
|
||
c.Set(constants.HeaderTraceID, traceID)
|
||
c.Set(constants.HeaderSpanID, spanID)
|
||
c.Set(constants.HeaderParentSpanID, parentSpanID)
|
||
c.Next()
|
||
}
|
||
}
|
||
|
||
type bodyLogWriter struct {
|
||
gin.ResponseWriter
|
||
body *bytes.Buffer
|
||
}
|
||
|
||
// TODO 包装一下 gin.ResponseWriter，通过这种方式拦截写响应
|
||
// 让gin写响应的时候先写到 bodyLogWriter 再写gin.ResponseWriter ，
|
||
// 这样利用中间件里输出访问日志时就能拿到响应了
|
||
// https://stackoverflow.com/questions/38501325/how-to-log-response-body-in-gin
|
||
func (w bodyLogWriter) Write(b []byte) (int, error) {
|
||
w.body.Write(b)
|
||
return w.ResponseWriter.Write(b)
|
||
}
|
||
|
||
// TODO 用于访问request与reponse的日志记录
|
||
// LogAccess define func of log access info
|
||
func LogAccess() gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
// 保存body
|
||
var reqBody []byte
|
||
contentType := c.GetHeader("Content-Type")
|
||
// multipart/form-data 文件上传请求, 不在日志里记录body
|
||
if !strings.Contains(contentType, "multipart/form-data") {
|
||
reqBody, _ = io.ReadAll(c.Request.Body)
|
||
c.Request.Body = io.NopCloser(bytes.NewReader(reqBody))
|
||
|
||
// var request map[string]interface{}
|
||
// if err := c.ShouldBindBodyWith(&request, binding.JSON); err != nil {
|
||
// c.JSON(400, gin.H{"error": err.Error()})
|
||
// return
|
||
// }
|
||
}
|
||
start := time.Now()
|
||
blw := &bodyLogWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer}
|
||
c.Writer = blw
|
||
|
||
accessLog(c, "access_start", time.Since(start), reqBody, nil)
|
||
defer func() {
|
||
var responseLogging string
|
||
if c.Writer.Size() > 10*1024 { // 响应大于10KB 不记录
|
||
responseLogging = "Response data size is too Large to log"
|
||
} else {
|
||
responseLogging = blw.body.String()
|
||
}
|
||
accessLog(c, "access_end", time.Since(start), reqBody, responseLogging)
|
||
}()
|
||
c.Next()
|
||
return
|
||
}
|
||
}
|
||
|
||
func accessLog(c *gin.Context, accessType string, dur time.Duration, body []byte, dataOut interface{}) {
|
||
req := c.Request
|
||
bodyStr := string(body)
|
||
query := req.URL.RawQuery
|
||
path := req.URL.Path
|
||
// TODO: 实现Token认证后再把访问日志里也加上token记录
|
||
// token := c.Request.Header.Get("token")
|
||
logger.New(c).Info("AccessLog",
|
||
"type", accessType,
|
||
"ip", c.ClientIP(),
|
||
//"token", token,
|
||
"method", req.Method,
|
||
"path", path,
|
||
"query", query,
|
||
"body", bodyStr,
|
||
"output", dataOut,
|
||
"time(ms)", int64(dur/time.Millisecond))
|
||
}
|