ceremonyclient/go-libp2p/p2p/host/resource-manager/conn_limiter_test.go
Cassie Heart ebb31fd792 v1.4.19
2024-06-08 11:32:45 +00:00

215 lines
5.2 KiB
Go

package rcmgr
import (
"encoding/binary"
"net"
"net/netip"
"testing"
"github.com/stretchr/testify/require"
)
func TestItLimits(t *testing.T) {
t.Run("IPv4", func(t *testing.T) {
ip, err := netip.ParseAddr("1.2.3.4")
require.NoError(t, err)
cl := newConnLimiter()
cl.connLimitPerSubnetV4[0].ConnCount = 1
require.True(t, cl.addConn(ip))
// should fail the second time
require.False(t, cl.addConn(ip))
otherIP, err := netip.ParseAddr("1.2.3.5")
require.NoError(t, err)
require.True(t, cl.addConn(otherIP))
})
t.Run("IPv6", func(t *testing.T) {
ip, err := netip.ParseAddr("1:2:3:4::1")
require.NoError(t, err)
cl := newConnLimiter()
original := cl.connLimitPerSubnetV6[0].ConnCount
cl.connLimitPerSubnetV6[0].ConnCount = 1
defer func() {
cl.connLimitPerSubnetV6[0].ConnCount = original
}()
require.True(t, cl.addConn(ip))
// should fail the second time
require.False(t, cl.addConn(ip))
otherIPSameSubnet := netip.MustParseAddr("1:2:3:4::2")
require.False(t, cl.addConn(otherIPSameSubnet))
otherIP := netip.MustParseAddr("2:2:3:4::2")
require.True(t, cl.addConn(otherIP))
})
t.Run("IPv6 with multiple limits", func(t *testing.T) {
cl := newConnLimiter()
for i := 0; i < defaultMaxConcurrentConns; i++ {
ip := net.ParseIP("ff:2:3:4::1")
binary.BigEndian.PutUint16(ip[14:], uint16(i))
ipAddr := netip.MustParseAddr(ip.String())
require.True(t, cl.addConn(ipAddr))
}
// Next one should fail
ip := net.ParseIP("ff:2:3:4::1")
binary.BigEndian.PutUint16(ip[14:], uint16(defaultMaxConcurrentConns+1))
require.False(t, cl.addConn(netip.MustParseAddr(ip.String())))
// But on a different root subnet should work
otherIP := netip.MustParseAddr("ffef:2:3::1")
require.True(t, cl.addConn(otherIP))
// But too many on the next subnet limit will fail too
for i := 0; i < defaultMaxConcurrentConns*8; i++ {
ip := net.ParseIP("ffef:2:3:4::1")
binary.BigEndian.PutUint16(ip[5:7], uint16(i))
ipAddr := netip.MustParseAddr(ip.String())
require.True(t, cl.addConn(ipAddr))
}
ip = net.ParseIP("ffef:2:3:4::1")
binary.BigEndian.PutUint16(ip[5:7], uint16(defaultMaxConcurrentConns*8+1))
ipAddr := netip.MustParseAddr(ip.String())
require.False(t, cl.addConn(ipAddr))
})
t.Run("IPv4 with localhost", func(t *testing.T) {
cl := &connLimiter{
networkPrefixLimitV4: DefaultNetworkPrefixLimitV4,
connLimitPerSubnetV4: []ConnLimitPerSubnet{
{PrefixLength: 0, ConnCount: 1}, // 1 connection for the whole IPv4 space
},
}
ip := netip.MustParseAddr("1.2.3.4")
require.True(t, cl.addConn(ip))
ip = netip.MustParseAddr("4.3.2.1")
// should fail the second time, we only allow 1 connection for the whole IPv4 space
require.False(t, cl.addConn(ip))
ip = netip.MustParseAddr("127.0.0.1")
// Succeeds because we defined an explicit limit for the loopback subnet
require.True(t, cl.addConn(ip))
})
}
func genIP(data *[]byte) (netip.Addr, bool) {
if len(*data) < 1 {
return netip.Addr{}, false
}
genIP6 := (*data)[0]&0x01 == 1
bytesRequired := 4
if genIP6 {
bytesRequired = 16
}
if len((*data)[1:]) < bytesRequired {
return netip.Addr{}, false
}
*data = (*data)[1:]
ip, ok := netip.AddrFromSlice((*data)[:bytesRequired])
*data = (*data)[bytesRequired:]
return ip, ok
}
func FuzzConnLimiter(f *testing.F) {
// The goal is to try to enter a state where the count is incorrectly 0
f.Fuzz(func(t *testing.T, data []byte) {
ips := make([]netip.Addr, 0, len(data)/5)
for {
ip, ok := genIP(&data)
if !ok {
break
}
ips = append(ips, ip)
}
cl := newConnLimiter()
addedConns := make([]netip.Addr, 0, len(ips))
for _, ip := range ips {
if cl.addConn(ip) {
addedConns = append(addedConns, ip)
}
}
addedCount := 0
for _, ip := range cl.ip4connsPerLimit {
for _, count := range ip {
addedCount += count
}
}
for _, ip := range cl.ip6connsPerLimit {
for _, count := range ip {
addedCount += count
}
}
for _, count := range cl.connsPerNetworkPrefixV4 {
addedCount += count
}
for _, count := range cl.connsPerNetworkPrefixV6 {
addedCount += count
}
if addedCount == 0 && len(addedConns) > 0 {
t.Fatalf("added count: %d", addedCount)
}
for _, ip := range addedConns {
cl.rmConn(ip)
}
leftoverCount := 0
for _, ip := range cl.ip4connsPerLimit {
for _, count := range ip {
leftoverCount += count
}
}
for _, ip := range cl.ip6connsPerLimit {
for _, count := range ip {
leftoverCount += count
}
}
for _, count := range cl.connsPerNetworkPrefixV4 {
addedCount += count
}
for _, count := range cl.connsPerNetworkPrefixV6 {
addedCount += count
}
if leftoverCount != 0 {
t.Fatalf("leftover count: %d", leftoverCount)
}
})
}
func TestSortedNetworkPrefixLimits(t *testing.T) {
npLimits := []NetworkPrefixLimit{
{
Network: netip.MustParsePrefix("1.2.0.0/16"),
},
{
Network: netip.MustParsePrefix("1.2.3.0/28"),
},
{
Network: netip.MustParsePrefix("1.2.3.4/32"),
},
}
npLimits = sortNetworkPrefixes(npLimits)
sorted := []NetworkPrefixLimit{
{
Network: netip.MustParsePrefix("1.2.3.4/32"),
},
{
Network: netip.MustParsePrefix("1.2.3.0/28"),
},
{
Network: netip.MustParsePrefix("1.2.0.0/16"),
},
}
require.EqualValues(t, sorted, npLimits)
}