package middleware
import (
	"net/http"
	"time"
	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)
var Limiter *rate.Limiter
// 定義一個(gè)中間件函數(shù)來進(jìn)行限流
func RateLimiterMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		if !Limiter.AllowN(time.Now(), 1) {
			c.JSON(http.StatusTooManyRequests, gin.H{"message": "Rate limit exceeded"})
			// 設(shè)置休眠和業(yè)務(wù)時(shí)長(zhǎng)一樣，為了更好從日志出看出規(guī)則
			time.Sleep(50 * time.Millisecond)
			c.Abort()
			return
		}
		c.Next()
	}
}

func main() {
	r := gin.Default()
	// 創(chuàng)建一個(gè)限流器，每秒允許最多10個(gè)請(qǐng)求
	middleware.Limiter = rate.NewLimiter(rate.Limit(10), 1)
	// 使用限流中間件
	r.Use(middleware.RateLimiterMiddleware())
	r.GET("/api/resource", func(c *gin.Context) {
		time.Sleep(50 * time.Millisecond)
		c.JSON(http.StatusOK, gin.H{"message": "Resource accessed"})
	})
	r.Run(":8080")
}

package middleware
import (
	"net/http"
	"sync"
	"time"
	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)
var IPLimiter *IPRateLimiter
func NewIPRateLimiter() *IPRateLimiter {
	return &IPRateLimiter{
		limiter: make(map[string]*rate.Limiter),
	}
}
type IPRateLimiter struct {
	mu      sync.Mutex
	limiter map[string]*rate.Limiter
}
func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
	i.mu.Lock()
	defer i.mu.Unlock()
	limiter, exists := i.limiter[ip]
	if !exists {
		limiter = rate.NewLimiter(2, 5) // 每秒2個(gè)請(qǐng)求，桶容量為5
		i.limiter[ip] = limiter
	}
	return limiter
}
// 定義一個(gè)中間件函數(shù)來進(jìn)行限流
func IPRateLimiterMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		limiter := IPLimiter.GetLimiter(ip)
		if !limiter.Allow() {
			c.JSON(http.StatusTooManyRequests, gin.H{"message": "Rate limit exceeded"})
			// 設(shè)置休眠和業(yè)務(wù)時(shí)長(zhǎng)一樣，為了更好從日志出看出規(guī)則
			time.Sleep(50 * time.Millisecond)
			c.Abort()
			return
		}
		c.Next()
	}
}

func main() {
	r := gin.Default()
	// 創(chuàng)建IP限流器
	middleware.IPLimiter = middleware.NewIPRateLimiter()
	// 使用限流中間件
	r.Use(middleware.IPRateLimiterMiddleware())
	r.GET("/api/resource", func(c *gin.Context) {
		time.Sleep(50 * time.Millisecond)
		c.JSON(http.StatusOK, gin.H{"message": "Resource accessed"})
	})
	r.Run(":8080")
}