github.com/MetalBlockchain/metalgo@v1.11.9/network/throttling/inbound_conn_upgrade_throttler.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package throttling
     5  
     6  import (
     7  	"net/netip"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/MetalBlockchain/metalgo/utils/logging"
    12  	"github.com/MetalBlockchain/metalgo/utils/set"
    13  	"github.com/MetalBlockchain/metalgo/utils/timer/mockable"
    14  )
    15  
    16  var (
    17  	_ InboundConnUpgradeThrottler = (*inboundConnUpgradeThrottler)(nil)
    18  	_ InboundConnUpgradeThrottler = (*noInboundConnUpgradeThrottler)(nil)
    19  )
    20  
    21  // InboundConnUpgradeThrottler returns whether we should upgrade an inbound connection from IP [ipStr].
    22  // If ShouldUpgrade(ipStr) returns false, the connection to that IP should be closed.
    23  // Note that InboundConnUpgradeThrottler rate-limits _upgrading_ of
    24  // inbound connections, whereas throttledListener rate-limits
    25  // _acceptance_ of inbound connections.
    26  type InboundConnUpgradeThrottler interface {
    27  	// Dispatch starts this InboundConnUpgradeThrottler.
    28  	// Must be called before [ShouldUpgrade].
    29  	// Blocks until [Stop] is called (i.e. should be called in a goroutine.)
    30  	Dispatch()
    31  	// Stop this InboundConnUpgradeThrottler and causes [Dispatch] to return.
    32  	// Should be called when we're done with this InboundConnUpgradeThrottler.
    33  	// This InboundConnUpgradeThrottler must not be used after [Stop] is called.
    34  	Stop()
    35  	// Returns whether we should upgrade an inbound connection from [ipStr].
    36  	// Must only be called after [Dispatch] has been called.
    37  	// If [ip] is a local IP, this method always returns true.
    38  	// Must not be called after [Stop] has been called.
    39  	ShouldUpgrade(ip netip.AddrPort) bool
    40  }
    41  
    42  type InboundConnUpgradeThrottlerConfig struct {
    43  	// ShouldUpgrade(ipStr) returns true if it has been at least [UpgradeCooldown]
    44  	// since the last time ShouldUpgrade(ipStr) returned true or if
    45  	// ShouldUpgrade(ipStr) has never been called.
    46  	// If <= 0, inbound connections not rate-limited.
    47  	UpgradeCooldown time.Duration `json:"upgradeCooldown"`
    48  	// Maximum number of inbound connections upgraded within [UpgradeCooldown].
    49  	// (As implemented in inboundConnUpgradeThrottler, may actually upgrade
    50  	// [MaxRecentConnsUpgraded+1] due to a race condition but that's fine.)
    51  	// If <= 0, inbound connections not rate-limited.
    52  	MaxRecentConnsUpgraded int `json:"maxRecentConnsUpgraded"`
    53  }
    54  
    55  // Returns an InboundConnUpgradeThrottler that upgrades an inbound
    56  // connection from a given IP at most every [UpgradeCooldown].
    57  func NewInboundConnUpgradeThrottler(log logging.Logger, config InboundConnUpgradeThrottlerConfig) InboundConnUpgradeThrottler {
    58  	if config.UpgradeCooldown <= 0 || config.MaxRecentConnsUpgraded <= 0 {
    59  		return &noInboundConnUpgradeThrottler{}
    60  	}
    61  	return &inboundConnUpgradeThrottler{
    62  		InboundConnUpgradeThrottlerConfig: config,
    63  		log:                               log,
    64  		done:                              make(chan struct{}),
    65  		recentIPsAndTimes:                 make(chan ipAndTime, config.MaxRecentConnsUpgraded),
    66  	}
    67  }
    68  
    69  // noInboundConnUpgradeThrottler upgrades all inbound connections
    70  type noInboundConnUpgradeThrottler struct{}
    71  
    72  func (*noInboundConnUpgradeThrottler) Dispatch() {}
    73  
    74  func (*noInboundConnUpgradeThrottler) Stop() {}
    75  
    76  func (*noInboundConnUpgradeThrottler) ShouldUpgrade(netip.AddrPort) bool {
    77  	return true
    78  }
    79  
    80  type ipAndTime struct {
    81  	ip                netip.Addr
    82  	cooldownElapsedAt time.Time
    83  }
    84  
    85  type inboundConnUpgradeThrottler struct {
    86  	InboundConnUpgradeThrottlerConfig
    87  	log  logging.Logger
    88  	lock sync.Mutex
    89  	// Useful for faking time in tests
    90  	clock mockable.Clock
    91  	// When [done] is closed, Dispatch returns.
    92  	done chan struct{}
    93  	// IP --> Present if ShouldUpgrade(ipStr) returned true
    94  	// within the last [UpgradeCooldown].
    95  	recentIPs set.Set[netip.Addr]
    96  	// Sorted in order of increasing time
    97  	// of last call to ShouldUpgrade that returned true.
    98  	// For each IP in this channel, ShouldUpgrade(ipStr)
    99  	// returned true within the last [UpgradeCooldown].
   100  	recentIPsAndTimes chan ipAndTime
   101  }
   102  
   103  // Returns whether we should upgrade an inbound connection from [ipStr].
   104  func (n *inboundConnUpgradeThrottler) ShouldUpgrade(addrPort netip.AddrPort) bool {
   105  	// Only use addr (not port). This mitigates DoS attacks from many nodes on one
   106  	// host.
   107  	addr := addrPort.Addr()
   108  	if addr.IsLoopback() {
   109  		// Don't rate-limit loopback IPs
   110  		return true
   111  	}
   112  
   113  	n.lock.Lock()
   114  	defer n.lock.Unlock()
   115  
   116  	if n.recentIPs.Contains(addr) {
   117  		// We recently upgraded an inbound connection from this IP
   118  		return false
   119  	}
   120  
   121  	select {
   122  	case n.recentIPsAndTimes <- ipAndTime{
   123  		ip:                addr,
   124  		cooldownElapsedAt: n.clock.Time().Add(n.UpgradeCooldown),
   125  	}:
   126  		n.recentIPs.Add(addr)
   127  		return true
   128  	default:
   129  		return false
   130  	}
   131  }
   132  
   133  func (n *inboundConnUpgradeThrottler) Dispatch() {
   134  	timer := time.NewTimer(0)
   135  	if !timer.Stop() {
   136  		<-timer.C
   137  	}
   138  
   139  	defer timer.Stop()
   140  	for {
   141  		select {
   142  		case next := <-n.recentIPsAndTimes:
   143  			// Sleep until it's time to remove the next IP
   144  			timer.Reset(next.cooldownElapsedAt.Sub(n.clock.Time()))
   145  
   146  			select {
   147  			case <-timer.C:
   148  				// Remove the next IP (we'd upgrade another inbound connection from it)
   149  				n.lock.Lock()
   150  				n.recentIPs.Remove(next.ip)
   151  				n.lock.Unlock()
   152  			case <-n.done:
   153  				return
   154  			}
   155  		case <-n.done:
   156  			return
   157  		}
   158  	}
   159  }
   160  
   161  func (n *inboundConnUpgradeThrottler) Stop() {
   162  	close(n.done)
   163  }