From d9ccd0c18c8f41a506c60eaa70ab43c4fd3024c9 Mon Sep 17 00:00:00 2001
From: mayiming <1627832236@qq.com>
Date: Sat, 13 Sep 2025 22:57:20 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E4=BB=A4=E7=89=8C=E6=A1=B6?=
 =?UTF-8?q?=E4=B8=AD=E9=97=B4=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 init/ratelimit/ratelimit.go | 41 +++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 init/ratelimit/ratelimit.go

diff --git a/init/ratelimit/ratelimit.go b/init/ratelimit/ratelimit.go
new file mode 100644
index 0000000..0c9df34
--- /dev/null
+++ b/init/ratelimit/ratelimit.go
@@ -0,0 +1,41 @@
+package ratelimit
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"go.uber.org/ratelimit"
+)
+
+var rl = ratelimit.New(20, ratelimit.WithoutSlack)
+
+func RateLimitMiddleware() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		// 1. 给请求设置超时（如3秒，超过3秒未拿到令牌直接返回）
+		ctx, cancel := context.WithTimeout(c.Request.Context(), 3*time.Second)
+		defer cancel()
+
+		// 2. 用带超时的方式获取令牌（需自定义限流器的 Take 逻辑）
+		ch := make(chan struct{}, 1)
+		go func() {
+			rl.Take()        // 阻塞拿令牌
+			ch <- struct{}{} // 拿到令牌后通知主协程
+		}()
+
+		select {
+		case <-ch:
+			// 3. 拿到令牌，继续执行
+			c.Next()
+		case <-ctx.Done():
+			// 4. 超时未拿到令牌，返回请求频繁
+			c.JSON(http.StatusTooManyRequests, gin.H{
+				"code":    429,
+				"message": "请求过于频繁，请稍后重试",
+				"reason":  "等待令牌超时",
+			})
+			c.Abort()
+		}
+	}
+}