dyn/internal/middleware/ratelimit.go

package middleware

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

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

type RateLimiter struct {
	ipLimits      map[string]*RateLimitEntry
	tokenLimits   map[string]*RateLimitEntry
	mu            sync.RWMutex
	limitPerIP    int
	limitPerToken int
}

type RateLimitEntry struct {
	Count     int
	ResetTime time.Time
}

func NewRateLimiter(perIP, perToken int) *RateLimiter {
	return &RateLimiter{
		ipLimits:      make(map[string]*RateLimitEntry),
		tokenLimits:   make(map[string]*RateLimitEntry),
		limitPerIP:    perIP,
		limitPerToken: perToken,
	}
}

func (rl *RateLimiter) RateLimitByIP() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()

		rl.mu.Lock()
		entry, exists := rl.ipLimits[ip]
		now := time.Now()

		if !exists || now.After(entry.ResetTime) {
			rl.ipLimits[ip] = &RateLimitEntry{
				Count:     1,
				ResetTime: now.Add(time.Minute),
			}
			rl.mu.Unlock()
			c.Next()
			return
		}

		if entry.Count >= rl.limitPerIP {
			rl.mu.Unlock()
			c.String(http.StatusTooManyRequests, "rate limit exceeded")
			c.Abort()
			return
		}

		entry.Count++
		rl.mu.Unlock()
		c.Next()
	}
}

func (rl *RateLimiter) RateLimitByToken() gin.HandlerFunc {
	return func(c *gin.Context) {
		token := extractToken(c)
		if token == "" {
			c.Next()
			return
		}

		rl.mu.Lock()
		entry, exists := rl.tokenLimits[token]
		now := time.Now()

		if !exists || now.After(entry.ResetTime) {
			rl.tokenLimits[token] = &RateLimitEntry{
				Count:     1,
				ResetTime: now.Add(time.Minute),
			}
			rl.mu.Unlock()
			c.Next()
			return
		}

		if entry.Count >= rl.limitPerToken {
			rl.mu.Unlock()
			c.String(http.StatusTooManyRequests, "rate limit exceeded")
			c.Abort()
			return
		}

		entry.Count++
		rl.mu.Unlock()
		c.Next()
	}
}

func extractToken(c *gin.Context) string {
	auth := c.GetHeader("Authorization")
	if auth == "" {
		return ""
	}

	parts := strings.SplitN(auth, " ", 2)
	if len(parts) != 2 || parts[0] != "Basic" {
		return ""
	}

	return parts[1]
}