package rcmgr
import (
"crypto/rand"
"fmt"
"net"
"testing"
"github.com/libp2p/go-libp2p/core/peer"
"github.com/libp2p/go-libp2p/core/test"
"github.com/multiformats/go-multiaddr"
)
func ExampleWithAllowlistedMultiaddrs() {
somePeer, err := test.RandPeerID()
if err != nil {
panic("Failed to generate somePeer")
}
limits := DefaultLimits.AutoScale()
rcmgr, err := NewResourceManager(NewFixedLimiter(limits), WithAllowlistedMultiaddrs([]multiaddr.Multiaddr{
// Any peer connecting from this IP address
multiaddr.StringCast("/ip4/1.2.3.4"),
// Only the specified peer from this address
multiaddr.StringCast("/ip4/2.2.3.4/p2p/" + somePeer.String()),
// Only peers from this 1.2.3.0/24 IP address range
multiaddr.StringCast("/ip4/1.2.3.0/ipcidr/24"),
}))
if err != nil {
panic("Failed to start resource manager")
}
// Use rcmgr as before
_ = rcmgr
}
func TestAllowedSimple(t *testing.T) {
allowlist := newAllowlist()
ma := multiaddr.StringCast("/ip4/1.2.3.4/tcp/1234")
err := allowlist.Add(ma)
if err != nil {
t.Fatalf("failed to add ip4: %s", err)
}
if !allowlist.Allowed(ma) {
t.Fatalf("addr should be allowed")
}
}
func TestAllowedWithPeer(t *testing.T) {
type testcase struct {
name string
allowlist []string
endpoint multiaddr.Multiaddr
peer peer.ID
// Is this endpoint allowed? (We don't have peer info yet)
isConnAllowed bool
// Is this peer + endpoint allowed?
isAllowedWithPeer bool
}
peerA := test.RandPeerIDFatal(t)
peerB := test.RandPeerIDFatal(t)
multiaddrA := multiaddr.StringCast("/ip4/1.2.3.4/tcp/1234")
multiaddrB := multiaddr.StringCast("/ip4/2.2.3.4/tcp/1234")
testcases := []testcase{
{
name: "Blocked",
isConnAllowed: false,
isAllowedWithPeer: false,
allowlist: []string{"/ip4/1.2.3.1"},
endpoint: multiaddrA,
peer: peerA,
},
{
name: "Blocked wrong peer",
isConnAllowed: true,
isAllowedWithPeer: false,
allowlist: []string{"/ip4/1.2.3.4" + "/p2p/" + peerB.String()},
endpoint: multiaddrA,
peer: peerA,
},
{
name: "allowed on network",
isConnAllowed: true,
isAllowedWithPeer: true,
allowlist: []string{"/ip4/1.2.3.0/ipcidr/24"},
endpoint: multiaddrA,
peer: peerA,
},
{
name: "Blocked peer not on network",
isConnAllowed: true,
isAllowedWithPeer: true,
allowlist: []string{"/ip4/1.2.3.0/ipcidr/24"},
endpoint: multiaddrA,
peer: peerA,
}, {
name: "allowed. right network, right peer",
isConnAllowed: true,
isAllowedWithPeer: true,
allowlist: []string{"/ip4/1.2.3.0/ipcidr/24" + "/p2p/" + peerA.String()},
endpoint: multiaddrA,
peer: peerA,
}, {
name: "allowed. right network, no peer",
isConnAllowed: true,
isAllowedWithPeer: true,
allowlist: []string{"/ip4/1.2.3.0/ipcidr/24"},
endpoint: multiaddrA,
peer: peerA,
},
{
name: "Blocked. right network, wrong peer",
isConnAllowed: true,
isAllowedWithPeer: false,
allowlist: []string{"/ip4/1.2.3.0/ipcidr/24" + "/p2p/" + peerB.String()},
endpoint: multiaddrA,
peer: peerA,
},
{
name: "allowed peer any ip",
isConnAllowed: true,
isAllowedWithPeer: true,
allowlist: []string{"/ip4/0.0.0.0/ipcidr/0"},
endpoint: multiaddrA,
peer: peerA,
},
{
name: "allowed peer multiple ips in allowlist",
isConnAllowed: true,
isAllowedWithPeer: true,
allowlist: []string{"/ip4/1.2.3.4/p2p/" + peerA.String(), "/ip4/2.2.3.4/p2p/" + peerA.String()},
endpoint: multiaddrA,
peer: peerA,
},
{
name: "allowed peer multiple ips in allowlist",
isConnAllowed: true,
isAllowedWithPeer: true,
allowlist: []string{"/ip4/1.2.3.4/p2p/" + peerA.String(), "/ip4/1.2.3.4/p2p/" + peerA.String()},
endpoint: multiaddrA,
peer: peerA,
},
{
name: "allowed peer multiple ips in allowlist",
isConnAllowed: true,
isAllowedWithPeer: true,
allowlist: []string{"/ip4/1.2.3.4/p2p/" + peerA.String(), "/ip4/2.2.3.4/p2p/" + peerA.String()},
endpoint: multiaddrB,
peer: peerA,
},
}
for _, tc := range testcases {
t.Run(tc.name, func(t *testing.T) {
allowlist := newAllowlist()
for _, maStr := range tc.allowlist {
ma, err := multiaddr.NewMultiaddr(maStr)
if err != nil {
fmt.Printf("failed to parse multiaddr: %s", err)
}
allowlist.Add(ma)
}
if allowlist.Allowed(tc.endpoint) != tc.isConnAllowed {
t.Fatalf("%v: expected %v", !tc.isConnAllowed, tc.isConnAllowed)
}
if allowlist.AllowedPeerAndMultiaddr(tc.peer, tc.endpoint) != tc.isAllowedWithPeer {
t.Fatalf("%v: expected %v", !tc.isAllowedWithPeer, tc.isAllowedWithPeer)
}
})
}
}
func TestRemoved(t *testing.T) {
type testCase struct {
name string
allowedMA string
}
peerA := test.RandPeerIDFatal(t)
maA := multiaddr.StringCast("/ip4/1.2.3.4")
testCases := []testCase{
{name: "ip4", allowedMA: "/ip4/1.2.3.4"},
{name: "ip4 with peer", allowedMA: "/ip4/1.2.3.4/p2p/" + peerA.String()},
{name: "ip4 network", allowedMA: "/ip4/0.0.0.0/ipcidr/0"},
{name: "ip4 network with peer", allowedMA: "/ip4/0.0.0.0/ipcidr/0/p2p/" + peerA.String()},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
allowlist := newAllowlist()
ma := multiaddr.StringCast(tc.allowedMA)
err := allowlist.Add(ma)
if err != nil {
t.Fatalf("failed to add ip4: %s", err)
}
if !allowlist.AllowedPeerAndMultiaddr(peerA, maA) {
t.Fatalf("addr should be allowed")
}
allowlist.Remove((ma))
if allowlist.AllowedPeerAndMultiaddr(peerA, maA) {
t.Fatalf("addr should not be allowed")
}
})
}
}
// BenchmarkAllowlistCheck benchmarks the allowlist with plausible conditions.
func BenchmarkAllowlistCheck(b *testing.B) {
allowlist := newAllowlist()
// How often do we expect a peer to be specified? 1 in N
ratioOfSpecifiedPeers := 10
// How often do we expect an allowlist hit? 1 in N
ratioOfAllowlistHit := 100
// How many multiaddrs in our allowlist?
howManyMultiaddrsInAllowList := 1_000
// How often is the IP addr an IPV6? 1 in N
ratioOfIPV6 := 20
countOfTotalPeersForTest := 100_000
mas := make([]multiaddr.Multiaddr, countOfTotalPeersForTest)
for i := 0; i < countOfTotalPeersForTest; i++ {
ip := make([]byte, 16)
n, err := rand.Reader.Read(ip)
if err != nil || n != 16 {
b.Fatalf("Failed to generate IP address")
}
var ipString string
if i%ratioOfIPV6 == 0 {
// IPv6
ip6 := net.IP(ip)
ipString = "/ip6/" + ip6.String()
} else {
// IPv4
ip4 := net.IPv4(ip[0], ip[1], ip[2], ip[3])
ipString = "/ip4/" + ip4.String()
}
var ma multiaddr.Multiaddr
if i%ratioOfSpecifiedPeers == 0 {
ma = multiaddr.StringCast(ipString + "/p2p/" + test.RandPeerIDFatal(b).String())
} else {
ma = multiaddr.StringCast(ipString)
}
if err != nil {
b.Fatalf("Failed to generate multiaddr: %v", ipString)
}
mas[i] = ma
}
for _, ma := range mas[:howManyMultiaddrsInAllowList] {
err := allowlist.Add(ma)
if err != nil {
b.Fatalf("Failed to add multiaddr")
}
}
masInAllowList := mas[:howManyMultiaddrsInAllowList]
masNotInAllowList := mas[howManyMultiaddrsInAllowList:]
b.ResetTimer()
for n := 0; n < b.N; n++ {
if n%ratioOfAllowlistHit == 0 {
allowlist.Allowed(masInAllowList[n%len(masInAllowList)])
} else {
allowlist.Allowed(masNotInAllowList[n%len(masNotInAllowList)])
}
}
}