go-proxy/runtime_support.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),
	}
}

// SetConfig replaces the active rate-limit settings.
//
//	rl.SetConfig(proxy.RateLimit{MaxConnectionsPerMinute: 30, BanDurationSeconds: 300})
func (rl *RateLimiter) SetConfig(config RateLimit) {
	if rl == nil {
		return
	}

	rl.mu.Lock()
	rl.cfg = config
	rl.mu.Unlock()
}

// 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 {
		return true
	}

	host := remoteHost(ip)
	now := time.Now().UTC()

	rl.mu.Lock()
	defer rl.mu.Unlock()

	if rl.cfg.MaxConnectionsPerMinute <= 0 {
		return true
	}

	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 {
		return
	}

	now := time.Now().UTC()
	rl.mu.Lock()
	defer rl.mu.Unlock()

	if rl.cfg.MaxConnectionsPerMinute <= 0 {
		return
	}

	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}
}

// SetGlobalDiff updates the default custom difficulty override.
//
//	cd.SetGlobalDiff(100000)
func (cd *CustomDiff) SetGlobalDiff(globalDiff uint64) {
	if cd == nil {
		return
	}

	cd.mu.Lock()
	cd.globalDiff = globalDiff
	cd.mu.Unlock()
}

// 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 && index < len(user)-1 {
		if value, errorValue := strconv.ParseUint(user[index+1:], 10, 64); errorValue == nil {
			event.Miner.SetUser(user[:index])
			event.Miner.SetCustomDiff(value)
			return
		}
		return
	}

	if cd == nil {
		return
	}

	cd.mu.RLock()
	globalDiff := cd.globalDiff
	cd.mu.RUnlock()
	if globalDiff > 0 {
		event.Miner.SetCustomDiff(globalDiff)
	}
}