Push code

ratelimit/ratelimit.go

@@ -0,0 +1,117 @@
+package ratelimit
+
+import (
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+type client struct {
+	tokens     int
+	lastRefill time.Time
+}
+
+type RateLimiter struct {
+	clients       map[string]*client
+	mu            sync.RWMutex
+	maxTokens     int
+	refillRate    time.Duration
+	tokensPerHit  int
+	cleanupTicker *time.Ticker
+}
+
+// NewRateLimiter creates a new rate limiter
+// maxTokens: maximum number of tokens a client can have
+// refillRate: how often to add tokens back
+// tokensPerHit: tokens consumed per request
+func NewRateLimiter(maxTokens int, refillRate time.Duration, tokensPerHit int) *RateLimiter {
+	rl := &RateLimiter{
+		clients:      make(map[string]*client),
+		maxTokens:    maxTokens,
+		refillRate:   refillRate,
+		tokensPerHit: tokensPerHit,
+	}
+
+	// Cleanup old clients every 5 minutes
+	rl.cleanupTicker = time.NewTicker(5 * time.Minute)
+	go rl.cleanupClients()
+
+	return rl
+}
+
+func (rl *RateLimiter) getClient(ip string) *client {
+	rl.mu.RLock()
+	c, exists := rl.clients[ip]
+	rl.mu.RUnlock()
+
+	if !exists {
+		rl.mu.Lock()
+		c = &client{
+			tokens:     rl.maxTokens,
+			lastRefill: time.Now(),
+		}
+		rl.clients[ip] = c
+		rl.mu.Unlock()
+	}
+
+	return c
+}
+
+func (rl *RateLimiter) refillTokens(c *client) {
+	now := time.Now()
+	elapsed := now.Sub(c.lastRefill)
+	tokensToAdd := int(elapsed / rl.refillRate)
+
+	if tokensToAdd > 0 {
+		c.tokens += tokensToAdd
+		if c.tokens > rl.maxTokens {
+			c.tokens = rl.maxTokens
+		}
+		c.lastRefill = now
+	}
+}
+
+func (rl *RateLimiter) cleanupClients() {
+	for range rl.cleanupTicker.C {
+		rl.mu.Lock()
+		now := time.Now()
+		for ip, c := range rl.clients {
+			if now.Sub(c.lastRefill) > 10*time.Minute {
+				delete(rl.clients, ip)
+			}
+		}
+		rl.mu.Unlock()
+	}
+}
+
+// Middleware returns a Gin middleware handler
+func (rl *RateLimiter) Middleware() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		ip := c.ClientIP()
+		client := rl.getClient(ip)
+
+		rl.mu.Lock()
+		rl.refillTokens(client)
+
+		if client.tokens >= rl.tokensPerHit {
+			client.tokens -= rl.tokensPerHit
+			rl.mu.Unlock()
+			c.Next()
+		} else {
+			rl.mu.Unlock()
+			c.JSON(http.StatusTooManyRequests, gin.H{
+				"error": "Rate limit exceeded. Please try again later.",
+			})
+			c.Abort()
+		}
+	}
+}
+
+// Stop stops the cleanup ticker
+func (rl *RateLimiter) Stop() {
+	if rl.cleanupTicker != nil {
+		rl.cleanupTicker.Stop()
+	}
+}