打造一个优雅的 Go 操作日志中间件 - 轻松实现请求响应全链路追踪

程序员爱钓鱼
37 阅读4分钟

前言

在企业应用开发中,操作日志审计记录是一个非常重要的功能需求。本文将详细介绍如何实现一个完整的 Go 操作日志审计中间件,包含所有必要的代码实现。

完整代码实现

1. 首先定义操作记录模型 (model/system/sys_operation_record.go)

package system

import (
    "time"
)

type SysOperationRecord struct {
    ID           uint          `json:"id" gorm:"primarykey"`
    Ip           string        `json:"ip" gorm:"column:ip"`             // 请求ip
    Method       string        `json:"method" gorm:"column:method"`     // 请求方法
    Path         string        `json:"path" gorm:"column:path"`         // 请求路径
    Status       int           `json:"status" gorm:"column:status"`     // 请求状态
    Latency      time.Duration `json:"latency" gorm:"column:latency"`   // 延迟
    Agent        string        `json:"agent" gorm:"column:agent"`       // 代理
    ErrorMessage string        `json:"error_message" gorm:"column:error_message"` // 错误信息
    Body         string        `json:"body" gorm:"type:text;column:body"`         // 请求Body
    Resp         string        `json:"resp" gorm:"type:text;column:resp"`         // 响应Body
    UserID       int           `json:"user_id" gorm:"column:user_id"`   // 用户id
    CreatedAt    time.Time     `json:"created_at" gorm:"column:created_at"`
}

2. 操作记录服务层 (service/system/sys_operation_record.go)

package system

import (
    "project/server/app/model/system"
    "project/server/global"
)

type OperationRecordService struct{}

func (o *OperationRecordService) CreateSysOperationRecord(record system.SysOperationRecord) error {
    return global.GVA_DB.Create(&record).Error
}

3. 中间件实现 (middleware/operation.go)

package middleware

import (
    "bytes"
    "encoding/json"
    "github.com/gin-gonic/gin"
    "go.uber.org/zap"
    "io"
    "net/http"
    "net/url"
    "project/server/app/model/system"
    "project/server/app/services"
    "project/server/global"
    "project/server/utils"
    "strconv"
    "strings"
    "sync"
    "time"
)

var operationRecordService = services.ServiceGroupApp.SystemServiceGroup.OperationRecordService

var respPool sync.Pool
var bufferSize = 1024

func init() {
    respPool.New = func() interface{} {
        return make([]byte, bufferSize)
    }
}

func OperationRecord() gin.HandlerFunc {
    return func(c *gin.Context) {
        var body []byte
        var userId int
        
        // 处理请求体
        if c.Request.Method != http.MethodGet {
            var err error
            body, err = io.ReadAll(c.Request.Body)
            if err != nil {
                global.GVA_LOG.Error("read body from request error:", zap.Error(err))
            } else {
                c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
            }
        } else {
            query := c.Request.URL.RawQuery
            query, _ = url.QueryUnescape(query)
            split := strings.Split(query, "&")
            m := make(map[string]string)
            for _, v := range split {
                kv := strings.Split(v, "=")
                if len(kv) == 2 {
                    m[kv[0]] = kv[1]
                }
            }
            body, _ = json.Marshal(&m)
        }

        // 获取用户ID
        claims, _ := utils.GetClaims(c)
        if claims != nil && claims.BaseClaims.ID != 0 {
            userId = int(claims.BaseClaims.ID)
        } else {
            id, err := strconv.Atoi(c.Request.Header.Get("x-user-id"))
            if err != nil {
                userId = 0
            }
            userId = id
        }

        record := system.SysOperationRecord{
            Ip:     c.ClientIP(),
            Method: c.Request.Method,
            Path:   c.Request.URL.Path,
            Agent:  c.Request.UserAgent(),
            Body:   "",
            UserID: userId,
        }

        // 处理文件上传
        if strings.Contains(c.GetHeader("Content-Type"), "multipart/form-data") {
            record.Body = "[文件]"
        } else {
            if len(body) > bufferSize {
                record.Body = "[超出记录长度]"
            } else {
                record.Body = string(body)
            }
        }

        writer := responseBodyWriter{
            ResponseWriter: c.Writer,
            body:          &bytes.Buffer{},
        }
        c.Writer = writer
        now := time.Now()

        c.Next()

        // 记录响应信息
        latency := time.Since(now)
        record.ErrorMessage = c.Errors.ByType(gin.ErrorTypePrivate).String()
        record.Status = c.Writer.Status()
        record.Latency = latency
        record.Resp = writer.body.String()

        // 处理下载文件的情况
        if strings.Contains(c.Writer.Header().Get("Pragma"), "public") ||
            strings.Contains(c.Writer.Header().Get("Expires"), "0") ||
            strings.Contains(c.Writer.Header().Get("Cache-Control"), "must-revalidate, post-check=0, pre-check=0") ||
            strings.Contains(c.Writer.Header().Get("Content-Type"), "application/force-download") ||
            strings.Contains(c.Writer.Header().Get("Content-Type"), "application/octet-stream") ||
            strings.Contains(c.Writer.Header().Get("Content-Type"), "application/vnd.ms-excel") ||
            strings.Contains(c.Writer.Header().Get("Content-Type"), "application/download") ||
            strings.Contains(c.Writer.Header().Get("Content-Disposition"), "attachment") ||
            strings.Contains(c.Writer.Header().Get("Content-Transfer-Encoding"), "binary") {
            if len(record.Resp) > bufferSize {
                record.Body = "超出记录长度"
            }
        }

        // 保存记录
        if err := operationRecordService.CreateSysOperationRecord(record); err != nil {
            global.GVA_LOG.Error("create operation record error:", zap.Error(err))
        }
    }
}

type responseBodyWriter struct {
    gin.ResponseWriter
    body *bytes.Buffer
}

func (r responseBodyWriter) Write(b []byte) (int, error) {
    r.body.Write(b)
    return r.ResponseWriter.Write(b)
}

4. 使用示例 (main.go)

package main

import (
    "github.com/gin-gonic/gin"
    "project/server/middleware"
)

func main() {
    r := gin.Default()
    
    // 注册中间件
    r.Use(middleware.OperationRecord())
    
    // API路由示例
    r.GET("/api/users", func(c *gin.Context) {
        c.JSON(200, gin.H{
            "message": "获取用户列表成功",
            "data": []map[string]interface{}{
                {"id": 1, "name": "张三"},
                {"id": 2, "name": "李四"},
            },
        })
    })
    
    r.POST("/api/users", func(c *gin.Context) {
        c.JSON(200, gin.H{
            "message": "创建用户成功",
            "data": map[string]interface{}{
                "id": 3,
                "name": "王五",
            },
        })
    })
    
    // 启动服务
    r.Run(":8080")
}

5. 数据库初始化 (initialize/gorm.go)

package initialize

import (
    "gorm.io/gorm"
    "project/server/app/model/system"
)

func RegisterTables(db *gorm.DB) {
    db.AutoMigrate(
        system.SysOperationRecord{},
    )
}

使用方法

  1. 首先确保项目中已经配置好了数据库连接

  2. 在 main.go 中注册中间件:

r := gin.Default()
r.Use(middleware.OperationRecord())
  1. 现在,所有的 API 请求都会被记录到数据库中,包括:
    • 请求的IP地址
    • 请求方法(GET, POST等)
    • 请求路径
    • 请求参数
    • 响应状态码
    • 响应内容
    • 执行时间
    • 用户ID
    • 错误信息(如果有)

测试示例

使用 curl 发送请求:

# GET请求测试
curl http://localhost:8080/api/users

# POST请求测试
curl -X POST http://localhost:8080/api/users \
  -H "Content-Type: application/json" \
  -d '{"name":"王五","age":25}'

查看数据库中的记录:

SELECT * FROM sys_operation_records ORDER BY created_at DESC LIMIT 1;

特点总结

  1. 完整的请求信息记录

    • 支持所有 HTTP 方法
    • 自动处理 GET 查询参数
    • 记录请求体和响应内容

  2. 特殊场景处理

    • 文件上传请求的特殊处理
    • 文件下载响应的特殊处理
    • 大内容的自动截断

  3. 性能优化

    • 使用 sync.Pool 复用缓冲区
    • 响应内容大小限制
    • 高效的响应内容捕获

  4. 用户识别

    • 支持 JWT 认证
    • 支持自定义 Header 中的用户ID

注意事项

  1. 数据库性能

    • 建议对 sys_operation_records 表建立适当的索引
    • 考虑数据定期清理策略

  2. 敏感信息

    • 注意在记录请求体时过滤敏感信息
    • 可以针对特定接口关闭记录

  3. 存储容量

    • 注意控制记录的内容大小
    • 建议配置合理的 bufferSize

结语

这个中间件提供了一个完整的操作审计解决方案,可以帮助我们追踪系统中的所有操作。它不仅可以用于问题排查,还可以用于安全审计、性能分析等多个场景。

希望这个详细的实现能够帮助到需要此功能的开发者！

avatar
文章
阅读
粉丝