modelRT/middleware/trace.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))
}