mirror of https://github.com/status-im/go-waku.git
87 lines
2.4 KiB
Go
87 lines
2.4 KiB
Go
|
package peermanager
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"github.com/libp2p/go-libp2p/core/control"
|
||
|
|
"github.com/libp2p/go-libp2p/core/network"
|
||
|
|
"github.com/libp2p/go-libp2p/core/peerstore"
|
||
|
|
ma "github.com/multiformats/go-multiaddr"
|
||
|
|
manet "github.com/multiformats/go-multiaddr/net"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
"github.com/waku-org/go-waku/waku/v2/utils"
|
||
|
|
"testing"
|
||
|
|
)
|
||
|
|
|
||
|
|
type mockConnMultiaddrs struct {
|
||
|
|
local, remote ma.Multiaddr
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m mockConnMultiaddrs) LocalMultiaddr() ma.Multiaddr {
|
||
|
|
return m.local
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m mockConnMultiaddrs) RemoteMultiaddr() ma.Multiaddr {
|
||
|
|
return m.remote
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestConnectionGater(t *testing.T) {
|
||
|
|
|
||
|
|
log := utils.Logger()
|
||
|
|
|
||
|
|
_, h1, _ := makeWakuRelay(t, log)
|
||
|
|
_, h2, _ := makeWakuRelay(t, log)
|
||
|
|
|
||
|
|
h1.Peerstore().AddAddrs(h2.ID(), h2.Addrs(), peerstore.PermanentAddrTTL)
|
||
|
|
err := h1.Connect(context.Background(), h2.Peerstore().PeerInfo(h2.ID()))
|
||
|
|
require.NoError(t, err)
|
||
|
|
|
||
|
|
h2.Peerstore().AddAddrs(h1.ID(), h1.Addrs(), peerstore.PermanentAddrTTL)
|
||
|
|
err = h2.Connect(context.Background(), h1.Peerstore().PeerInfo(h1.ID()))
|
||
|
|
require.NoError(t, err)
|
||
|
|
|
||
|
|
peerA := h1.ID()
|
||
|
|
|
||
|
|
remoteAddr1 := ma.StringCast("/ip4/1.2.3.4/tcp/1234")
|
||
|
|
|
||
|
|
connGater := NewConnectionGater(2, log)
|
||
|
|
|
||
|
|
// Test peer blocking
|
||
|
|
allow := connGater.InterceptPeerDial(peerA)
|
||
|
|
require.True(t, allow)
|
||
|
|
|
||
|
|
// Test connection was secured and upgraded
|
||
|
|
allow = connGater.InterceptSecured(network.DirInbound, peerA, &mockConnMultiaddrs{local: nil, remote: nil})
|
||
|
|
require.True(t, allow)
|
||
|
|
|
||
|
|
connection1 := h1.Network().Conns()[0]
|
||
|
|
|
||
|
|
allow, reason := connGater.InterceptUpgraded(connection1)
|
||
|
|
require.True(t, allow)
|
||
|
|
require.Equal(t, control.DisconnectReason(0), reason)
|
||
|
|
|
||
|
|
// Test addr and subnet blocking
|
||
|
|
allow = connGater.InterceptAddrDial(peerA, remoteAddr1)
|
||
|
|
require.True(t, allow)
|
||
|
|
|
||
|
|
// Bellow the connection limit
|
||
|
|
allow = connGater.InterceptAccept(&mockConnMultiaddrs{local: nil, remote: remoteAddr1})
|
||
|
|
require.True(t, allow)
|
||
|
|
|
||
|
|
ip, err := manet.ToIP(remoteAddr1)
|
||
|
|
require.NoError(t, err)
|
||
|
|
connGater.limiter[ip.String()] = 3
|
||
|
|
|
||
|
|
// Above the connection limit
|
||
|
|
allow = connGater.InterceptAccept(&mockConnMultiaddrs{local: nil, remote: remoteAddr1})
|
||
|
|
require.False(t, allow)
|
||
|
|
|
||
|
|
// Call twice NotifyDisconnect to get bellow the limit(2): 3 -> 1
|
||
|
|
connGater.NotifyDisconnect(remoteAddr1)
|
||
|
|
connGater.NotifyDisconnect(remoteAddr1)
|
||
|
|
|
||
|
|
// Bellow the connection limit again
|
||
|
|
allow = connGater.validateInboundConn(remoteAddr1)
|
||
|
|
require.True(t, allow)
|
||
|
|
|
||
|
|
}
|