【Golang】gin框架如何在中间件中捕获响应并修改后返回

本文讲述如何捕获中间件响应以及重写响应

在gin框架中，在控制器里面调用c.JSON(code, jsonObj)后，向HTTP响应中写入JSON格式的数据，并且设置相应的HTTP状态码。当这个函数被调用时，数据并不会被“保存”到某个特定的位置，而是被直接写入到HTTP响应体中，并通过网络发送给客户端。

如果想在中间件中记录响应日志等操作，我们该如何获取响应数据呢？

package mainimport ("bytes" // 引入bytes包，用于处理字节缓冲区，帮助我们缓存响应体内容"net/http""github.com/gin-gonic/gin" // 导入Gin框架包
)// 定义一个responseWriterWrapper类型，用于包裹gin.ResponseWriter，以扩展其功能
type responseWriterWrapper struct {gin.ResponseWriter // 继承gin.ResponseWriter，保留原有功能body       *bytes.Buffer // 新增一个缓冲区，用于存储响应体的内容statusCode int           // 用于记录响应的状态码
}// 重写WriteHeader方法，用于在响应头被写入之前记录状态码
func (w *responseWriterWrapper) WriteHeader(statusCode int) {w.statusCode = statusCode // 记录状态码w.ResponseWriter.WriteHeader(statusCode) // 调用原始的WriteHeader方法发送状态码
}// 重写Write方法，实现在响应体内容被写入时同时缓存这些内容
func (w *responseWriterWrapper) Write(b []byte) (int, error) {w.body.Write(b) // 将响应体内容写入缓冲区进行缓存return w.ResponseWriter.Write(b) // 调用原始的Write方法将内容写入实际的响应体
}// 定义loggingMiddleware中间件，用于在每个请求结束时打印响应的状态码和内容
func loggingMiddleware(c *gin.Context) {// 创建一个responseWriterWrapper实例，用于替换当前的ResponseWriterwriter := &responseWriterWrapper{ResponseWriter: c.Writer, // 使用原ResponseWriter初始化body:           &bytes.Buffer{}, // 初始化一个空的缓冲区}c.Writer = writer // 将上下文中的Writer替换为我们自定义的writer// 继续执行后续的请求处理链c.Next()// 在所有的处理完成后，可以从writer中获取并打印响应的状态码和内容status := writer.statusCodebody := writer.bodyprintln("Response Status:", status) // 打印状态码println("Response Body:", body.String()) // 将缓冲区内容转换为字符串并打印
}func main() {// 初始化Gin引擎，默认使用Logger和Recovery中间件r := gin.Default()// 使用我们自定义的loggingMiddleware中间件r.Use(loggingMiddleware)// 定义一个简单的路由，返回JSON响应r.GET("/", func(c *gin.Context) {c.JSON(http.StatusOK, gin.H{"message": "Hello, world!"}) // 返回成功状态码和一条消息})// 启动服务器，监听0.0.0.0:8080r.Run()
}

假如需要统一对响应数据做加密，如何修改这个返回数据再响应给客户端呢？

package mainimport ("bytes"         // 引入字节缓冲区处理包，用于缓存响应体"encoding/json" // 引入JSON编码解码包，用于处理JSON数据"net/http""github.com/gin-gonic/gin" // 导入Gin框架包
)// 定义responseWriterWrapper结构体，用于封装gin.ResponseWriter并添加缓冲区以存储响应体内容
type responseWriterWrapper struct {gin.ResponseWriter               // 继承gin.ResponseWriter接口body               *bytes.Buffer // 使用字节缓冲区存储响应体
}// 重写Write方法，将响应体内容写入缓冲区
func (w *responseWriterWrapper) Write(b []byte) (int, error) {return w.body.Write(b)
}// encryptMiddleware 是自定义中间件，用于在响应发送前进行日志记录或数据处理（例如加密）
func encryptMiddleware(c *gin.Context) {// 创建responseWriterWrapper实例，替换默认的ResponseWriterw := &responseWriterWrapper{ResponseWriter: c.Writer,body:           &bytes.Buffer{},}c.Writer = w// 标记，指示是否已经对响应数据进行了加密处理isEncrypt := false// 使用defer确保无论函数如何退出都能重置缓冲区并最终写出响应defer func() {if !isEncrypt {// 如果没有加密，则直接将缓存的内容写出w.ResponseWriter.Write(w.body.Bytes())}w.body.Reset() // 重置缓冲区以备后续请求使用}()// 继续执行后续的处理链，这里是重复调用了c.Next()，在实际应用中应避免，这里为了示例简化处理c.Next()// 解析缓冲区中的JSON数据到gin.H类型变量result中var result gin.Hif err := json.Unmarshal(w.body.Bytes(), &result); err != nil {return // 如果解析出错，直接返回不作处理}// 检查响应中是否存在code字段，并判断其值是否为0codeValue, ok := result["code"].(float64) // JSON解码时int可能转为float64if !ok || int(codeValue) != 0 {return // 如果code不是预期值，则不进行加密处理}// 获取响应中的"data"字段dataValue, ok := result["data"]if !ok {return // 如果"data"不存在，则不进行处理}// 加密逻辑，这里仅为示例，实际加密过程应替换此简单字符串替换逻辑encryptFunc := func(data any) string {return "我是加密后字符串"}encryptedData := encryptFunc(dataValue)// 修改响应体中的"data"为加密后的数据，并增加"is_encrypt"字段result["data"] = encryptedDataresult["is_encrypt"] = trueisEncrypt = true // 设置标记表示已加密// 将修改后的结果重新序列化为JSON格式newBody, err := json.Marshal(result)if err != nil {// 序列化出错则取消加密标记，避免写出错误数据isEncrypt = falsereturn}// 将加密后的新响应体写回客户端_, _ = w.ResponseWriter.Write(newBody)
}func main() {// 初始化Gin路由器，并使用自定义中间件r := gin.Default()r.Use(encryptMiddleware)// 定义一个GET路由，返回JSON响应r.GET("/", func(c *gin.Context) {c.JSON(http.StatusOK, gin.H{"code": 0, "message": "Hello, world!", "data": "..."})})// 启动HTTP服务器r.Run() // 监听0.0.0.0:8080并开始服务
}

执行：curl "http:/127.0.0.1:8080/" 返回

{"code": 0,"data": "我是加密后字符串","is_encrypt": true,"message": "Hello, world!"
}

参考

[1]: How to rewrite response body in middleware? #3384