follows_throttle.go raw

   1  package acl
   2  
   3  import (
   4  	"sync"
   5  	"time"
   6  )
   7  
   8  // ThrottleState tracks accumulated delay for an identity (IP or pubkey)
   9  type ThrottleState struct {
  10  	AccumulatedDelay time.Duration
  11  	LastEventTime    time.Time
  12  }
  13  
  14  // ProgressiveThrottle implements linear delay with time decay.
  15  // Each event adds perEvent delay, and delay decays at 1:1 ratio with elapsed time.
  16  // This creates a natural rate limit that averages to 1 event per perEvent interval.
  17  type ProgressiveThrottle struct {
  18  	mu           sync.Mutex
  19  	ipStates     map[string]*ThrottleState
  20  	pubkeyStates map[string]*ThrottleState
  21  	perEvent     time.Duration // delay increment per event (default 200ms)
  22  	maxDelay     time.Duration // cap (default 60s)
  23  }
  24  
  25  // NewProgressiveThrottle creates a new throttle with the given parameters.
  26  // perEvent is the delay added per event (e.g., 200ms).
  27  // maxDelay is the maximum accumulated delay cap (e.g., 60s).
  28  func NewProgressiveThrottle(perEvent, maxDelay time.Duration) *ProgressiveThrottle {
  29  	return &ProgressiveThrottle{
  30  		ipStates:     make(map[string]*ThrottleState),
  31  		pubkeyStates: make(map[string]*ThrottleState),
  32  		perEvent:     perEvent,
  33  		maxDelay:     maxDelay,
  34  	}
  35  }
  36  
  37  // GetDelay returns accumulated delay for this identity and updates state.
  38  // It tracks both IP and pubkey independently and returns the maximum of both.
  39  // This prevents evasion via different pubkeys from same IP or vice versa.
  40  func (pt *ProgressiveThrottle) GetDelay(ip, pubkeyHex string) time.Duration {
  41  	pt.mu.Lock()
  42  	defer pt.mu.Unlock()
  43  
  44  	now := time.Now()
  45  	var ipDelay, pubkeyDelay time.Duration
  46  
  47  	if ip != "" {
  48  		ipDelay = pt.updateState(pt.ipStates, ip, now)
  49  	}
  50  	if pubkeyHex != "" {
  51  		pubkeyDelay = pt.updateState(pt.pubkeyStates, pubkeyHex, now)
  52  	}
  53  
  54  	// Return max of both to prevent evasion
  55  	if ipDelay > pubkeyDelay {
  56  		return ipDelay
  57  	}
  58  	return pubkeyDelay
  59  }
  60  
  61  // updateState calculates and updates the delay for a single identity.
  62  // The algorithm:
  63  // 1. Decay: subtract elapsed time from accumulated delay (1:1 ratio)
  64  // 2. Add: add perEvent for this new event
  65  // 3. Cap: limit to maxDelay
  66  func (pt *ProgressiveThrottle) updateState(states map[string]*ThrottleState, key string, now time.Time) time.Duration {
  67  	state, exists := states[key]
  68  	if !exists {
  69  		// First event from this identity
  70  		states[key] = &ThrottleState{
  71  			AccumulatedDelay: pt.perEvent,
  72  			LastEventTime:    now,
  73  		}
  74  		return pt.perEvent
  75  	}
  76  
  77  	// Decay: subtract elapsed time (1:1 ratio)
  78  	elapsed := now.Sub(state.LastEventTime)
  79  	state.AccumulatedDelay -= elapsed
  80  	if state.AccumulatedDelay < 0 {
  81  		state.AccumulatedDelay = 0
  82  	}
  83  
  84  	// Add new event's delay
  85  	state.AccumulatedDelay += pt.perEvent
  86  	state.LastEventTime = now
  87  
  88  	// Cap at max
  89  	if state.AccumulatedDelay > pt.maxDelay {
  90  		state.AccumulatedDelay = pt.maxDelay
  91  	}
  92  
  93  	return state.AccumulatedDelay
  94  }
  95  
  96  // Cleanup removes entries that have fully decayed (no remaining delay).
  97  // This should be called periodically to prevent unbounded memory growth.
  98  func (pt *ProgressiveThrottle) Cleanup() {
  99  	pt.mu.Lock()
 100  	defer pt.mu.Unlock()
 101  
 102  	now := time.Now()
 103  
 104  	// Remove IP entries that have fully decayed
 105  	for k, v := range pt.ipStates {
 106  		elapsed := now.Sub(v.LastEventTime)
 107  		if elapsed >= v.AccumulatedDelay {
 108  			delete(pt.ipStates, k)
 109  		}
 110  	}
 111  
 112  	// Remove pubkey entries that have fully decayed
 113  	for k, v := range pt.pubkeyStates {
 114  		elapsed := now.Sub(v.LastEventTime)
 115  		if elapsed >= v.AccumulatedDelay {
 116  			delete(pt.pubkeyStates, k)
 117  		}
 118  	}
 119  }
 120  
 121  // Stats returns the current number of tracked IPs and pubkeys (for monitoring)
 122  func (pt *ProgressiveThrottle) Stats() (ipCount, pubkeyCount int) {
 123  	pt.mu.Lock()
 124  	defer pt.mu.Unlock()
 125  	return len(pt.ipStates), len(pt.pubkeyStates)
 126  }
 127