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 }