ceremonyclient/go-libp2p/p2p/net/upgrader/upgrader_test.go

205 lines
6.4 KiB
Go
Raw Normal View History

2023-08-21 03:50:38 +00:00
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"
"github.com/golang/mock/gomock"
ma "github.com/multiformats/go-multiaddr"
manet "github.com/multiformats/go-multiaddr/net"
"github.com/stretchr/testify/require"
)
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)
})
}