【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!"
}