package upgrader_test import ( "context" "crypto/rand" "errors" "net" "testing" "github.com/libp2p/go-libp2p/core/connmgr" "github.com/libp2p/go-libp2p/core/crypto" "github.com/libp2p/go-libp2p/core/network" mocknetwork "github.com/libp2p/go-libp2p/core/network/mocks" "github.com/libp2p/go-libp2p/core/peer" "github.com/libp2p/go-libp2p/core/sec" "github.com/libp2p/go-libp2p/core/sec/insecure" "github.com/libp2p/go-libp2p/core/transport" "github.com/libp2p/go-libp2p/p2p/muxer/yamux" "github.com/libp2p/go-libp2p/p2p/net/upgrader" ma "github.com/multiformats/go-multiaddr" manet "github.com/multiformats/go-multiaddr/net" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" ) func createUpgrader(t *testing.T) (peer.ID, transport.Upgrader) { return createUpgraderWithMuxers(t, []upgrader.StreamMuxer{{ID: "negotiate", Muxer: &negotiatingMuxer{}}}, nil, nil) } func createUpgraderWithConnGater(t *testing.T, connGater connmgr.ConnectionGater) (peer.ID, transport.Upgrader) { return createUpgraderWithMuxers(t, []upgrader.StreamMuxer{{ID: "negotiate", Muxer: &negotiatingMuxer{}}}, nil, connGater) } func createUpgraderWithResourceManager(t *testing.T, rcmgr network.ResourceManager) (peer.ID, transport.Upgrader) { return createUpgraderWithMuxers(t, []upgrader.StreamMuxer{{ID: "negotiate", Muxer: &negotiatingMuxer{}}}, rcmgr, nil) } func createUpgraderWithOpts(t *testing.T, opts ...upgrader.Option) (peer.ID, transport.Upgrader) { return createUpgraderWithMuxers(t, []upgrader.StreamMuxer{{ID: "negotiate", Muxer: &negotiatingMuxer{}}}, nil, nil, opts...) } func newPeer(t *testing.T) (peer.ID, crypto.PrivKey) { t.Helper() priv, _, err := crypto.GenerateEd25519Key(rand.Reader) require.NoError(t, err) id, err := peer.IDFromPrivateKey(priv) require.NoError(t, err) return id, priv } func createUpgraderWithMuxers(t *testing.T, muxers []upgrader.StreamMuxer, rcmgr network.ResourceManager, connGater connmgr.ConnectionGater, opts ...upgrader.Option) (peer.ID, transport.Upgrader) { id, priv := newPeer(t) u, err := upgrader.New([]sec.SecureTransport{insecure.NewWithIdentity(insecure.ID, id, priv)}, muxers, nil, rcmgr, connGater, opts...) require.NoError(t, err) return id, u } // negotiatingMuxer sets up a new yamux connection // It makes sure that this happens at the same time for client and server. type negotiatingMuxer struct{} func (m *negotiatingMuxer) NewConn(c net.Conn, isServer bool, scope network.PeerScope) (network.MuxedConn, error) { var err error // run a fake muxer negotiation if isServer { _, err = c.Write([]byte("setup")) } else { _, err = c.Read(make([]byte, 5)) } if err != nil { return nil, err } return yamux.DefaultTransport.NewConn(c, isServer, scope) } // blockingMuxer blocks the muxer negotiation until the contained chan is closed type blockingMuxer struct { unblock chan struct{} } var _ network.Multiplexer = &blockingMuxer{} func newBlockingMuxer() *blockingMuxer { return &blockingMuxer{unblock: make(chan struct{})} } func (m *blockingMuxer) NewConn(c net.Conn, isServer bool, scope network.PeerScope) (network.MuxedConn, error) { <-m.unblock return (&negotiatingMuxer{}).NewConn(c, isServer, scope) } func (m *blockingMuxer) Unblock() { close(m.unblock) } // errorMuxer is a muxer that errors while setting up type errorMuxer struct{} var _ network.Multiplexer = &errorMuxer{} func (m *errorMuxer) NewConn(c net.Conn, isServer bool, scope network.PeerScope) (network.MuxedConn, error) { return nil, errors.New("mux error") } func testConn(t *testing.T, clientConn, serverConn transport.CapableConn) { t.Helper() require := require.New(t) cstr, err := clientConn.OpenStream(context.Background()) require.NoError(err) _, err = cstr.Write([]byte("foobar")) require.NoError(err) sstr, err := serverConn.AcceptStream() require.NoError(err) b := make([]byte, 6) _, err = sstr.Read(b) require.NoError(err) require.Equal([]byte("foobar"), b) } func dial(t *testing.T, upgrader transport.Upgrader, raddr ma.Multiaddr, p peer.ID, scope network.ConnManagementScope) (transport.CapableConn, error) { t.Helper() macon, err := manet.Dial(raddr) if err != nil { return nil, err } return upgrader.Upgrade(context.Background(), nil, macon, network.DirOutbound, p, scope) } func TestOutboundConnectionGating(t *testing.T) { require := require.New(t) id, u := createUpgrader(t) ln := createListener(t, u) defer ln.Close() testGater := &testGater{} _, dialUpgrader := createUpgraderWithConnGater(t, testGater) conn, err := dial(t, dialUpgrader, ln.Multiaddr(), id, &network.NullScope{}) require.NoError(err) require.NotNil(conn) _ = conn.Close() // blocking accepts doesn't affect the dialling side, only the listener. testGater.BlockAccept(true) conn, err = dial(t, dialUpgrader, ln.Multiaddr(), id, &network.NullScope{}) require.NoError(err) require.NotNil(conn) _ = conn.Close() // now let's block all connections after being secured. testGater.BlockSecured(true) conn, err = dial(t, dialUpgrader, ln.Multiaddr(), id, &network.NullScope{}) require.Error(err) require.Contains(err.Error(), "gater rejected connection") require.Nil(conn) } func TestOutboundResourceManagement(t *testing.T) { t.Run("successful handshake", func(t *testing.T) { id, upgrader := createUpgrader(t) ln := createListener(t, upgrader) defer ln.Close() ctrl := gomock.NewController(t) defer ctrl.Finish() connScope := mocknetwork.NewMockConnManagementScope(ctrl) gomock.InOrder( connScope.EXPECT().PeerScope(), connScope.EXPECT().SetPeer(id), connScope.EXPECT().PeerScope().Return(&network.NullScope{}), ) _, dialUpgrader := createUpgrader(t) conn, err := dial(t, dialUpgrader, ln.Multiaddr(), id, connScope) require.NoError(t, err) require.NotNil(t, conn) connScope.EXPECT().Done() require.NoError(t, conn.Close()) }) t.Run("failed negotiation", func(t *testing.T) { id, upgrader := createUpgraderWithMuxers(t, []upgrader.StreamMuxer{{ID: "errorMuxer", Muxer: &errorMuxer{}}}, nil, nil) ln := createListener(t, upgrader) defer ln.Close() ctrl := gomock.NewController(t) defer ctrl.Finish() connScope := mocknetwork.NewMockConnManagementScope(ctrl) gomock.InOrder( connScope.EXPECT().PeerScope(), connScope.EXPECT().SetPeer(id), connScope.EXPECT().PeerScope().Return(&network.NullScope{}), connScope.EXPECT().Done(), ) _, dialUpgrader := createUpgrader(t) _, err := dial(t, dialUpgrader, ln.Multiaddr(), id, connScope) require.Error(t, err) }) }