TGSS/internal/server/ratelimiter.go

package server

import (
	"context"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/fx"
	"go.uber.org/zap"
	"golang.org/x/time/rate"
)

type ClientLimiter struct {
	feedLimiter  *rate.Limiter
	imageLimiter *rate.Limiter
	lastSeen     time.Time
}

type RateLimiter struct {
	logger  *zap.Logger
	clients map[string]*ClientLimiter
	mu      sync.Mutex
}

func NewRateLimiter(logger *zap.Logger) *RateLimiter {
	return &RateLimiter{logger: logger, mu: sync.Mutex{}, clients: map[string]*ClientLimiter{}}
}

func (r *RateLimiter) getLimiter(ip string, kind string) *rate.Limiter {
	r.mu.Lock()
	defer r.mu.Unlock()

	if c, ok := r.clients[ip]; ok {
		c.lastSeen = time.Now()
		if kind == "image" {
			return c.imageLimiter
		}
		return c.feedLimiter
	}

	feedLim := rate.NewLimiter(rate.Every(3*time.Second), 5)
	imageLim := rate.NewLimiter(rate.Every(5*time.Second), 4) // TODO: get from env

	r.clients[ip] = &ClientLimiter{
		feedLimiter:  feedLim,
		imageLimiter: imageLim,
		lastSeen:     time.Now(),
	}
	if kind == "image" {
		return imageLim
	}
	return feedLim
}

func (r *RateLimiter) CleanupRateLimiter() {
	r.mu.Lock()
	defer r.mu.Unlock()

	for ip, clientLimiter := range r.clients {
		if time.Since(clientLimiter.lastSeen) > time.Hour { // Delete ip ratelimiters older that 1 hour
			delete(r.clients, ip)
		}

	}

}

func RegisterRateLimiterCleanup(lc fx.Lifecycle, logger *zap.Logger, r *RateLimiter) {
	var ticker *time.Ticker
	var stopChan chan struct{}

	lc.Append(fx.Hook{
		OnStart: func(ctx context.Context) error {
			logger.Info("Starting rate limiter cleanup task", zap.Duration("interval", 30*time.Minute))
			ticker = time.NewTicker(30 * time.Minute)
			stopChan = make(chan struct{})
			go func() {
				for {
					select {
					case <-ticker.C:
						logger.Info("Running rate limiter cleanup")
						r.CleanupRateLimiter()
						logger.Info("Rate limiter cleanup completed")
					case <-stopChan:
						return
					}
				}
			}()
			return nil
		},
		OnStop: func(ctx context.Context) error {
			logger.Info("Stopping rate limiter cleanup task")
			if ticker != nil {
				ticker.Stop()
			}
			if stopChan != nil {
				close(stopChan)
			}
			return nil
		},
	})
}

func (r *RateLimiter) FeedRateLimit() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		lim := r.getLimiter(ip, "feed")
		if !lim.Allow() {
			c.AbortWithStatusJSON(429, gin.H{"error": "Slow down pls"})
			return
		}
		c.Next()
	}
}

func (r *RateLimiter) ImageRateLimit() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		lim := r.getLimiter(ip, "image")
		if !lim.Allow() {
			c.AbortWithStatusJSON(429, gin.H{"error": "Slow down pls, these arent your mothers nudes"})
			return
		}
		c.Next()
	}
}