iloom/gateway/internal/middleware/ratelimit.go

package middleware

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

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

type bucket struct {
	tokens    float64
	lastCheck time.Time
	mu        sync.Mutex
}

var buckets sync.Map

func RateLimit(rps int) gin.HandlerFunc {
	rate := float64(rps)

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

		val, _ := buckets.LoadOrStore(ip, &bucket{
			tokens:    rate,
			lastCheck: time.Now(),
		})
		b := val.(*bucket)

		b.mu.Lock()
		now := time.Now()
		elapsed := now.Sub(b.lastCheck).Seconds()
		b.lastCheck = now
		b.tokens += elapsed * rate
		if b.tokens > rate {
			b.tokens = rate
		}

		if b.tokens < 1 {
			b.mu.Unlock()
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"code":    42900,
				"message": "rate limit exceeded",
			})
			return
		}

		b.tokens--
		b.mu.Unlock()

		c.Next()
	}
}