LokiStack/backend/middleware/ratelimit.go

package middleware

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

	"lokistack/utils"

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

// ipLimiter 内存滑动窗口 IP 限流（简化实现）。
// 用于登录防爆破：超限返回 429。单实例部署下内存实现即可，无需外部存储。
type ipLimiter struct {
	mu      sync.Mutex
	window  time.Duration
	max     int
	hits    map[string][]time.Time
	lastSweep time.Time
}

func newIPLimiter(max int, window time.Duration) *ipLimiter {
	return &ipLimiter{
		window:  window,
		max:     max,
		hits:    make(map[string][]time.Time),
	}
}

// Allow 记录一次命中并判断是否放行
func (l *ipLimiter) Allow(ip string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := time.Now()
	// 惰性清理过期记录，防止 map 无限增长
	if now.Sub(l.lastSweep) > time.Minute {
		for k, v := range l.hits {
			kept := v[:0]
			for _, t := range v {
				if now.Sub(t) < l.window {
					kept = append(kept, t)
				}
			}
			if len(kept) == 0 {
				delete(l.hits, k)
			} else {
				l.hits[k] = kept
			}
		}
		l.lastSweep = now
	}

	history := l.hits[ip]
	kept := history[:0]
	for _, t := range history {
		if now.Sub(t) < l.window {
			kept = append(kept, t)
		}
	}
	if len(kept) >= l.max {
		l.hits[ip] = kept
		return false
	}
	l.hits[ip] = append(kept, now)
	return true
}

// LoginRateLimit 登录接口 IP 限流：max 次 / window 时长
func LoginRateLimit(max int, window time.Duration) gin.HandlerFunc {
	limiter := newIPLimiter(max, window)
	return func(c *gin.Context) {
		if !limiter.Allow(c.ClientIP()) {
			c.JSON(http.StatusTooManyRequests, utils.Response{
				Code:    http.StatusTooManyRequests,
				Message: "登录尝试过于频繁，请稍后再试",
			})
			c.Abort()
			return
		}
		c.Next()
	}
}