waku-org/go-waku

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)

}