From 8d0f1e6aa31ea0dfdc8f1b47d4d266d3765bfb48 Mon Sep 17 00:00:00 2001
From: liangzai <2440983361@qq.com>
Date: Sat, 22 Nov 2025 15:43:24 +0800
Subject: [PATCH] fix

---
 config.yaml                             |  2 +
 internal/config/config.go               |  8 ++-
 internal/shared/middleware/ratelimit.go | 76 +++++++++++++++++++++++++
 3 files changed, 83 insertions(+), 3 deletions(-)

diff --git a/config.yaml b/config.yaml
index 4bd9f34..af908df 100644
--- a/config.yaml
+++ b/config.yaml
@@ -160,6 +160,8 @@ ocr:
 ratelimit:
     requests: 5000
     window: 60s
+    exclude_domains: # 排除频率限制的域名
+        - "api.*" # 排除所有 api.* 二级域名（业务核心接口）
 
 # 每日请求限制配置
 daily_ratelimit:
diff --git a/internal/config/config.go b/internal/config/config.go
index ba63e76..87a8c54 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -117,9 +117,11 @@ type JWTConfig struct {
 
 // RateLimitConfig 限流配置
 type RateLimitConfig struct {
-	Requests int           `mapstructure:"requests"`
-	Window   time.Duration `mapstructure:"window"`
-	Burst    int           `mapstructure:"burst"`
+	Requests      int           `mapstructure:"requests"`
+	Window        time.Duration `mapstructure:"window"`
+	Burst         int           `mapstructure:"burst"`
+	ExcludePaths   []string     `mapstructure:"exclude_paths"`   // 排除频率限制的路径
+	ExcludeDomains []string     `mapstructure:"exclude_domains"` // 排除频率限制的域名
 }
 
 // DailyRateLimitConfig 每日限流配置
diff --git a/internal/shared/middleware/ratelimit.go b/internal/shared/middleware/ratelimit.go
index b544b0d..6fff0b9 100644
--- a/internal/shared/middleware/ratelimit.go
+++ b/internal/shared/middleware/ratelimit.go
@@ -2,6 +2,7 @@ package middleware
 
 import (
 	"fmt"
+	"strings"
 	"sync"
 	"time"
 
@@ -42,6 +43,23 @@ func (m *RateLimitMiddleware) GetPriority() int {
 // Handle 返回中间件处理函数
 func (m *RateLimitMiddleware) Handle() gin.HandlerFunc {
 	return func(c *gin.Context) {
+		// 检查是否在排除域名中
+		host := c.Request.Host
+		// 移除端口部分
+		if idx := strings.Index(host, ":"); idx != -1 {
+			host = host[:idx]
+		}
+		if m.isExcludedDomain(host) {
+			c.Next()
+			return
+		}
+
+		// 检查是否在排除路径中
+		if m.isExcludedPath(c.Request.URL.Path) {
+			c.Next()
+			return
+		}
+
 		// 获取客户端标识（IP地址）
 		clientID := m.getClientID(c)
 
@@ -150,6 +168,64 @@ func (m *RateLimitMiddleware) cleanup() {
 	}
 }
 
+// isExcludedDomain 检查域名是否在排除列表中
+func (m *RateLimitMiddleware) isExcludedDomain(host string) bool {
+	for _, excludeDomain := range m.config.RateLimit.ExcludeDomains {
+		// 支持通配符匹配
+		if strings.HasPrefix(excludeDomain, "*") {
+			// 后缀匹配，如 "*.api.example.com" 匹配 "api.example.com"
+			if strings.HasSuffix(host, excludeDomain[1:]) {
+				return true
+			}
+		} else if strings.HasSuffix(excludeDomain, "*") {
+			// 前缀匹配，如 "api.*" 匹配 "api.example.com"
+			if strings.HasPrefix(host, excludeDomain[:len(excludeDomain)-1]) {
+				return true
+			}
+		} else {
+			// 精确匹配
+			if host == excludeDomain {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// isExcludedPath 检查路径是否在排除列表中
+func (m *RateLimitMiddleware) isExcludedPath(path string) bool {
+	for _, excludePath := range m.config.RateLimit.ExcludePaths {
+		// 支持多种匹配模式
+		if strings.HasPrefix(excludePath, "*") {
+			// 前缀匹配，如 "*api_name" 匹配 "/api/v1/any_api_name"
+			if strings.Contains(path, excludePath[1:]) {
+				return true
+			}
+		} else if strings.HasSuffix(excludePath, "*") {
+			// 后缀匹配，如 "/api/v1/*" 匹配 "/api/v1/any_api_name"
+			if strings.HasPrefix(path, excludePath[:len(excludePath)-1]) {
+				return true
+			}
+		} else if strings.Contains(excludePath, "*") {
+			// 中间通配符匹配，如 "/api/v1/*api_name" 匹配 "/api/v1/any_api_name"
+			parts := strings.Split(excludePath, "*")
+			if len(parts) == 2 {
+				prefix := parts[0]
+				suffix := parts[1]
+				if strings.HasPrefix(path, prefix) && strings.HasSuffix(path, suffix) {
+					return true
+				}
+			}
+		} else {
+			// 精确匹配
+			if path == excludePath {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // GetStats 获取限流统计
 func (m *RateLimitMiddleware) GetStats() map[string]interface{} {
 	m.mutex.RLock()