mirror of
https://github.com/libp2p/go-libp2p-resource-manager.git
synced 2025-03-10 09:10:38 +08:00
Add WithAllowlistedMultiaddr option
This commit is contained in:
parent
8c5643aa2f
commit
374b297cee
13
allowlist.go
13
allowlist.go
@ -26,6 +26,19 @@ type Allowlist struct {
|
||||
allowedPeerByNetwork map[peer.ID][]*net.IPNet
|
||||
}
|
||||
|
||||
// WithAllowlistedMultiaddrs sets the multiaddrs to be in the allowlist
|
||||
func WithAllowlistedMultiaddrs(mas []multiaddr.Multiaddr) Option {
|
||||
return func(rm *resourceManager) error {
|
||||
for _, ma := range mas {
|
||||
err := rm.allowlist.Add(ma)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func newAllowlist() Allowlist {
|
||||
return Allowlist{
|
||||
allowedPeerByNetwork: make(map[peer.ID][]*net.IPNet),
|
||||
|
@ -1004,16 +1004,19 @@ func TestResourceManager(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestResourceManagerWithAllowlist(t *testing.T) {
|
||||
peerA := test.RandPeerIDFatal(t)
|
||||
|
||||
limits := NewDefaultLimiter()
|
||||
limits.SystemLimits = limits.SystemLimits.WithConnLimit(0, 0, 0)
|
||||
limits.TransientLimits = limits.SystemLimits.WithConnLimit(0, 0, 0)
|
||||
rcmgr, err := NewResourceManager(limits)
|
||||
rcmgr, err := NewResourceManager(limits, WithAllowlistedMultiaddrs([]multiaddr.Multiaddr{
|
||||
multiaddr.StringCast("/ip4/1.2.3.4"),
|
||||
multiaddr.StringCast("/ip4/4.3.2.1/p2p/" + peerA.String()),
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
peerA := test.RandPeerIDFatal(t)
|
||||
|
||||
{
|
||||
// Setup allowlist. TODO, replace this with a config once config changes are in
|
||||
r := rcmgr.(*resourceManager)
|
||||
@ -1022,10 +1025,6 @@ func TestResourceManagerWithAllowlist(t *testing.T) {
|
||||
r.allowlistedSystem.IncRef()
|
||||
r.allowlistedTransient = newTransientScope(limits.GetTransientLimits().WithConnLimit(1, 1, 1), r, "allowlistedTransient", r.allowlistedSystem.resourceScope)
|
||||
r.allowlistedTransient.IncRef()
|
||||
allowlist := r.GetAllowlist()
|
||||
|
||||
allowlist.Add(multiaddr.StringCast("/ip4/1.2.3.4"))
|
||||
allowlist.Add(multiaddr.StringCast("/ip4/4.3.2.1/p2p/" + peerA.String()))
|
||||
}
|
||||
|
||||
// A connection comes in from a non-allowlisted ip address
|
||||
@ -1066,5 +1065,4 @@ func TestResourceManagerWithAllowlist(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user