Gin 架構在中間件中擷取 response body 的方法

在寫一個Gin架構日志中間件的時候，需要記錄請求和響應相關的一些資料，例如請求參數、請求方法、請求時間、請求頭、耗時、響應狀态碼、響應資料等，gin為擷取這些資料基本都提供了現成的方法，但是擷取響應資料還是有一定難度和複雜度的。

那麼，該如何擷取響應資料也就是 response body 呢？先上代碼。

代碼示例

1、先寫一個 middleware ，簡單列印一下 response body

package middleware

import (
   "bytes"
   "fmt"
   "github.com/gin-gonic/gin"
)

//自定義一個結構體，實作 gin.ResponseWriter interface
type responseWriter struct {
   gin.ResponseWriter
   b *bytes.Buffer
}

//重寫 Write([]byte) (int, error) 方法
func (w responseWriter) Write(b []byte) (int, error) {
   //向一個bytes.buffer中寫一份資料來為擷取body使用
   w.b.Write(b)
   //完成gin.Context.Writer.Write()原有功能
   return w.ResponseWriter.Write(b)
}

func PrintResponse(c *gin.Context) {
   writer := responseWriter{
      c.Writer,
      bytes.NewBuffer([]byte{}),
   }
   c.Writer = writer

   c.Next()

   fmt.Println("response body：" + writer.b.String())
}
           

2、引用 middleware ，看看效果

package main

import (
   "github.com/gin-gonic/gin"
   "hello/middleware"
   "net/http"
)

func main() {
   r := gin.New()

   //添加擷取響應内容 middleware
   r.Use(middleware.PrintResponse)

   r.GET("/test", func(c *gin.Context) {
      c.JSON(http.StatusOK, gin.H{"name": "xiaoming"})
   })
   // Listen and Server in 0.0.0.0:8080
   r.Run(":8080")
}           

打開浏覽器通路 http://127.0.0.1:8080/test ，即可在控制台中看到 response body 的輸出内容

response body：{"name":"xiaoming"}           

原理

通過上面的代碼可以看出，gin 通過調用如下方法寫入了 response body

 c.JSON(http.StatusOK, gin.H{"name": "xiaoming"})           

追進 JSON 方法，源碼如下，調用了 Render 方法

package gin

func (c *Context) JSON(code int, obj interface{}) {
   c.Render(code, render.JSON{Data: obj})
}           

追進 Render 方法， 可以看出 gin.Context.Writer 作為參數傳給了 r.Render() 方法，這裡形參 r 的實參為 render.JSON{Data: obj} ，是以實際調用的是 func (r JSON) Render(w http.ResponseWriter) 方法。

package gin

// Render writes the response headers and calls render.Render to render data.
func (c *Context) Render(code int, r render.Render) {
   c.Status(code)

   if !bodyAllowedForStatus(code) {
      r.WriteContentType(c.Writer)
      c.Writer.WriteHeaderNow()
      return
   }

   if err := r.Render(c.Writer); err != nil {
      panic(err)
   }
}           

func (r JSON) Render(w http.ResponseWriter) 源碼如下

package render

// Render (JSON) writes data with custom ContentType.
func (r JSON) Render(w http.ResponseWriter) (err error) {
   if err = WriteJSON(w, r.Data); err != nil {
      panic(err)
   }
   return
}

// WriteJSON marshals the given interface object and writes it with custom ContentType.
func WriteJSON(w http.ResponseWriter, obj interface{}) error {
   writeContentType(w, jsonContentType)
   jsonBytes, err := json.Marshal(obj)
   if err != nil {
      return err
   }
   _, err = w.Write(jsonBytes)
   return err
}           

在 func (r JSON) Render() 方法中，gin.Context.Writer 被傳到了 WriteJSON() 方法中，最終寫入資料調用的是 gin.Context.Writer.Write() 方法。gin.ResponseWriter 源碼如下

package gin

type Context struct {
   //...
   writermem responseWriter
   Writer    ResponseWriter
   //...
}

// ResponseWriter ...
type ResponseWriter interface {
    //...
    http.ResponseWriter
    //...
}           

可以看出 gin.Context.Writer 類型為 interface gin.ResponseWriter。

package http

type ResponseWriter interface {
    //...
    Write([]byte) (int, error)
    //...
}           

要實作 gin.ResponseWriter 接口，必須實作Write([]byte) (int, error) 方法。是以寫入 response body 調用的是 gin.Context.Writer.Write() 方法，gin.Context.Writer 需要是type gin.ResponseWriter interface 的一個具體實作。

到此，可以看出上面代碼示例的思路：

實作 type gin.ResponseWriter interface 并重寫 Write([]byte) (int, error) 方法，該方法在實作 gin.Context.Writer.Write() 原有功能的同時，再向一個 bytes.buffer 中寫一份資料來用于擷取 response body 使用。