121 lines
2.9 KiB
Go
121 lines
2.9 KiB
Go
package mining
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// limiter := NewRateLimiter(10, 20) // 10 requests per second with a burst of 20.
|
|
// router.Use(limiter.Middleware())
|
|
type RateLimiter struct {
|
|
requestsPerSecond int
|
|
burst int
|
|
clients map[string]*rateLimitClient
|
|
mutex sync.RWMutex
|
|
stopChannel chan struct{}
|
|
stopped bool
|
|
}
|
|
|
|
type rateLimitClient struct {
|
|
tokens float64
|
|
lastCheck time.Time
|
|
}
|
|
|
|
// limiter := NewRateLimiter(10, 20) // 10 requests/second, burst of 20
|
|
// defer limiter.Stop()
|
|
func NewRateLimiter(requestsPerSecond, burst int) *RateLimiter {
|
|
limiter := &RateLimiter{
|
|
requestsPerSecond: requestsPerSecond,
|
|
burst: burst,
|
|
clients: make(map[string]*rateLimitClient),
|
|
stopChannel: make(chan struct{}),
|
|
}
|
|
|
|
// Start cleanup goroutine
|
|
go limiter.cleanupLoop()
|
|
|
|
return limiter
|
|
}
|
|
|
|
// go limiter.cleanupLoop() // started by NewRateLimiter; runs until limiter.Stop()
|
|
func (limiter *RateLimiter) cleanupLoop() {
|
|
ticker := time.NewTicker(time.Minute)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-limiter.stopChannel:
|
|
return
|
|
case <-ticker.C:
|
|
limiter.cleanup()
|
|
}
|
|
}
|
|
}
|
|
|
|
// limiter.cleanup() // called every minute by cleanupLoop; evicts IPs idle for >5 minutes
|
|
func (limiter *RateLimiter) cleanup() {
|
|
limiter.mutex.Lock()
|
|
defer limiter.mutex.Unlock()
|
|
|
|
for ip, client := range limiter.clients {
|
|
if time.Since(client.lastCheck) > 5*time.Minute {
|
|
delete(limiter.clients, ip)
|
|
}
|
|
}
|
|
}
|
|
|
|
// limiter.Stop() // call on shutdown to release the cleanup goroutine
|
|
func (limiter *RateLimiter) Stop() {
|
|
limiter.mutex.Lock()
|
|
defer limiter.mutex.Unlock()
|
|
|
|
if !limiter.stopped {
|
|
close(limiter.stopChannel)
|
|
limiter.stopped = true
|
|
}
|
|
}
|
|
|
|
// router.Use(limiter.Middleware()) // install before route handlers
|
|
func (limiter *RateLimiter) Middleware() gin.HandlerFunc {
|
|
return func(requestContext *gin.Context) {
|
|
clientAddress := requestContext.ClientIP()
|
|
|
|
limiter.mutex.Lock()
|
|
client, exists := limiter.clients[clientAddress]
|
|
if !exists {
|
|
client = &rateLimitClient{tokens: float64(limiter.burst), lastCheck: time.Now()}
|
|
limiter.clients[clientAddress] = client
|
|
}
|
|
|
|
// Token bucket algorithm
|
|
now := time.Now()
|
|
elapsed := now.Sub(client.lastCheck).Seconds()
|
|
client.tokens += elapsed * float64(limiter.requestsPerSecond)
|
|
if client.tokens > float64(limiter.burst) {
|
|
client.tokens = float64(limiter.burst)
|
|
}
|
|
client.lastCheck = now
|
|
|
|
if client.tokens < 1 {
|
|
limiter.mutex.Unlock()
|
|
respondWithError(requestContext, http.StatusTooManyRequests, "RATE_LIMITED",
|
|
"too many requests", "rate limit exceeded")
|
|
requestContext.Abort()
|
|
return
|
|
}
|
|
|
|
client.tokens--
|
|
limiter.mutex.Unlock()
|
|
requestContext.Next()
|
|
}
|
|
}
|
|
|
|
// if limiter.ClientCount() == 0 { /* no active clients */ }
|
|
func (limiter *RateLimiter) ClientCount() int {
|
|
limiter.mutex.RLock()
|
|
defer limiter.mutex.RUnlock()
|
|
return len(limiter.clients)
|
|
}
|