mirror of
https://source.quilibrium.com/quilibrium/ceremonyclient.git
synced 2024-11-15 04:35:17 +00:00
121 lines
3.6 KiB
Go
121 lines
3.6 KiB
Go
package basichost
|
|
|
|
import (
|
|
"context"
|
|
"net/netip"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
|
|
ma "github.com/multiformats/go-multiaddr"
|
|
|
|
swarmt "github.com/libp2p/go-libp2p/p2p/net/swarm/testing"
|
|
|
|
"go.uber.org/mock/gomock"
|
|
)
|
|
|
|
func setupMockNAT(t *testing.T) (mockNAT *MockNAT, reset func()) {
|
|
t.Helper()
|
|
ctrl := gomock.NewController(t)
|
|
mockNAT = NewMockNAT(ctrl)
|
|
origDiscoverNAT := discoverNAT
|
|
discoverNAT = func(ctx context.Context) (nat, error) { return mockNAT, nil }
|
|
return mockNAT, func() {
|
|
discoverNAT = origDiscoverNAT
|
|
ctrl.Finish()
|
|
}
|
|
}
|
|
|
|
func TestMapping(t *testing.T) {
|
|
mockNAT, reset := setupMockNAT(t)
|
|
defer reset()
|
|
|
|
sw := swarmt.GenSwarm(t)
|
|
defer sw.Close()
|
|
m := newNATManager(sw)
|
|
require.Eventually(t, func() bool {
|
|
m.natMx.Lock()
|
|
defer m.natMx.Unlock()
|
|
return m.nat != nil
|
|
}, time.Second, time.Millisecond)
|
|
externalAddr := netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 2, 3, 4}), 4321)
|
|
// pretend that we have a TCP mapping
|
|
mockNAT.EXPECT().GetMapping("tcp", 1234).Return(externalAddr, true)
|
|
m1, _ := ma.StringCast("/ip4/1.2.3.4/tcp/4321")
|
|
m2, _ := ma.StringCast("/ip4/0.0.0.0/tcp/1234")
|
|
require.Equal(t, m1, m.GetMapping(m2))
|
|
|
|
// pretend that we have a QUIC mapping
|
|
mockNAT.EXPECT().GetMapping("udp", 1234).Return(externalAddr, true)
|
|
m3, _ := ma.StringCast("/ip4/1.2.3.4/udp/4321/quic-v1")
|
|
m4, _ := ma.StringCast("/ip4/0.0.0.0/udp/1234/quic-v1")
|
|
require.Equal(t, m3, m.GetMapping(m4))
|
|
|
|
// pretend that there's no mapping
|
|
mockNAT.EXPECT().GetMapping("tcp", 1234).Return(netip.AddrPort{}, false)
|
|
m5, _ := ma.StringCast("/ip4/0.0.0.0/tcp/1234")
|
|
require.Nil(t, m.GetMapping(m5))
|
|
|
|
// make sure this works for WebSocket addresses as well
|
|
mockNAT.EXPECT().GetMapping("tcp", 1234).Return(externalAddr, true)
|
|
m6, _ := ma.StringCast("/ip4/1.2.3.4/tcp/4321/ws")
|
|
m7, _ := ma.StringCast("/ip4/0.0.0.0/tcp/1234/ws")
|
|
require.Equal(t, m6, m.GetMapping(m7))
|
|
|
|
// make sure this works for WebTransport addresses as well
|
|
mockNAT.EXPECT().GetMapping("udp", 1234).Return(externalAddr, true)
|
|
m8, _ := ma.StringCast("/ip4/1.2.3.4/udp/4321/quic-v1/webtransport")
|
|
m9, _ := ma.StringCast("/ip4/0.0.0.0/udp/1234/quic-v1/webtransport")
|
|
require.Equal(t, m8, m.GetMapping(m9))
|
|
}
|
|
|
|
func TestAddAndRemoveListeners(t *testing.T) {
|
|
mockNAT, reset := setupMockNAT(t)
|
|
defer reset()
|
|
|
|
sw := swarmt.GenSwarm(t)
|
|
defer sw.Close()
|
|
m := newNATManager(sw)
|
|
require.Eventually(t, func() bool {
|
|
m.natMx.Lock()
|
|
defer m.natMx.Unlock()
|
|
return m.nat != nil
|
|
}, time.Second, time.Millisecond)
|
|
|
|
added := make(chan struct{}, 1)
|
|
// add a TCP listener
|
|
mockNAT.EXPECT().AddMapping(gomock.Any(), "tcp", 1234).Do(func(context.Context, string, int) { added <- struct{}{} })
|
|
m1, _ := ma.StringCast("/ip4/0.0.0.0/tcp/1234")
|
|
require.NoError(t, sw.Listen(m1))
|
|
select {
|
|
case <-added:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("didn't receive call to AddMapping")
|
|
}
|
|
|
|
// add a QUIC listener
|
|
mockNAT.EXPECT().AddMapping(gomock.Any(), "udp", 1234).Do(func(context.Context, string, int) { added <- struct{}{} })
|
|
m2, _ := ma.StringCast("/ip4/0.0.0.0/udp/1234/quic-v1")
|
|
require.NoError(t, sw.Listen(m2))
|
|
select {
|
|
case <-added:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("didn't receive call to AddMapping")
|
|
}
|
|
|
|
// remove the QUIC listener
|
|
mockNAT.EXPECT().RemoveMapping(gomock.Any(), "udp", 1234).Do(func(context.Context, string, int) { added <- struct{}{} })
|
|
m3, _ := ma.StringCast("/ip4/0.0.0.0/udp/1234/quic-v1")
|
|
sw.ListenClose(m3)
|
|
select {
|
|
case <-added:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("didn't receive call to RemoveMapping")
|
|
}
|
|
|
|
// test shutdown
|
|
mockNAT.EXPECT().RemoveMapping(gomock.Any(), "tcp", 1234).MaxTimes(1)
|
|
mockNAT.EXPECT().Close().MaxTimes(1)
|
|
}
|