go-zero的日志中间件

package middleware

import (
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/rest/httpx"
	"net/http"
	"net/http/httputil"
	"strings"
)

const (
	// ApplicationJson stands for application/json.
	ApplicationJson = "application/json"
	// ContentType is the header key for Content-Type.
	ContentType = "Content-Type"
)

type LogMiddleware struct {
}

func NewLogMiddleware() *LogMiddleware {
	return &LogMiddleware{}
}

func (m *LogMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		proxy := &responseProxy{w: w}
		requestLog(r)
		next(proxy, r)
		logx.Infof("LogMiddleware response uri:%s jsonResult :%+v", r.RequestURI, string(proxy.body))
	}
}

type responseProxy struct {
	w    http.ResponseWriter
	body []byte
}

func (p *responseProxy) Header() http.Header {
	return p.w.Header()
}
func (p *responseProxy) Write(data []byte) (int, error) {
	p.body = append(p.body, data...)
	return p.w.Write(data)
}
func (p *responseProxy) WriteHeader(statusCode int) {
	p.w.WriteHeader(statusCode)
}

func requestLog(r *http.Request) {
	// 打印所有header
	logx.Infof("LogMiddleware request uri:%s header :%+v", r.RequestURI, r.Header)
	// json日志
	if withJsonBody(r) {
		requestDump, err := httputil.DumpRequest(r, true)
		logx.Infof("LogMiddleware request uri:%s jsonParams :%+v, err:%+v", r.RequestURI, string(requestDump), err)
	} else {
		// form表单日志和其他
		formParams, err := httpx.GetFormValues(r)
		logx.Infof("LogMiddleware request uri:%s formParams :%+v, err:%+v", r.RequestURI, formParams, err)
	}
}

func withJsonBody(r *http.Request) bool {
	return r.ContentLength > 0 && strings.Contains(r.Header.Get(ContentType), ApplicationJson)
}