light-apiengine-feature/middleware/logger/logger.go

package logger

import (
	"bytes"
	"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
	ReqPostData string
	LoginId     string

	RespTime time.Time
	RespBody string
	RespCode string // 由 resp body 解析
	RespMsg  string // 由 resp body 解析
	RespData string // 由 resp body 解析

	CostTime float64
}

// 日志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()

		// 设置对象
		logInfo.RespBody = respWriter.respBody.String()
		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.ReqPostData = c.Request.PostForm.Encode()
		logInfo.ReqClientIp = c.ClientIP()

		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() {
		for {
			select {
			case logInfo := <-__logInfoChan:
				if __logHandler == nil {
					cassandraLogHandler(logInfo)
				} else {
					__logHandler(logInfo)
				}
			}
		}
	}()
}