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

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

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

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

package main

import (
	"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实例，用于替换当前的ResponseWriter
	writer := &responseWriterWrapper{
		ResponseWriter: c.Writer, // 使用原ResponseWriter初始化
		body:           &bytes.Buffer{}, // 初始化一个空的缓冲区
	}
	c.Writer = writer // 将上下文中的Writer替换为我们自定义的writer

	// 继续执行后续的请求处理链
	c.Next()

	// 在所有的处理完成后，可以从writer中获取并打印响应的状态码和内容
	status := writer.statusCode
	body := writer.body

	println("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:8080
	r.Run()
}

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

package main

import (
	"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实例，替换默认的ResponseWriter
	w := &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.H
	if err := json.Unmarshal(w.body.Bytes(), &result); err != nil {
		return // 如果解析出错，直接返回不作处理
	}

	// 检查响应中是否存在code字段，并判断其值是否为0
	codeValue, ok := result["code"].(float64) // JSON解码时int可能转为float64
	if !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"] = encryptedData
	result["is_encrypt"] = true
	isEncrypt = true // 设置标记表示已加密

	// 将修改后的结果重新序列化为JSON格式
	newBody, err := json.Marshal(result)
	if err != nil {
		// 序列化出错则取消加密标记，避免写出错误数据
		isEncrypt = false
		return
	}

	// 将加密后的新响应体写回客户端
	_, _ = 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