go中间件获取请求和响应数据

go-zero中间件获取请求和响应数据

Go中，http.ResponseWriter对象无法直接获取响应内容和响应头，因为这些信息是在HTTP响应被发送到客户端后才能得到。

如果需要获取HTTP响应内容和响应头，可以使用一个缓冲区（例如bytes.Buffer类型）来存储响应内容，并创建一个新的http.ResponseWriter对象，将响应写入缓冲区和原始的http.ResponseWriter对象。然后，在HTTP响应发送到客户端后，可以从缓冲区中获取响应内容和响应头。
代码如下：

package middleware

import (
	"bytes"
	"io"
	"log"
	"net/http"
)

type LogMiddleware struct {
	// 可以在这里添加操作数据的依赖
}

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

type logResponseWriter struct {
	writer http.ResponseWriter
	code   int
	// 存响应数据
	buf    *bytes.Buffer
}

func newLogResponseWriter(writer http.ResponseWriter, code int) *logResponseWriter {
	var buf bytes.Buffer
	return &logResponseWriter{
		writer: writer,
		code:   code,
		buf:    &buf,
	}
}

func (w *logResponseWriter) Write(bs []byte) (int, error) {
	w.buf.Write(bs)
	return w.writer.Write(bs)
}

func (w *logResponseWriter) Header() http.Header {
	return w.writer.Header()
}

func (w *logResponseWriter) WriteHeader(code int) {
	w.code = code
	w.writer.WriteHeader(code)
}

func (m *LogMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var dup io.ReadCloser
		dup, r.Body, _ = drainBody(r.Body)
		lwr := newLogResponseWriter(w, http.StatusOK)
		next(lwr, r)
		r.Body = dup
		// 将自己逻辑写在logDetailLogic
		m.logDetailLogic(r, lwr)
	}
}

func (m *LogMiddleware) logDetailLogic(request *http.Request, response *logResponseWriter) {
	if request.Body != nil {
		bs, _ := io.ReadAll(request.Body)
		log.Println("请求内容：", string(bs))
	}
	if len(response.buf.Bytes()) > 0 {
		log.Println("响应内容：", response.buf.String())
	}
}

// drainBody from httputil.drainBody
func drainBody(b io.ReadCloser) (r1, r2 io.ReadCloser, err error) {
	if b == nil || b == http.NoBody {
		// No copying needed. Preserve the magic sentinel meaning of NoBody.
		return http.NoBody, http.NoBody, nil
	}
	var buf bytes.Buffer
	if _, err = buf.ReadFrom(b); err != nil {
		return nil, b, err
	}
	if err = b.Close(); err != nil {
		return nil, b, err
	}
	return io.NopCloser(&buf), io.NopCloser(bytes.NewReader(buf.Bytes())), nil
}