light-apiengine-feature/middleware/logger.go

package middleware

import (
	"bytes"
	"fmt"
	"time"

	"git.qianqiusoft.com/public/glog"
	"git.qianqiusoft.com/qianqiusoft/light-apiengine/config"
	"git.qianqiusoft.com/qianqiusoft/light-apiengine/utils"
	"github.com/gin-gonic/gin"
)

/**
 * 自定义日志响应类，主要是把resp的响应内容写入到body中
 */
type loggerRespWriter struct {
	gin.ResponseWriter
	respBody *bytes.Buffer
}

// 写字节数组
func (w loggerRespWriter) Write(b []byte) (int, error) {
	w.respBody.Write(b) // 保存
	return w.ResponseWriter.Write(b)
}

// 写字符串
func (w loggerRespWriter) WriteString(s string) (int, error) {
	w.respBody.WriteString(s) // 保存
	return w.ResponseWriter.WriteString(s)
}

/**
 * 日志信息
 */
type LogInfo struct {
	ReqClientIp string
	ReqTime     time.Time
	ReqMethod   string
	ReqUrl      string
	ReqProto    string
	ReqUa       string
	ReqReferer  string
	Header      []byte
	LoginId     string
	ReqHeader   []byte
	ReqBody     []byte
	RespTime    time.Time
	RespBody    []byte

	CostTime float64
}

var bucket utils.SyncSlice

// 日志chan
var __logInfoChan chan *LogInfo

// app名称
var __appName string = ""

// 日志处理函数
var __logHandler func(info []LogInfo) = nil

func init() {
	defer func() {
		if p := recover(); p != nil {
			glog.Errorln("ecover", p)
		}
	}()

	// 设置app名称
	__appName = config.AppConfig.GetKey("app_name")
	__logInfoChan = make(chan *LogInfo, 2000)

	// 日志处理
	logProcess()
}

// 设置日志处理接口
func SetLogHandler(handler func(info []LogInfo)) {
	__logHandler = handler
}

// 获取日志中间件
func LoggerMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		respWriter := &loggerRespWriter{
			ResponseWriter: c.Writer,
			respBody:       bytes.NewBuffer([]byte{}),
		}
		c.Writer = respWriter // 注入自定writer

		logInfo := &LogInfo{}
		logInfo.ReqTime = time.Now()

		// 下一个请求
		c.Next()
		hByte := bytes.Buffer{}
		bByte := make([]byte, 1024)

		cReq := c.Request.Clone(c)
		cReq.Header.Write(&hByte)
		rbn, _ := cReq.Body.Read(bByte)

		// 设置对象
		logInfo.RespBody = respWriter.respBody.Bytes()
		logInfo.RespTime = time.Now()
		logInfo.ReqMethod = c.Request.Method
		logInfo.ReqUrl = c.Request.RequestURI
		logInfo.ReqProto = c.Request.Proto
		logInfo.ReqUa = c.Request.UserAgent()
		logInfo.ReqReferer = c.Request.Referer()
		logInfo.ReqBody = bByte[0:rbn]
		logInfo.ReqHeader = hByte.Bytes()
		logInfo.ReqClientIp = c.ClientIP()
		logInfo.CostTime = float64(logInfo.RespTime.Sub(logInfo.ReqTime).Seconds())

		token := c.GetHeader("token")
		if token != "" {
			tk := utils.GetGlobalTokenStore().Get(token)
			if tk != nil {
				logInfo.LoginId = tk.LoginID
			}
		}
		if logInfo.LoginId == "" {
			logInfo.LoginId = "__no_login__"
		}

		__logInfoChan <- logInfo
	}
}

// 日志处理进程
func logProcess() {
	go func() {
		defer func() {
			if rec := recover(); rec != nil {
				fmt.Println(rec)
			}
		}()
		for {
			select {
			case <-time.Tick(6 * time.Second):
				if __logHandler != nil {
					logs := []LogInfo{}
					for _, v := range bucket.Reset() {
						logs = append(logs, *v.(*LogInfo))
					}
					__logHandler(logs)
				}
			case logInfo := <-__logInfoChan:
				if __logHandler != nil {
					bucket.Append(logInfo)
				}
			}
		}
	}()
}