添加 API 限流和支付渠道熔断

2026-06-10 15:48:04 +08:00
parent 0fbe30b377
commit df7d6a2675
7 changed files with 171 additions and 4 deletions
@@ -76,6 +76,7 @@ require (
 	github.com/quic-go/qpack v0.6.0 // indirect
 	github.com/quic-go/quic-go v0.59.1 // indirect
 	github.com/rs/xid v1.6.0 // indirect
+	github.com/sony/gobreaker/v2 v2.4.0 // indirect
 	github.com/tinylib/msgp v1.6.1 // indirect
 	github.com/tjfoc/gmsm v1.4.1 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
@@ -235,6 +235,8 @@ github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
 github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d/go.mod h1:OnSkiWE9lh6wB0YB77sQom3nweQdgAjqCqsofrRNTgc=
 github.com/smartystreets/assertions v1.1.0/go.mod h1:tcbTF8ujkAEcZ8TElKY+i30BzYlVhC/LOxJk7iOWnoo=
 github.com/smartystreets/goconvey v1.6.4/go.mod h1:syvi0/a8iFYH4r/RixwvyeAJjdLS9QV7WQ/tjFTllLA=
+github.com/sony/gobreaker/v2 v2.4.0 h1:g2KJRW1Ubty3+ZOcSEUN7K+REQJdN6yo6XvaML+jptg=
+github.com/sony/gobreaker/v2 v2.4.0/go.mod h1:pTyFJgcZ3h2tdQVLZZruK2C0eoFL1fb/G83wK1ZQl+s=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoHMkEqE=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
@@ -17,6 +17,7 @@ type Config struct {
 	SMS           SMSConfig
 	Realname      RealnameConfig
 	Log           LogConfig
+	RateLimit     RateLimitConfig
 }

 type StorageConfig struct {
@@ -48,6 +49,11 @@ type LogConfig struct {
 	EnableFile    bool
 }

+type RateLimitConfig struct {
+	Enabled           bool
+	RequestsPerMinute int
+}
+
 func Load() Config {
 	return Config{
 		AppEnv:        getEnv("APP_ENV", "development"),
@@ -82,6 +88,10 @@ func Load() Config {
 			EnableConsole: getEnvBool("LOG_ENABLE_CONSOLE", true),
 			EnableFile:    getEnvBool("LOG_ENABLE_FILE", true),
 		},
+		RateLimit: RateLimitConfig{
+			Enabled:           getEnvBool("RATE_LIMIT_ENABLED", true),
+			RequestsPerMinute: getEnvInt("RATE_LIMIT_REQUESTS_PER_MINUTE", 300),
+		},
 	}
 }

@@ -0,0 +1,83 @@
+package middleware
+
+import (
+	"net/http"
+	"strconv"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+type rateLimitBucket struct {
+	count    int
+	resetAt  time.Time
+	lastSeen time.Time
+}
+
+type rateLimiter struct {
+	mu      sync.Mutex
+	limit   int
+	window  time.Duration
+	buckets map[string]rateLimitBucket
+}
+
+func RateLimitPerMinute(limit int) gin.HandlerFunc {
+	if limit <= 0 {
+		return func(c *gin.Context) {
+			c.Next()
+		}
+	}
+	limiter := &rateLimiter{
+		limit:   limit,
+		window:  time.Minute,
+		buckets: make(map[string]rateLimitBucket),
+	}
+	return limiter.handle
+}
+
+func (l *rateLimiter) handle(c *gin.Context) {
+	now := time.Now()
+	key := c.ClientIP()
+	allowed, resetAt := l.allow(key, now)
+	if !allowed {
+		c.Header("Retry-After", retryAfterSeconds(now, resetAt))
+		c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
+			"code":    "rate_limited",
+			"message": "请求过于频繁，请稍后再试",
+		})
+		return
+	}
+	c.Next()
+}
+
+func (l *rateLimiter) allow(key string, now time.Time) (bool, time.Time) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	l.cleanup(now)
+	bucket := l.buckets[key]
+	if bucket.resetAt.IsZero() || !now.Before(bucket.resetAt) {
+		bucket = rateLimitBucket{resetAt: now.Add(l.window)}
+	}
+	bucket.count++
+	bucket.lastSeen = now
+	l.buckets[key] = bucket
+	return bucket.count <= l.limit, bucket.resetAt
+}
+
+func (l *rateLimiter) cleanup(now time.Time) {
+	for key, bucket := range l.buckets {
+		if now.Sub(bucket.lastSeen) > 2*l.window {
+			delete(l.buckets, key)
+		}
+	}
+}
+
+func retryAfterSeconds(now time.Time, resetAt time.Time) string {
+	seconds := int(resetAt.Sub(now).Seconds())
+	if seconds < 1 {
+		seconds = 1
+	}
+	return strconv.Itoa(seconds)
+}
@@ -338,7 +338,7 @@ func (c lakalaChannel) VerifyNotify(params map[string]string, rawPayload string,
 func buildChannelClient(dto *paymentconfig.ConfigDTO) (channelClient, error) {
 	switch strings.ToLower(strings.TrimSpace(dto.Provider)) {
 	case "leshua":
-		return newLeshuaChannel(leshua.Config{
+		return withChannelBreaker("leshua", newLeshuaChannel(leshua.Config{
 			GatewayURL: dto.GatewayURL,
 			MerchantID: dto.MerchantID,
 			SignKey:    dto.SignKey,
@@ -348,9 +348,9 @@ func buildChannelClient(dto *paymentconfig.ConfigDTO) (channelClient, error) {
 			PayWay:     firstNonEmpty(dto.PayWay, "ZFBZF"),
 			JSPayFlag:  firstNonEmpty(dto.JSPayFlag, "2"),
 			SignType:   firstNonEmpty(dto.SignType, "MD5"),
-		}), nil
+		})), nil
 	case "lakala":
-		return newLakalaChannel(lakala.Config{
+		return withChannelBreaker("lakala", newLakalaChannel(lakala.Config{
 			GatewayURL:         dto.GatewayURL,
 			AppID:              extraString(dto.ExtraConfig, "app_id"),
 			SerialNo:           extraString(dto.ExtraConfig, "serial_no"),
@@ -365,7 +365,7 @@ func buildChannelClient(dto *paymentconfig.ConfigDTO) (channelClient, error) {
 			JSPayFlag:          firstNonEmpty(dto.JSPayFlag, "2"),
 			PayMode:            extraString(dto.ExtraConfig, "pay_mode"),
 			OrderExpireMinutes: extraInt(dto.ExtraConfig, "order_expire_minutes"),
-		}), nil
+		})), nil
 	case "mock":
 		return nil, nil
 	default:
@@ -0,0 +1,68 @@
+package payment
+
+import (
+	"context"
+	"time"
+
+	"github.com/sony/gobreaker/v2"
+)
+
+type breakerChannel struct {
+	next          channelClient
+	createPayment *gobreaker.CircuitBreaker[*channelCreatePaymentResponse]
+	queryPayment  *gobreaker.CircuitBreaker[*channelQueryPaymentResponse]
+	createRefund  *gobreaker.CircuitBreaker[*channelCreateRefundResponse]
+	queryRefund   *gobreaker.CircuitBreaker[*channelQueryRefundResponse]
+}
+
+func withChannelBreaker(name string, next channelClient) channelClient {
+	if next == nil {
+		return nil
+	}
+	return breakerChannel{
+		next:          next,
+		createPayment: newPaymentBreaker[*channelCreatePaymentResponse](name + ".create_payment"),
+		queryPayment:  newPaymentBreaker[*channelQueryPaymentResponse](name + ".query_payment"),
+		createRefund:  newPaymentBreaker[*channelCreateRefundResponse](name + ".create_refund"),
+		queryRefund:   newPaymentBreaker[*channelQueryRefundResponse](name + ".query_refund"),
+	}
+}
+
+func newPaymentBreaker[T any](name string) *gobreaker.CircuitBreaker[T] {
+	return gobreaker.NewCircuitBreaker[T](gobreaker.Settings{
+		Name:        name,
+		MaxRequests: 3,
+		Timeout:     10 * time.Second,
+		ReadyToTrip: func(counts gobreaker.Counts) bool {
+			return counts.ConsecutiveFailures >= 5
+		},
+	})
+}
+
+func (c breakerChannel) CreatePayment(ctx context.Context, req channelCreatePaymentRequest) (*channelCreatePaymentResponse, error) {
+	return c.createPayment.Execute(func() (*channelCreatePaymentResponse, error) {
+		return c.next.CreatePayment(ctx, req)
+	})
+}
+
+func (c breakerChannel) QueryPayment(ctx context.Context, thirdOrderID string, providerOrderID string) (*channelQueryPaymentResponse, error) {
+	return c.queryPayment.Execute(func() (*channelQueryPaymentResponse, error) {
+		return c.next.QueryPayment(ctx, thirdOrderID, providerOrderID)
+	})
+}
+
+func (c breakerChannel) CreateRefund(ctx context.Context, req channelCreateRefundRequest) (*channelCreateRefundResponse, error) {
+	return c.createRefund.Execute(func() (*channelCreateRefundResponse, error) {
+		return c.next.CreateRefund(ctx, req)
+	})
+}
+
+func (c breakerChannel) QueryRefund(ctx context.Context, req channelQueryRefundRequest) (*channelQueryRefundResponse, error) {
+	return c.queryRefund.Execute(func() (*channelQueryRefundResponse, error) {
+		return c.next.QueryRefund(ctx, req)
+	})
+}
+
+func (c breakerChannel) VerifyNotify(params map[string]string, rawPayload string, contentType string, authorization string) (channelVerifyNotifyResult, error) {
+	return c.next.VerifyNotify(params, rawPayload, contentType, authorization)
+}
@@ -59,6 +59,9 @@ func New(cfg config.Config, deps Dependencies, logger *zap.Logger) *gin.Engine {
 	if cfg.AppEnv != "production" {
 		engine.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler))
 	}
+	if cfg.RateLimit.Enabled {
+		engine.Use(middleware.RateLimitPerMinute(cfg.RateLimit.RequestsPerMinute))
+	}

 	jwtManager := auth.NewJWTManager(cfg.JWTSecret)
 	var userRepo *auth.UserRepository