136 lines
3 KiB
Go
136 lines
3 KiB
Go
package proxy
|
|
|
|
import (
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// NewRateLimiter allocates a per-IP token bucket limiter.
|
|
//
|
|
// rl := proxy.NewRateLimiter(cfg.RateLimit)
|
|
func NewRateLimiter(config RateLimit) *RateLimiter {
|
|
return &RateLimiter{
|
|
cfg: config,
|
|
buckets: make(map[string]*tokenBucket),
|
|
banned: make(map[string]time.Time),
|
|
}
|
|
}
|
|
|
|
// Allow returns true if the IP address is permitted to open a new connection. Thread-safe.
|
|
//
|
|
// if rl.Allow(conn.RemoteAddr().String()) { proceed() }
|
|
func (rl *RateLimiter) Allow(ip string) bool {
|
|
if rl == nil || rl.cfg.MaxConnectionsPerMinute <= 0 {
|
|
return true
|
|
}
|
|
|
|
host := remoteHost(ip)
|
|
now := time.Now().UTC()
|
|
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
if bannedUntil, exists := rl.banned[host]; exists {
|
|
if bannedUntil.After(now) {
|
|
return false
|
|
}
|
|
delete(rl.banned, host)
|
|
}
|
|
|
|
bucket, exists := rl.buckets[host]
|
|
if !exists {
|
|
bucket = &tokenBucket{
|
|
tokens: rl.cfg.MaxConnectionsPerMinute,
|
|
lastRefill: now,
|
|
}
|
|
rl.buckets[host] = bucket
|
|
}
|
|
|
|
rl.refillBucket(bucket, now)
|
|
if bucket.tokens <= 0 {
|
|
if rl.cfg.BanDurationSeconds > 0 {
|
|
rl.banned[host] = now.Add(time.Duration(rl.cfg.BanDurationSeconds) * time.Second)
|
|
}
|
|
return false
|
|
}
|
|
|
|
bucket.tokens--
|
|
return true
|
|
}
|
|
|
|
// Tick removes expired ban entries and refills all token buckets. Called every second.
|
|
//
|
|
// rl.Tick()
|
|
func (rl *RateLimiter) Tick() {
|
|
if rl == nil || rl.cfg.MaxConnectionsPerMinute <= 0 {
|
|
return
|
|
}
|
|
|
|
now := time.Now().UTC()
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
for host, bannedUntil := range rl.banned {
|
|
if !bannedUntil.After(now) {
|
|
delete(rl.banned, host)
|
|
}
|
|
}
|
|
|
|
for _, bucket := range rl.buckets {
|
|
rl.refillBucket(bucket, now)
|
|
}
|
|
}
|
|
|
|
func (rl *RateLimiter) refillBucket(bucket *tokenBucket, now time.Time) {
|
|
if bucket == nil || rl.cfg.MaxConnectionsPerMinute <= 0 {
|
|
return
|
|
}
|
|
|
|
refillEvery := time.Minute / time.Duration(rl.cfg.MaxConnectionsPerMinute)
|
|
if refillEvery <= 0 {
|
|
refillEvery = time.Second
|
|
}
|
|
|
|
elapsed := now.Sub(bucket.lastRefill)
|
|
if elapsed < refillEvery {
|
|
return
|
|
}
|
|
|
|
tokensToAdd := int(elapsed / refillEvery)
|
|
bucket.tokens += tokensToAdd
|
|
if bucket.tokens > rl.cfg.MaxConnectionsPerMinute {
|
|
bucket.tokens = rl.cfg.MaxConnectionsPerMinute
|
|
}
|
|
bucket.lastRefill = bucket.lastRefill.Add(time.Duration(tokensToAdd) * refillEvery)
|
|
}
|
|
|
|
// NewCustomDiff stores the default custom difficulty override.
|
|
//
|
|
// cd := proxy.NewCustomDiff(50000)
|
|
func NewCustomDiff(globalDiff uint64) *CustomDiff {
|
|
return &CustomDiff{globalDiff: globalDiff}
|
|
}
|
|
|
|
// OnLogin parses miner.User for a "+{number}" suffix and sets miner.CustomDiff.
|
|
//
|
|
// cd.OnLogin(proxy.Event{Miner: miner})
|
|
func (cd *CustomDiff) OnLogin(event Event) {
|
|
if event.Miner == nil {
|
|
return
|
|
}
|
|
|
|
user := event.Miner.User()
|
|
index := strings.LastIndex(user, "+")
|
|
if index > 0 {
|
|
if value, errorValue := strconv.ParseUint(user[index+1:], 10, 64); errorValue == nil {
|
|
event.Miner.SetUser(user[:index])
|
|
event.Miner.SetCustomDiff(value)
|
|
return
|
|
}
|
|
}
|
|
|
|
if cd != nil && cd.globalDiff > 0 {
|
|
event.Miner.SetCustomDiff(cd.globalDiff)
|
|
}
|
|
}
|