package rcmgr import ( "math" "net/netip" "slices" "sync" ) type ConnLimitPerSubnet struct { // This defines how big the subnet is. For example, a /24 subnet has a // PrefixLength of 24. All IPs that share the same 24 bit prefix are in the same // subnet. Are in the same subnet, and bound to the same limit. PrefixLength int // The maximum number of connections allowed for each subnet. ConnCount int } type NetworkPrefixLimit struct { // The Network prefix for which this limit applies. Network netip.Prefix // The maximum number of connections allowed for this subnet. ConnCount int } // 8 for now so that it matches the number of concurrent dials we may do // in swarm_dial.go. With future smart dialing work we should bring this // down var defaultMaxConcurrentConns = 8 var defaultIP4Limit = ConnLimitPerSubnet{ ConnCount: defaultMaxConcurrentConns, PrefixLength: 32, } var defaultIP6Limits = []ConnLimitPerSubnet{ { ConnCount: defaultMaxConcurrentConns, PrefixLength: 56, }, { ConnCount: 8 * defaultMaxConcurrentConns, PrefixLength: 48, }, } var DefaultNetworkPrefixLimitV4 = sortNetworkPrefixes([]NetworkPrefixLimit{ { // Loopback address for v4 https://datatracker.ietf.org/doc/html/rfc6890#section-2.2.2 Network: netip.MustParsePrefix("127.0.0.0/8"), ConnCount: math.MaxInt, // Unlimited }, }) var DefaultNetworkPrefixLimitV6 = sortNetworkPrefixes([]NetworkPrefixLimit{ { // Loopback address for v6 https://datatracker.ietf.org/doc/html/rfc6890#section-2.2.3 Network: netip.MustParsePrefix("::1/128"), ConnCount: math.MaxInt, // Unlimited }, }) // Network prefixes limits must be sorted by most specific to least specific. This lets us // actually use the more specific limits, otherwise only the less specific ones // would be matched. e.g. 1.2.3.0/24 must come before 1.2.0.0/16. func sortNetworkPrefixes(limits []NetworkPrefixLimit) []NetworkPrefixLimit { slices.SortStableFunc(limits, func(a, b NetworkPrefixLimit) int { return b.Network.Bits() - a.Network.Bits() }) return limits } // WithNetworkPrefixLimit sets the limits for the number of connections allowed // for a specific Network Prefix. Use this when you want to set higher limits // for a specific subnet than the default limit per subnet. func WithNetworkPrefixLimit(ipv4 []NetworkPrefixLimit, ipv6 []NetworkPrefixLimit) Option { return func(rm *resourceManager) error { if ipv4 != nil { rm.connLimiter.networkPrefixLimitV4 = sortNetworkPrefixes(ipv4) } if ipv6 != nil { rm.connLimiter.networkPrefixLimitV6 = sortNetworkPrefixes(ipv6) } return nil } } // WithLimitPerSubnet sets the limits for the number of connections allowed per // subnet. This will limit the number of connections per subnet if that subnet // is not defined in the NetworkPrefixLimit option. Think of this as a default // limit for any given subnet. func WithLimitPerSubnet(ipv4 []ConnLimitPerSubnet, ipv6 []ConnLimitPerSubnet) Option { return func(rm *resourceManager) error { if ipv4 != nil { rm.connLimiter.connLimitPerSubnetV4 = ipv4 } if ipv6 != nil { rm.connLimiter.connLimitPerSubnetV6 = ipv6 } return nil } } type connLimiter struct { mu sync.Mutex // Specific Network Prefix limits. If these are set, they take precedence over the // subnet limits. // These must be sorted by most specific to least specific. networkPrefixLimitV4 []NetworkPrefixLimit networkPrefixLimitV6 []NetworkPrefixLimit connsPerNetworkPrefixV4 []int connsPerNetworkPrefixV6 []int // Subnet limits. connLimitPerSubnetV4 []ConnLimitPerSubnet connLimitPerSubnetV6 []ConnLimitPerSubnet ip4connsPerLimit []map[string]int ip6connsPerLimit []map[string]int } func newConnLimiter() *connLimiter { return &connLimiter{ networkPrefixLimitV4: DefaultNetworkPrefixLimitV4, networkPrefixLimitV6: DefaultNetworkPrefixLimitV6, connLimitPerSubnetV4: []ConnLimitPerSubnet{defaultIP4Limit}, connLimitPerSubnetV6: defaultIP6Limits, } } func (cl *connLimiter) addNetworkPrefixLimit(isIP6 bool, npLimit NetworkPrefixLimit) { cl.mu.Lock() defer cl.mu.Unlock() if isIP6 { cl.networkPrefixLimitV6 = append(cl.networkPrefixLimitV6, npLimit) cl.networkPrefixLimitV6 = sortNetworkPrefixes(cl.networkPrefixLimitV6) } else { cl.networkPrefixLimitV4 = append(cl.networkPrefixLimitV4, npLimit) cl.networkPrefixLimitV4 = sortNetworkPrefixes(cl.networkPrefixLimitV4) } } // addConn adds a connection for the given IP address. It returns true if the connection is allowed. func (cl *connLimiter) addConn(ip netip.Addr) bool { cl.mu.Lock() defer cl.mu.Unlock() networkPrefixLimits := cl.networkPrefixLimitV4 connsPerNetworkPrefix := cl.connsPerNetworkPrefixV4 limits := cl.connLimitPerSubnetV4 connsPerLimit := cl.ip4connsPerLimit isIP6 := ip.Is6() if isIP6 { networkPrefixLimits = cl.networkPrefixLimitV6 connsPerNetworkPrefix = cl.connsPerNetworkPrefixV6 limits = cl.connLimitPerSubnetV6 connsPerLimit = cl.ip6connsPerLimit } // Check Network Prefix limits first if len(connsPerNetworkPrefix) == 0 && len(networkPrefixLimits) > 0 { // Initialize the counts connsPerNetworkPrefix = make([]int, len(networkPrefixLimits)) if isIP6 { cl.connsPerNetworkPrefixV6 = connsPerNetworkPrefix } else { cl.connsPerNetworkPrefixV4 = connsPerNetworkPrefix } } for i, limit := range networkPrefixLimits { if limit.Network.Contains(ip) { if connsPerNetworkPrefix[i]+1 > limit.ConnCount { return false } connsPerNetworkPrefix[i]++ // Done. If we find a match in the network prefix limits, we use // that and don't use the general subnet limits. return true } } if len(connsPerLimit) == 0 && len(limits) > 0 { connsPerLimit = make([]map[string]int, len(limits)) if isIP6 { cl.ip6connsPerLimit = connsPerLimit } else { cl.ip4connsPerLimit = connsPerLimit } } for i, limit := range limits { prefix, err := ip.Prefix(limit.PrefixLength) if err != nil { return false } masked := prefix.String() counts, ok := connsPerLimit[i][masked] if !ok { if connsPerLimit[i] == nil { connsPerLimit[i] = make(map[string]int) } connsPerLimit[i][masked] = 0 } if counts+1 > limit.ConnCount { return false } } // All limit checks passed, now we update the counts for i, limit := range limits { prefix, _ := ip.Prefix(limit.PrefixLength) masked := prefix.String() connsPerLimit[i][masked]++ } return true } func (cl *connLimiter) rmConn(ip netip.Addr) { cl.mu.Lock() defer cl.mu.Unlock() networkPrefixLimits := cl.networkPrefixLimitV4 connsPerNetworkPrefix := cl.connsPerNetworkPrefixV4 limits := cl.connLimitPerSubnetV4 connsPerLimit := cl.ip4connsPerLimit isIP6 := ip.Is6() if isIP6 { networkPrefixLimits = cl.networkPrefixLimitV6 connsPerNetworkPrefix = cl.connsPerNetworkPrefixV6 limits = cl.connLimitPerSubnetV6 connsPerLimit = cl.ip6connsPerLimit } // Check NetworkPrefix limits first if len(connsPerNetworkPrefix) == 0 && len(networkPrefixLimits) > 0 { // Initialize just in case. We should have already initialized in // addConn, but if the callers calls rmConn first we don't want to panic connsPerNetworkPrefix = make([]int, len(networkPrefixLimits)) if isIP6 { cl.connsPerNetworkPrefixV6 = connsPerNetworkPrefix } else { cl.connsPerNetworkPrefixV4 = connsPerNetworkPrefix } } for i, limit := range networkPrefixLimits { if limit.Network.Contains(ip) { count := connsPerNetworkPrefix[i] if count <= 0 { log.Errorf("unexpected conn count for ip %s. Was this not added with addConn first?", ip) return } connsPerNetworkPrefix[i]-- // Done. We updated the count in the defined network prefix limit. return } } if len(connsPerLimit) == 0 && len(limits) > 0 { // Initialize just in case. We should have already initialized in // addConn, but if the callers calls rmConn first we don't want to panic connsPerLimit = make([]map[string]int, len(limits)) if isIP6 { cl.ip6connsPerLimit = connsPerLimit } else { cl.ip4connsPerLimit = connsPerLimit } } for i, limit := range limits { prefix, err := ip.Prefix(limit.PrefixLength) if err != nil { // Unexpected since we should have seen this IP before in addConn log.Errorf("unexpected error getting prefix: %v", err) continue } masked := prefix.String() counts, ok := connsPerLimit[i][masked] if !ok || counts == 0 { // Unexpected, but don't panic log.Errorf("unexpected conn count for %s ok=%v count=%v", masked, ok, counts) continue } connsPerLimit[i][masked]-- if connsPerLimit[i][masked] <= 0 { delete(connsPerLimit[i], masked) } } }