Merge pull request #258 from libp2p/fix-connection-gating-tests

run connection gating tests on both TCP and QUIC

swarm_test.go

@@ -292,6 +292,7 @@ func TestConnectionGating(t *testing.T) {
 		p1ConnectednessToP2 network.Connectedness
 		p2ConnectednessToP1 network.Connectedness
 		isP1OutboundErr     bool
+		disableOnQUIC       bool
 	}{
 		"no gating": {
 			p1ConnectednessToP2: network.Connected,
@@ -324,6 +325,8 @@ func TestConnectionGating(t *testing.T) {
 			p1ConnectednessToP2: network.NotConnected,
 			p2ConnectednessToP1: network.NotConnected,
 			isP1OutboundErr:     true,
+			// QUIC gates the connection after completion of the handshake
+			disableOnQUIC: true,
 		},
 		"p2 gates inbound peer dial before multiplexing": {
 			p1Gater: func(c *MockConnectionGater) *MockConnectionGater {
@@ -355,33 +358,43 @@ func TestConnectionGating(t *testing.T) {
 	}
 
 	for n, tc := range tcs {
-		t.Run(n, func(t *testing.T) {
-			p1Gater := DefaultMockConnectionGater()
-			p2Gater := DefaultMockConnectionGater()
-			if tc.p1Gater != nil {
-				p1Gater = tc.p1Gater(p1Gater)
-			}
-			if tc.p2Gater != nil {
-				p2Gater = tc.p2Gater(p2Gater)
+		for _, useQuic := range []bool{false, true} {
+			trString := "TCP"
+			optTransport := OptDisableQUIC
+			if useQuic {
+				if tc.disableOnQUIC {
+					continue
+				}
+				trString = "QUIC"
+				optTransport = OptDisableTCP
 			}
+			t.Run(fmt.Sprintf("%s %s", n, trString), func(t *testing.T) {
+				p1Gater := DefaultMockConnectionGater()
+				p2Gater := DefaultMockConnectionGater()
+				if tc.p1Gater != nil {
+					p1Gater = tc.p1Gater(p1Gater)
+				}
+				if tc.p2Gater != nil {
+					p2Gater = tc.p2Gater(p2Gater)
+				}
 
-			sw1 := GenSwarm(t, ctx, OptConnGater(p1Gater))
-			sw2 := GenSwarm(t, ctx, OptConnGater(p2Gater))
-
-			p1 := sw1.LocalPeer()
-			p2 := sw2.LocalPeer()
-			sw1.Peerstore().AddAddr(p2, sw2.ListenAddresses()[0], peerstore.PermanentAddrTTL)
-			// 1 -> 2
-			_, err := sw1.DialPeer(ctx, p2)
+				sw1 := GenSwarm(t, ctx, OptConnGater(p1Gater), optTransport)
+				sw2 := GenSwarm(t, ctx, OptConnGater(p2Gater), optTransport)
 
-			require.Equal(t, tc.isP1OutboundErr, err != nil, n)
-			require.Equal(t, tc.p1ConnectednessToP2, sw1.Connectedness(p2), n)
+				p1 := sw1.LocalPeer()
+				p2 := sw2.LocalPeer()
+				sw1.Peerstore().AddAddr(p2, sw2.ListenAddresses()[0], peerstore.PermanentAddrTTL)
+				// 1 -> 2
+				_, err := sw1.DialPeer(ctx, p2)
 
-			require.Eventually(t, func() bool {
-				return tc.p2ConnectednessToP1 == sw2.Connectedness(p1)
-			}, 2*time.Second, 100*time.Millisecond, n)
-		})
+				require.Equal(t, tc.isP1OutboundErr, err != nil, n)
+				require.Equal(t, tc.p1ConnectednessToP2, sw1.Connectedness(p2), n)
 
+				require.Eventually(t, func() bool {
+					return tc.p2ConnectednessToP1 == sw2.Connectedness(p1)
+				}, 2*time.Second, 100*time.Millisecond, n)
+			})
+		}
 	}
 }
 

testing/testing.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"testing"
 
+	csms "github.com/libp2p/go-conn-security-multistream"
 	"github.com/libp2p/go-libp2p-core/connmgr"
 	"github.com/libp2p/go-libp2p-core/control"
 	"github.com/libp2p/go-libp2p-core/crypto"
@@ -12,10 +13,8 @@ import (
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/libp2p/go-libp2p-core/peerstore"
 	"github.com/libp2p/go-libp2p-core/sec/insecure"
+	"github.com/libp2p/go-libp2p-peerstore/pstoremem"
 	quic "github.com/libp2p/go-libp2p-quic-transport"
-
-	csms "github.com/libp2p/go-conn-security-multistream"
-	pstoremem "github.com/libp2p/go-libp2p-peerstore/pstoremem"
 	swarm "github.com/libp2p/go-libp2p-swarm"
 	"github.com/libp2p/go-libp2p-testing/net"
 	tptu "github.com/libp2p/go-libp2p-transport-upgrader"
@@ -23,13 +22,15 @@ import (
 	msmux "github.com/libp2p/go-stream-muxer-multistream"
 	"github.com/libp2p/go-tcp-transport"
 
-	goprocess "github.com/jbenet/goprocess"
+	"github.com/jbenet/goprocess"
 	ma "github.com/multiformats/go-multiaddr"
 )
 
 type config struct {
 	disableReuseport bool
 	dialOnly         bool
+	disableTCP       bool
+	disableQUIC      bool
 	connectionGater  connmgr.ConnectionGater
 	sk               crypto.PrivKey
 }
@@ -47,6 +48,16 @@ var OptDialOnly Option = func(_ *testing.T, c *config) {
 	c.dialOnly = true
 }
 
+// OptDisableTCP disables TCP.
+var OptDisableTCP Option = func(_ *testing.T, c *config) {
+	c.disableTCP = true
+}
+
+// OptDisableQUIC disables QUIC.
+var OptDisableQUIC Option = func(_ *testing.T, c *config) {
+	c.disableQUIC = true
+}
+
 // OptConnGater configures the given connection gater on the test
 func OptConnGater(cg connmgr.ConnectionGater) Option {
 	return func(_ *testing.T, c *config) {
@@ -111,33 +122,36 @@ func GenSwarm(t *testing.T, ctx context.Context, opts ...Option) *swarm.Swarm {
 
 	upgrader := GenUpgrader(s)
 	upgrader.ConnGater = cfg.connectionGater
-	tcpTransport := tcp.NewTCPTransport(upgrader)
-	tcpTransport.DisableReuseport = cfg.disableReuseport
-
-	quicTransport, err := quic.NewTransport(p.PrivKey, nil, nil)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	if err := s.AddTransport(tcpTransport); err != nil {
-		t.Fatal(err)
-	}
 
-	if err := s.AddTransport(quicTransport); err != nil {
-		t.Fatal(err)
+	if !cfg.disableTCP {
+		tcpTransport := tcp.NewTCPTransport(upgrader)
+		tcpTransport.DisableReuseport = cfg.disableReuseport
+		if err := s.AddTransport(tcpTransport); err != nil {
+			t.Fatal(err)
+		}
+		if !cfg.dialOnly {
+			if err := s.Listen(p.Addr); err != nil {
+				t.Fatal(err)
+			}
+		}
 	}
-
-	if !cfg.dialOnly {
-		if err := s.Listen(p.Addr); err != nil {
+	if !cfg.disableQUIC {
+		quicTransport, err := quic.NewTransport(p.PrivKey, nil, cfg.connectionGater)
+		if err != nil {
 			t.Fatal(err)
 		}
-		if err := s.Listen(ma.StringCast("/ip4/127.0.0.1/udp/0/quic")); err != nil {
+		if err := s.AddTransport(quicTransport); err != nil {
 			t.Fatal(err)
 		}
-
+		if !cfg.dialOnly {
+			if err := s.Listen(ma.StringCast("/ip4/127.0.0.1/udp/0/quic")); err != nil {
+				t.Fatal(err)
+			}
+		}
+	}
+	if !cfg.dialOnly {
 		s.Peerstore().AddAddrs(p.ID, s.ListenAddresses(), peerstore.PermanentAddrTTL)
 	}
-
 	return s
 }