Mining/pkg/mining/ratelimiter.go

package mining

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

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

// RateLimiter provides token bucket rate limiting per IP address
type RateLimiter struct {
	requestsPerSecond int
	burst             int
	clients           map[string]*rateLimitClient
	mu                sync.RWMutex
	stopChan          chan struct{}
	stopped           bool
}

type rateLimitClient struct {
	tokens    float64
	lastCheck time.Time
}

// NewRateLimiter creates a new rate limiter with the specified limits
func NewRateLimiter(requestsPerSecond, burst int) *RateLimiter {
	rl := &RateLimiter{
		requestsPerSecond: requestsPerSecond,
		burst:             burst,
		clients:           make(map[string]*rateLimitClient),
		stopChan:          make(chan struct{}),
	}

	// Start cleanup goroutine
	go rl.cleanupLoop()

	return rl
}

// cleanupLoop removes stale clients periodically
func (rl *RateLimiter) cleanupLoop() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-rl.stopChan:
			return
		case <-ticker.C:
			rl.cleanup()
		}
	}
}

// cleanup removes clients that haven't made requests in 5 minutes
func (rl *RateLimiter) cleanup() {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	for ip, c := range rl.clients {
		if time.Since(c.lastCheck) > 5*time.Minute {
			delete(rl.clients, ip)
		}
	}
}

// Stop stops the rate limiter's cleanup goroutine
func (rl *RateLimiter) Stop() {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	if !rl.stopped {
		close(rl.stopChan)
		rl.stopped = true
	}
}

// Middleware returns a Gin middleware handler for rate limiting
func (rl *RateLimiter) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()

		rl.mu.Lock()
		cl, exists := rl.clients[ip]
		if !exists {
			cl = &rateLimitClient{tokens: float64(rl.burst), lastCheck: time.Now()}
			rl.clients[ip] = cl
		}

		// Token bucket algorithm
		now := time.Now()
		elapsed := now.Sub(cl.lastCheck).Seconds()
		cl.tokens += elapsed * float64(rl.requestsPerSecond)
		if cl.tokens > float64(rl.burst) {
			cl.tokens = float64(rl.burst)
		}
		cl.lastCheck = now

		if cl.tokens < 1 {
			rl.mu.Unlock()
			respondWithError(c, http.StatusTooManyRequests, "RATE_LIMITED",
				"too many requests", "rate limit exceeded")
			c.Abort()
			return
		}

		cl.tokens--
		rl.mu.Unlock()
		c.Next()
	}
}

// ClientCount returns the number of tracked clients (for testing/monitoring)
func (rl *RateLimiter) ClientCount() int {
	rl.mu.RLock()
	defer rl.mu.RUnlock()
	return len(rl.clients)
}