mirror of
https://source.quilibrium.com/quilibrium/ceremonyclient.git
synced 2024-11-15 20:55:17 +00:00
205 lines
6.4 KiB
Go
205 lines
6.4 KiB
Go
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)
|
|
})
|
|
}
|