Add WithAllowlistedMultiaddr option

This commit is contained in:
Marco Munizaga 2022-06-27 15:47:51 -07:00
parent 8c5643aa2f
commit 374b297cee
2 changed files with 19 additions and 8 deletions

View File

@ -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),

View File

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