76 lines
1.3 KiB
Go
76 lines
1.3 KiB
Go
package ratelimiter
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
type TokenBucketRateLimiter struct {
|
|
sync.RWMutex
|
|
clients map[string]*Client
|
|
config Config
|
|
}
|
|
|
|
func NewTokenBucketRateLimiter(config Config) *TokenBucketRateLimiter {
|
|
rl := &TokenBucketRateLimiter{
|
|
clients: make(map[string]*Client),
|
|
config: config,
|
|
}
|
|
|
|
go rl.cleanupClients()
|
|
|
|
return rl
|
|
}
|
|
|
|
type Client struct {
|
|
Limiter *rate.Limiter
|
|
LastSeen time.Time
|
|
}
|
|
|
|
func (rl *TokenBucketRateLimiter) getClient(ip string) *Client {
|
|
rl.RLock()
|
|
_, exists := rl.clients[ip]
|
|
rl.RUnlock()
|
|
|
|
if !exists {
|
|
limiter := rate.NewLimiter(rate.Every(rl.config.TimeFrame), rl.config.RequestsPerTimeFrame)
|
|
rl.Lock()
|
|
rl.clients[ip] = &Client{
|
|
Limiter: limiter,
|
|
LastSeen: time.Now()}
|
|
|
|
rl.Unlock()
|
|
}
|
|
|
|
client := rl.clients[ip]
|
|
client.LastSeen = time.Now()
|
|
|
|
return client
|
|
}
|
|
|
|
func (rl *TokenBucketRateLimiter) Allow(ip string) (bool, float64) {
|
|
client := rl.getClient(ip)
|
|
|
|
allowed := client.Limiter.Allow()
|
|
tokens := client.Limiter.Tokens()
|
|
|
|
return allowed, tokens
|
|
}
|
|
|
|
func (rl *TokenBucketRateLimiter) cleanupClients() {
|
|
for {
|
|
time.Sleep(time.Minute)
|
|
//log cleaning up clients
|
|
|
|
rl.Lock()
|
|
for ip, client := range rl.clients {
|
|
if time.Since(client.LastSeen) > 3*time.Minute { // timeout period
|
|
delete(rl.clients, ip)
|
|
}
|
|
}
|
|
rl.Unlock()
|
|
}
|
|
}
|