第16章：自定义中间件

本章目标

学完本章后，你将能够：

1. 从零编写自定义中间件
2. 使用 c.Set() / c.Get() 在中间件间传递数据
3. 实现请求耗时统计、TraceID、降级等实用中间件
4. 控制中间件执行顺序

前置知识

• 需要先学习：第14-15章（中间件概念与内置中间件）

---

16.1 中间件模板

// 自定义中间件的标准模板
func MyMiddleware() gin.HandlerFunc {
	// 初始化代码：程序启动时执行一次
	
	return func(c *gin.Context) {
		// 前置代码：请求到达时执行
		
		c.Next()  // 交给下一个处理器
		
		// 后置代码：响应返回时执行
	}
}

---

16.2 实用中间件示例

请求耗时统计

func RequestTimer() gin.HandlerFunc {
	return func(c *gin.Context) {
		start := time.Now()
		c.Next()
		elapsed := time.Since(start)
		
		// 超过 100ms 打 Warn
		if elapsed > 100*time.Millisecond {
			fmt.Printf("⚠️ 慢请求: %s %s 耗时 %v\n",
				c.Request.Method, c.Request.URL.Path, elapsed)
		}
	}
}

TraceID 中间件（请求追踪）

func TraceID() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 先从请求头获取，没有就生成
		traceID := c.GetHeader("X-Trace-ID")
		if traceID == "" {
			traceID = uuid.New().String()
		}
		
		// 存到 Context（后续 Handler 和中间件都能取到）
		c.Set("trace_id", traceID)
		
		// 返回给客户端
		c.Header("X-Trace-ID", traceID)
		
		c.Next()
	}
}

// 在 Handler 中获取
func handler(c *gin.Context) {
	traceID, _ := c.Get("trace_id")
	fmt.Println("TraceID:", traceID)
}

简单限流（基于计数器）

func SimpleRateLimiter(limit int, window time.Duration) gin.HandlerFunc {
	var mu sync.Mutex
	counters := make(map[string][]time.Time)
	
	return func(c *gin.Context) {
		ip := c.ClientIP()
		now := time.Now()
		
		mu.Lock()
		// 清理过期记录
		var valid []time.Time
		for _, t := range counters[ip] {
			if now.Sub(t) < window {
				valid = append(valid, t)
			}
		}
		counters[ip] = valid
		
		if len(valid) >= limit {
			mu.Unlock()
			c.AbortWithStatusJSON(429, gin.H{"error": "请求太频繁"})
			return
		}
		
		counters[ip] = append(counters[ip], now)
		mu.Unlock()
		
		c.Next()
	}
}

// 使用：每个 IP 每分钟最多 60 次
r.Use(SimpleRateLimiter(60, time.Minute))

---

16.3 c.Set() / c.Get() 数据传递

// 中间件 A：设置数据
func MiddlewareA() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Set("user_id", 123)
		c.Set("role", "admin")
		c.Next()
	}
}

// 中间件 B：读取数据
func MiddlewareB() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Next()
		userID, _ := c.Get("user_id")
		role, _ := c.Get("role")
		fmt.Printf("请求来自用户 %v, 角色 %v\n", userID, role)
	}
}

---

16.4 完整示例

package main

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

func Trace() gin.HandlerFunc {
	return func(c *gin.Context) {
		traceID := fmt.Sprintf("%d", time.Now().UnixNano())
		c.Set("trace_id", traceID)
		c.Header("X-Trace-ID", traceID)
		fmt.Printf("[%s] → %s %s\n", traceID, c.Request.Method, c.Request.URL.Path)
		c.Next()
	}
}

func Timer() gin.HandlerFunc {
	return func(c *gin.Context) {
		start := time.Now()
		c.Next()
		traceID, _ := c.Get("trace_id")
		fmt.Printf("[%s] ← 耗时 %v\n", traceID, time.Since(start))
	}
}

func main() {
	r := gin.New()
	r.Use(Trace(), Timer(), gin.Recovery())

	r.GET("/hello", func(c *gin.Context) {
		time.Sleep(50 * time.Millisecond)  // 模拟业务处理
		c.JSON(200, gin.H{"msg": "hello"})
	})

	r.Run(":8080")
}

---

常见错误

错误1：中间件中 panic 没有被捕获

// ❌ gin.New() 没有 Recovery
r := gin.New()
r.Use(MyMiddleware)  // 如果中间件 panic，服务崩溃

// ✅ 加 Recovery
r := gin.New()
r.Use(gin.Recovery())
r.Use(MyMiddleware)

错误2：c.Set 的 key 冲突

// ❌ 多个中间件用相同的 key
func M1(c *gin.Context) { c.Set("user", "admin") }
func M2(c *gin.Context) { c.Set("user", "guest") }  // 覆盖了 M1

// ✅ 用明确的命名空间
func M1(c *gin.Context) { c.Set("auth:user", "admin") }
func M2(c *gin.Context) { c.Set("trace:id", "abc123") }

---

本章小结

• 中间件模板：初始化 → 返回 func(c *gin.Context) → c.Next() 分界
• c.Set(key, value) / c.Get(key) 跨中间件传数据
• 常见自定义中间件：耗时、TraceID、限流
• 中间件注册顺序 = 执行顺序

练习题

1. 写一个打印请求耗时的中间件。
2. 写一个 TraceID 中间件，从请求头读取或自动生成，并存入 Context 和响应头。
3. 写一个简单的 IP 限流中间件（每个 IP 每分钟最多 10 次）。
4. 使用 c.Set 和 c.Get 在两个中间件间传递用户信息。
5. （思考题）如果要在中间件中修改请求体（Body），需要注意什么？