backend/internal/middleware/ratelimit.go

package middleware

import (
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
)

type RateLimiter struct {
	requests map[string][]time.Time
	mu       sync.Mutex
	limit    int
	window   time.Duration
}

func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
	rl := &RateLimiter{
		requests: make(map[string][]time.Time),
		limit:    limit,
		window:   window,
	}

	go func() {
		for {
			time.Sleep(window)
			rl.cleanup()
		}
	}()

	return rl
}

func (rl *RateLimiter) cleanup() {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	for key, times := range rl.requests {
		var valid []time.Time
		for _, t := range times {
			if now.Sub(t) < rl.window {
				valid = append(valid, t)
			}
		}
		if len(valid) == 0 {
			delete(rl.requests, key)
		} else {
			rl.requests[key] = valid
		}
	}
}

func (rl *RateLimiter) isAllowed(key string) bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	times := rl.requests[key]

	var valid []time.Time
	for _, t := range times {
		if now.Sub(t) < rl.window {
			valid = append(valid, t)
		}
	}

	if len(valid) >= rl.limit {
		rl.requests[key] = valid
		return false
	}

	rl.requests[key] = append(valid, now)
	return true
}

func (rl *RateLimiter) count(key string) int {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	times := rl.requests[key]

	var valid []time.Time
	for _, t := range times {
		if now.Sub(t) < rl.window {
			valid = append(valid, t)
		}
	}
	rl.requests[key] = valid
	return len(valid)
}

func RateLimit(requestsPerMinute int) gin.HandlerFunc {
	return RateLimitWithDuration(requestsPerMinute, time.Minute)
}

func RateLimitWithDuration(limit int, window time.Duration) gin.HandlerFunc {
	limiter := NewRateLimiter(limit, window)

	return func(c *gin.Context) {
		ip := c.ClientIP()

		if !limiter.isAllowed(ip) {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"code":    429,
				"message": "too many requests, please try again later",
			})
			c.Abort()
			return
		}

		c.Next()
	}
}

func RateLimitWithKey(requestsPerMinute int, keyFunc func(c *gin.Context) string) gin.HandlerFunc {
	return RateLimitWithDurationAndKey(requestsPerMinute, time.Minute, keyFunc)
}

func RateLimitWithDurationAndKey(limit int, window time.Duration, keyFunc func(c *gin.Context) string) gin.HandlerFunc {
	limiter := NewRateLimiter(limit, window)

	return func(c *gin.Context) {
		key := keyFunc(c)
		if key == "" {
			key = c.ClientIP()
		}

		if !limiter.isAllowed(key) {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"code":    429,
				"message": "too many requests, please try again later",
			})
			c.Abort()
			return
		}

		c.Next()
	}
}

// ============================================================
// IP 封禁与登录失败追踪
// ============================================================

const (
	DefaultMaxLoginFailures = 10
	DefaultBanDuration      = 15 * time.Minute
	DefaultFailureWindow    = 15 * time.Minute
)

type ipBanEntry struct {
	unbanAt time.Time
}

type loginFailureEntry struct {
	count    int
	firstAt  time.Time
	lastAt   time.Time
}

type IPBanStore struct {
	mu       sync.RWMutex
	bans     map[string]*ipBanEntry
	failures map[string]*loginFailureEntry

	maxFailures int
	banDuration time.Duration
	failWindow  time.Duration
}

var defaultIPBanStore = NewIPBanStore(DefaultMaxLoginFailures, DefaultBanDuration, DefaultFailureWindow)

func NewIPBanStore(maxFailures int, banDuration, failWindow time.Duration) *IPBanStore {
	store := &IPBanStore{
		bans:        make(map[string]*ipBanEntry),
		failures:    make(map[string]*loginFailureEntry),
		maxFailures: maxFailures,
		banDuration: banDuration,
		failWindow:  failWindow,
	}

	go store.cleanupLoop()

	return store
}

func (s *IPBanStore) cleanupLoop() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for range ticker.C {
		s.cleanup()
	}
}

func (s *IPBanStore) cleanup() {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()

	for ip, entry := range s.bans {
		if now.After(entry.unbanAt) {
			delete(s.bans, ip)
		}
	}

	for ip, entry := range s.failures {
		if now.Sub(entry.firstAt) > s.failWindow && now.Sub(entry.lastAt) > s.failWindow {
			delete(s.failures, ip)
		}
	}
}

func (s *IPBanStore) IsBanned(ip string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	entry, exists := s.bans[ip]
	if !exists {
		return false
	}

	return time.Now().Before(entry.unbanAt)
}

func (s *IPBanStore) RecordFailure(ip string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()

	entry, exists := s.failures[ip]
	if !exists || now.Sub(entry.firstAt) > s.failWindow {
		entry = &loginFailureEntry{
			count:   1,
			firstAt: now,
			lastAt:  now,
		}
		s.failures[ip] = entry
		return false
	}

	entry.count++
	entry.lastAt = now

	if entry.count >= s.maxFailures {
		s.bans[ip] = &ipBanEntry{
			unbanAt: now.Add(s.banDuration),
		}
		delete(s.failures, ip)
		return true
	}

	return false
}

func (s *IPBanStore) ResetFailures(ip string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	delete(s.failures, ip)
}

func RecordLoginFailure(ip string) bool {
	return defaultIPBanStore.RecordFailure(ip)
}

func ResetLoginFailures(ip string) {
	defaultIPBanStore.ResetFailures(ip)
}

func IPBanMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()

		if defaultIPBanStore.IsBanned(ip) {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"code":    429,
				"message": "your IP has been temporarily banned due to too many failed attempts, please try again later",
			})
			c.Abort()
			return
		}

		c.Next()
	}
}

// ============================================================
// 预配置的速率限制中间件
// ============================================================

func LoginRateLimit() gin.HandlerFunc {
	return RateLimitWithDuration(10, time.Minute)
}

func RegisterRateLimit() gin.HandlerFunc {
	return RateLimitWithDuration(5, time.Hour)
}

func PasswordResetRateLimit() gin.HandlerFunc {
	return RateLimitWithDuration(3, time.Hour)
}

func SendCodeRateLimit() gin.HandlerFunc {
	return RateLimitWithDuration(5, time.Hour)
}

func GlobalAuthRateLimit() gin.HandlerFunc {
	return RateLimitWithDuration(30, time.Minute)
}