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)]) } } }