From af9bca50131f952a42aa2f194c4c5fa3f9a04d47 Mon Sep 17 00:00:00 2001
From: William Banfield <4561443+williambanfield@users.noreply.github.com>
Date: Tue, 28 Jun 2022 16:07:15 -0400
Subject: [PATCH] p2p: set empty timeouts to configed values. (manual backport
 of #8847) (#8869)

* regenerate mocks using newer style

* p2p: set empty timeouts to small values. (#8847)

These timeouts default to 'do not time out' if they are not set. This times up resources, potentially indefinitely. If node on the other side of the the handshake is up but unresponsive, the[ handshake call](https://github.com/tendermint/tendermint/blob/edec79448aa1d62b84683b1b22e12e145dbdda7c/internal/p2p/router.go#L720) will _never_ return.

* fix light client select statement
---
 abci/client/mocks/client.go                   | 15 ++++++++
 internal/consensus/mocks/cons_sync_reactor.go | 15 ++++++++
 internal/evidence/mocks/block_store.go        | 15 ++++++++
 internal/p2p/mocks/connection.go              | 35 ++++++++++++++-----
 internal/p2p/mocks/peer.go                    | 15 ++++++++
 internal/p2p/mocks/transport.go               | 15 ++++++++
 internal/p2p/peer_test.go                     |  6 ++--
 internal/p2p/router.go                        |  8 +----
 internal/p2p/router_test.go                   | 14 ++++----
 internal/p2p/switch.go                        |  6 ++--
 internal/p2p/switch_test.go                   |  8 ++---
 internal/p2p/test_util.go                     |  2 +-
 internal/p2p/transport.go                     |  3 +-
 internal/p2p/transport_mconn.go               | 19 ++++++++--
 internal/p2p/transport_memory.go              |  8 +++++
 internal/p2p/transport_test.go                | 12 +++----
 internal/proxy/mocks/app_conn_consensus.go    | 15 ++++++++
 internal/proxy/mocks/app_conn_mempool.go      | 15 ++++++++
 internal/proxy/mocks/app_conn_query.go        | 15 ++++++++
 internal/proxy/mocks/app_conn_snapshot.go     | 15 ++++++++
 internal/state/indexer/mocks/event_sink.go    | 15 ++++++++
 internal/state/mocks/block_store.go           | 15 ++++++++
 internal/state/mocks/evidence_pool.go         | 15 ++++++++
 internal/state/mocks/store.go                 | 15 ++++++++
 internal/statesync/mocks/state_provider.go    | 15 ++++++++
 light/client.go                               | 11 +++---
 light/provider/mocks/provider.go              | 15 ++++++++
 light/rpc/mocks/light_client.go               | 15 ++++++++
 node/node.go                                  |  4 ++-
 rpc/client/mocks/client.go                    | 15 ++++++++
 scripts/mockery_generate.sh                   | 14 +++++++-
 31 files changed, 354 insertions(+), 51 deletions(-)

diff --git a/abci/client/mocks/client.go b/abci/client/mocks/client.go
index 664646e61..cfe5ea7af 100644
--- a/abci/client/mocks/client.go
+++ b/abci/client/mocks/client.go
@@ -801,3 +801,18 @@ func (_m *Client) String() string {
 func (_m *Client) Wait() {
 	_m.Called()
 }
+
+type NewClientT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewClient creates a new instance of Client. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewClient(t NewClientT) *Client {
+	mock := &Client{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/consensus/mocks/cons_sync_reactor.go b/internal/consensus/mocks/cons_sync_reactor.go
index 5ac592f0d..2c694742b 100644
--- a/internal/consensus/mocks/cons_sync_reactor.go
+++ b/internal/consensus/mocks/cons_sync_reactor.go
@@ -26,3 +26,18 @@ func (_m *ConsSyncReactor) SetStateSyncingMetrics(_a0 float64) {
 func (_m *ConsSyncReactor) SwitchToConsensus(_a0 state.State, _a1 bool) {
 	_m.Called(_a0, _a1)
 }
+
+type NewConsSyncReactorT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewConsSyncReactor creates a new instance of ConsSyncReactor. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewConsSyncReactor(t NewConsSyncReactorT) *ConsSyncReactor {
+	mock := &ConsSyncReactor{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/evidence/mocks/block_store.go b/internal/evidence/mocks/block_store.go
index ef3346b2a..b0c67ff87 100644
--- a/internal/evidence/mocks/block_store.go
+++ b/internal/evidence/mocks/block_store.go
@@ -57,3 +57,18 @@ func (_m *BlockStore) LoadBlockMeta(height int64) *types.BlockMeta {
 
 	return r0
 }
+
+type NewBlockStoreT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewBlockStore creates a new instance of BlockStore. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewBlockStore(t NewBlockStoreT) *BlockStore {
+	mock := &BlockStore{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/p2p/mocks/connection.go b/internal/p2p/mocks/connection.go
index 6c6174117..e5ba9584a 100644
--- a/internal/p2p/mocks/connection.go
+++ b/internal/p2p/mocks/connection.go
@@ -13,6 +13,8 @@ import (
 
 	p2p "github.com/tendermint/tendermint/internal/p2p"
 
+	time "time"
+
 	types "github.com/tendermint/tendermint/types"
 )
 
@@ -49,20 +51,20 @@ func (_m *Connection) FlushClose() error {
 	return r0
 }
 
-// Handshake provides a mock function with given fields: _a0, _a1, _a2
-func (_m *Connection) Handshake(_a0 context.Context, _a1 types.NodeInfo, _a2 crypto.PrivKey) (types.NodeInfo, crypto.PubKey, error) {
-	ret := _m.Called(_a0, _a1, _a2)
+// Handshake provides a mock function with given fields: _a0, _a1, _a2, _a3
+func (_m *Connection) Handshake(_a0 context.Context, _a1 time.Duration, _a2 types.NodeInfo, _a3 crypto.PrivKey) (types.NodeInfo, crypto.PubKey, error) {
+	ret := _m.Called(_a0, _a1, _a2, _a3)
 
 	var r0 types.NodeInfo
-	if rf, ok := ret.Get(0).(func(context.Context, types.NodeInfo, crypto.PrivKey) types.NodeInfo); ok {
-		r0 = rf(_a0, _a1, _a2)
+	if rf, ok := ret.Get(0).(func(context.Context, time.Duration, types.NodeInfo, crypto.PrivKey) types.NodeInfo); ok {
+		r0 = rf(_a0, _a1, _a2, _a3)
 	} else {
 		r0 = ret.Get(0).(types.NodeInfo)
 	}
 
 	var r1 crypto.PubKey
-	if rf, ok := ret.Get(1).(func(context.Context, types.NodeInfo, crypto.PrivKey) crypto.PubKey); ok {
-		r1 = rf(_a0, _a1, _a2)
+	if rf, ok := ret.Get(1).(func(context.Context, time.Duration, types.NodeInfo, crypto.PrivKey) crypto.PubKey); ok {
+		r1 = rf(_a0, _a1, _a2, _a3)
 	} else {
 		if ret.Get(1) != nil {
 			r1 = ret.Get(1).(crypto.PubKey)
@@ -70,8 +72,8 @@ func (_m *Connection) Handshake(_a0 context.Context, _a1 types.NodeInfo, _a2 cry
 	}
 
 	var r2 error
-	if rf, ok := ret.Get(2).(func(context.Context, types.NodeInfo, crypto.PrivKey) error); ok {
-		r2 = rf(_a0, _a1, _a2)
+	if rf, ok := ret.Get(2).(func(context.Context, time.Duration, types.NodeInfo, crypto.PrivKey) error); ok {
+		r2 = rf(_a0, _a1, _a2, _a3)
 	} else {
 		r2 = ret.Error(2)
 	}
@@ -206,3 +208,18 @@ func (_m *Connection) TrySendMessage(_a0 p2p.ChannelID, _a1 []byte) (bool, error
 
 	return r0, r1
 }
+
+type NewConnectionT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewConnection creates a new instance of Connection. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewConnection(t NewConnectionT) *Connection {
+	mock := &Connection{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/p2p/mocks/peer.go b/internal/p2p/mocks/peer.go
index b905c1156..021c905f2 100644
--- a/internal/p2p/mocks/peer.go
+++ b/internal/p2p/mocks/peer.go
@@ -332,3 +332,18 @@ func (_m *Peer) TrySend(_a0 byte, _a1 []byte) bool {
 func (_m *Peer) Wait() {
 	_m.Called()
 }
+
+type NewPeerT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewPeer creates a new instance of Peer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewPeer(t NewPeerT) *Peer {
+	mock := &Peer{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/p2p/mocks/transport.go b/internal/p2p/mocks/transport.go
index 82bd670cb..46d825501 100644
--- a/internal/p2p/mocks/transport.go
+++ b/internal/p2p/mocks/transport.go
@@ -119,3 +119,18 @@ func (_m *Transport) String() string {
 
 	return r0
 }
+
+type NewTransportT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewTransport creates a new instance of Transport. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewTransport(t NewTransportT) *Transport {
+	mock := &Transport{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/p2p/peer_test.go b/internal/p2p/peer_test.go
index dfe7bc798..dad7b98b5 100644
--- a/internal/p2p/peer_test.go
+++ b/internal/p2p/peer_test.go
@@ -90,7 +90,7 @@ func createOutboundPeerAndPerformHandshake(
 	if err != nil {
 		return nil, err
 	}
-	peerInfo, _, err := pc.conn.Handshake(context.Background(), ourNodeInfo, pk)
+	peerInfo, _, err := pc.conn.Handshake(context.Background(), 0, ourNodeInfo, pk)
 	if err != nil {
 		return nil, err
 	}
@@ -187,7 +187,7 @@ func (rp *remotePeer) Dial(addr *NetAddress) (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	_, _, err = pc.conn.Handshake(context.Background(), rp.nodeInfo(), rp.PrivKey)
+	_, _, err = pc.conn.Handshake(context.Background(), 0, rp.nodeInfo(), rp.PrivKey)
 	if err != nil {
 		return nil, err
 	}
@@ -213,7 +213,7 @@ func (rp *remotePeer) accept() {
 		if err != nil {
 			golog.Printf("Failed to create a peer: %+v", err)
 		}
-		_, _, err = pc.conn.Handshake(context.Background(), rp.nodeInfo(), rp.PrivKey)
+		_, _, err = pc.conn.Handshake(context.Background(), 0, rp.nodeInfo(), rp.PrivKey)
 		if err != nil {
 			golog.Printf("Failed to handshake a peer: %+v", err)
 		}
diff --git a/internal/p2p/router.go b/internal/p2p/router.go
index 81a8928a4..dc717565f 100644
--- a/internal/p2p/router.go
+++ b/internal/p2p/router.go
@@ -829,13 +829,7 @@ func (r *Router) handshakePeer(
 	expectID types.NodeID,
 ) (types.NodeInfo, crypto.PubKey, error) {
 
-	if r.options.HandshakeTimeout > 0 {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, r.options.HandshakeTimeout)
-		defer cancel()
-	}
-
-	peerInfo, peerKey, err := conn.Handshake(ctx, r.nodeInfo, r.privKey)
+	peerInfo, peerKey, err := conn.Handshake(ctx, r.options.HandshakeTimeout, r.nodeInfo, r.privKey)
 	if err != nil {
 		return peerInfo, peerKey, err
 	}
diff --git a/internal/p2p/router_test.go b/internal/p2p/router_test.go
index e8494fdf4..37ce95768 100644
--- a/internal/p2p/router_test.go
+++ b/internal/p2p/router_test.go
@@ -352,7 +352,7 @@ func TestRouter_AcceptPeers(t *testing.T) {
 			closer := tmsync.NewCloser()
 			mockConnection := &mocks.Connection{}
 			mockConnection.On("String").Maybe().Return("mock")
-			mockConnection.On("Handshake", mock.Anything, selfInfo, selfKey).
+			mockConnection.On("Handshake", mock.Anything, mock.Anything, selfInfo, selfKey).
 				Return(tc.peerInfo, tc.peerKey, nil)
 			mockConnection.On("Close").Run(func(_ mock.Arguments) { closer.Close() }).Return(nil)
 			mockConnection.On("RemoteEndpoint").Return(p2p.Endpoint{})
@@ -462,7 +462,7 @@ func TestRouter_AcceptPeers_HeadOfLineBlocking(t *testing.T) {
 
 	mockConnection := &mocks.Connection{}
 	mockConnection.On("String").Maybe().Return("mock")
-	mockConnection.On("Handshake", mock.Anything, selfInfo, selfKey).
+	mockConnection.On("Handshake", mock.Anything, mock.Anything, selfInfo, selfKey).
 		WaitUntil(closeCh).Return(types.NodeInfo{}, nil, io.EOF)
 	mockConnection.On("Close").Return(nil)
 	mockConnection.On("RemoteEndpoint").Return(p2p.Endpoint{})
@@ -543,7 +543,7 @@ func TestRouter_DialPeers(t *testing.T) {
 			mockConnection := &mocks.Connection{}
 			mockConnection.On("String").Maybe().Return("mock")
 			if tc.dialErr == nil {
-				mockConnection.On("Handshake", mock.Anything, selfInfo, selfKey).
+				mockConnection.On("Handshake", mock.Anything, mock.Anything, selfInfo, selfKey).
 					Return(tc.peerInfo, tc.peerKey, nil)
 				mockConnection.On("Close").Run(func(_ mock.Arguments) { closer.Close() }).Return(nil)
 			}
@@ -630,7 +630,7 @@ func TestRouter_DialPeers_Parallel(t *testing.T) {
 
 	mockConnection := &mocks.Connection{}
 	mockConnection.On("String").Maybe().Return("mock")
-	mockConnection.On("Handshake", mock.Anything, selfInfo, selfKey).
+	mockConnection.On("Handshake", mock.Anything, mock.Anything, selfInfo, selfKey).
 		WaitUntil(closeCh).Return(types.NodeInfo{}, nil, io.EOF)
 	mockConnection.On("Close").Return(nil)
 
@@ -710,7 +710,7 @@ func TestRouter_EvictPeers(t *testing.T) {
 
 	mockConnection := &mocks.Connection{}
 	mockConnection.On("String").Maybe().Return("mock")
-	mockConnection.On("Handshake", mock.Anything, selfInfo, selfKey).
+	mockConnection.On("Handshake", mock.Anything, mock.Anything, selfInfo, selfKey).
 		Return(peerInfo, peerKey.PubKey(), nil)
 	mockConnection.On("ReceiveMessage").WaitUntil(closeCh).Return(chID, nil, io.EOF)
 	mockConnection.On("RemoteEndpoint").Return(p2p.Endpoint{})
@@ -779,7 +779,7 @@ func TestRouter_ChannelCompatability(t *testing.T) {
 
 	mockConnection := &mocks.Connection{}
 	mockConnection.On("String").Maybe().Return("mock")
-	mockConnection.On("Handshake", mock.Anything, selfInfo, selfKey).
+	mockConnection.On("Handshake", mock.Anything, mock.Anything, selfInfo, selfKey).
 		Return(incompatiblePeer, peerKey.PubKey(), nil)
 	mockConnection.On("RemoteEndpoint").Return(p2p.Endpoint{})
 	mockConnection.On("Close").Return(nil)
@@ -828,7 +828,7 @@ func TestRouter_DontSendOnInvalidChannel(t *testing.T) {
 
 	mockConnection := &mocks.Connection{}
 	mockConnection.On("String").Maybe().Return("mock")
-	mockConnection.On("Handshake", mock.Anything, selfInfo, selfKey).
+	mockConnection.On("Handshake", mock.Anything, mock.Anything, selfInfo, selfKey).
 		Return(peer, peerKey.PubKey(), nil)
 	mockConnection.On("RemoteEndpoint").Return(p2p.Endpoint{})
 	mockConnection.On("Close").Return(nil)
diff --git a/internal/p2p/switch.go b/internal/p2p/switch.go
index 60c0c7deb..9e1b0311b 100644
--- a/internal/p2p/switch.go
+++ b/internal/p2p/switch.go
@@ -865,11 +865,11 @@ func (sw *Switch) handshakePeer(
 	c Connection,
 	expectPeerID types.NodeID,
 ) (types.NodeInfo, crypto.PubKey, error) {
-	// Moved from transport and hardcoded until legacy P2P stack removal.
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	peerInfo, peerKey, err := c.Handshake(ctx, sw.nodeInfo, sw.nodeKey.PrivKey)
+	// Moved timeout from transport and hardcoded until legacy P2P stack removal.
+	peerInfo, peerKey, err := c.Handshake(ctx, 5*time.Second, sw.nodeInfo, sw.nodeKey.PrivKey)
 	if err != nil {
 		return peerInfo, peerKey, ErrRejected{
 			conn:          c.(*mConnConnection).conn,
diff --git a/internal/p2p/switch_test.go b/internal/p2p/switch_test.go
index 8cb755c9f..c68cfceaf 100644
--- a/internal/p2p/switch_test.go
+++ b/internal/p2p/switch_test.go
@@ -267,7 +267,7 @@ func TestSwitchPeerFilter(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	peerInfo, _, err := c.Handshake(ctx, sw.nodeInfo, sw.nodeKey.PrivKey)
+	peerInfo, _, err := c.Handshake(ctx, 0, sw.nodeInfo, sw.nodeKey.PrivKey)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -324,7 +324,7 @@ func TestSwitchPeerFilterTimeout(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	peerInfo, _, err := c.Handshake(ctx, sw.nodeInfo, sw.nodeKey.PrivKey)
+	peerInfo, _, err := c.Handshake(ctx, 0, sw.nodeInfo, sw.nodeKey.PrivKey)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -360,7 +360,7 @@ func TestSwitchPeerFilterDuplicate(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	peerInfo, _, err := c.Handshake(ctx, sw.nodeInfo, sw.nodeKey.PrivKey)
+	peerInfo, _, err := c.Handshake(ctx, 0, sw.nodeInfo, sw.nodeKey.PrivKey)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -415,7 +415,7 @@ func TestSwitchStopsNonPersistentPeerOnError(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	peerInfo, _, err := c.Handshake(ctx, sw.nodeInfo, sw.nodeKey.PrivKey)
+	peerInfo, _, err := c.Handshake(ctx, 0, sw.nodeInfo, sw.nodeKey.PrivKey)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/internal/p2p/test_util.go b/internal/p2p/test_util.go
index b2851646d..ae21ba4d7 100644
--- a/internal/p2p/test_util.go
+++ b/internal/p2p/test_util.go
@@ -126,7 +126,7 @@ func (sw *Switch) addPeerWithConnection(conn net.Conn) error {
 		}
 		return err
 	}
-	peerNodeInfo, _, err := pc.conn.Handshake(context.Background(), sw.nodeInfo, sw.nodeKey.PrivKey)
+	peerNodeInfo, _, err := pc.conn.Handshake(context.Background(), 0, sw.nodeInfo, sw.nodeKey.PrivKey)
 	if err != nil {
 		if err := conn.Close(); err != nil {
 			sw.Logger.Error("Error closing connection", "err", err)
diff --git a/internal/p2p/transport.go b/internal/p2p/transport.go
index a3245dfc8..2e4d26abd 100644
--- a/internal/p2p/transport.go
+++ b/internal/p2p/transport.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"time"
 
 	"github.com/tendermint/tendermint/crypto"
 	"github.com/tendermint/tendermint/internal/p2p/conn"
@@ -84,7 +85,7 @@ type Connection interface {
 	// FIXME: The handshake should really be the Router's responsibility, but
 	// that requires the connection interface to be byte-oriented rather than
 	// message-oriented (see comment above).
-	Handshake(context.Context, types.NodeInfo, crypto.PrivKey) (types.NodeInfo, crypto.PubKey, error)
+	Handshake(context.Context, time.Duration, types.NodeInfo, crypto.PrivKey) (types.NodeInfo, crypto.PubKey, error)
 
 	// ReceiveMessage returns the next message received on the connection,
 	// blocking until one is available. Returns io.EOF if closed.
diff --git a/internal/p2p/transport_mconn.go b/internal/p2p/transport_mconn.go
index eca261476..4d18d896b 100644
--- a/internal/p2p/transport_mconn.go
+++ b/internal/p2p/transport_mconn.go
@@ -9,6 +9,7 @@ import (
 	"net"
 	"strconv"
 	"sync"
+	"time"
 
 	"golang.org/x/net/netutil"
 
@@ -255,6 +256,7 @@ func newMConnConnection(
 // Handshake implements Connection.
 func (c *mConnConnection) Handshake(
 	ctx context.Context,
+	timeout time.Duration,
 	nodeInfo types.NodeInfo,
 	privKey crypto.PrivKey,
 ) (types.NodeInfo, crypto.PubKey, error) {
@@ -264,6 +266,12 @@ func (c *mConnConnection) Handshake(
 		peerKey  crypto.PubKey
 		errCh    = make(chan error, 1)
 	)
+	handshakeCtx := ctx
+	if timeout > 0 {
+		var cancel context.CancelFunc
+		handshakeCtx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+	}
 	// To handle context cancellation, we need to do the handshake in a
 	// goroutine and abort the blocking network calls by closing the connection
 	// when the context is canceled.
@@ -276,12 +284,17 @@ func (c *mConnConnection) Handshake(
 			}
 		}()
 		var err error
-		mconn, peerInfo, peerKey, err = c.handshake(ctx, nodeInfo, privKey)
-		errCh <- err
+		mconn, peerInfo, peerKey, err = c.handshake(handshakeCtx, nodeInfo, privKey)
+
+		select {
+		case errCh <- err:
+		case <-handshakeCtx.Done():
+		}
+
 	}()
 
 	select {
-	case <-ctx.Done():
+	case <-handshakeCtx.Done():
 		_ = c.Close()
 		return types.NodeInfo{}, nil, ctx.Err()
 
diff --git a/internal/p2p/transport_memory.go b/internal/p2p/transport_memory.go
index 09a387254..f2d1d0c72 100644
--- a/internal/p2p/transport_memory.go
+++ b/internal/p2p/transport_memory.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net"
 	"sync"
+	"time"
 
 	"github.com/tendermint/tendermint/crypto"
 	tmsync "github.com/tendermint/tendermint/internal/libs/sync"
@@ -270,9 +271,16 @@ func (c *MemoryConnection) Status() conn.ConnectionStatus {
 // Handshake implements Connection.
 func (c *MemoryConnection) Handshake(
 	ctx context.Context,
+	timeout time.Duration,
 	nodeInfo types.NodeInfo,
 	privKey crypto.PrivKey,
 ) (types.NodeInfo, crypto.PubKey, error) {
+	if timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+	}
+
 	select {
 	case c.sendCh <- memoryMessage{nodeInfo: &nodeInfo, pubKey: privKey.PubKey()}:
 		c.logger.Debug("sent handshake", "nodeInfo", nodeInfo)
diff --git a/internal/p2p/transport_test.go b/internal/p2p/transport_test.go
index 1b8ab77f5..de7405177 100644
--- a/internal/p2p/transport_test.go
+++ b/internal/p2p/transport_test.go
@@ -265,7 +265,7 @@ func TestConnection_Handshake(t *testing.T) {
 		errCh := make(chan error, 1)
 		go func() {
 			// Must use assert due to goroutine.
-			peerInfo, peerKey, err := ba.Handshake(ctx, bInfo, bKey)
+			peerInfo, peerKey, err := ba.Handshake(ctx, 0, bInfo, bKey)
 			if err == nil {
 				assert.Equal(t, aInfo, peerInfo)
 				assert.Equal(t, aKey.PubKey(), peerKey)
@@ -273,7 +273,7 @@ func TestConnection_Handshake(t *testing.T) {
 			errCh <- err
 		}()
 
-		peerInfo, peerKey, err := ab.Handshake(ctx, aInfo, aKey)
+		peerInfo, peerKey, err := ab.Handshake(ctx, 0, aInfo, aKey)
 		require.NoError(t, err)
 		require.Equal(t, bInfo, peerInfo)
 		require.Equal(t, bKey.PubKey(), peerKey)
@@ -291,7 +291,7 @@ func TestConnection_HandshakeCancel(t *testing.T) {
 		ab, ba := dialAccept(t, a, b)
 		timeoutCtx, cancel := context.WithTimeout(ctx, 1*time.Minute)
 		cancel()
-		_, _, err := ab.Handshake(timeoutCtx, types.NodeInfo{}, ed25519.GenPrivKey())
+		_, _, err := ab.Handshake(timeoutCtx, 0, types.NodeInfo{}, ed25519.GenPrivKey())
 		require.Error(t, err)
 		require.Equal(t, context.Canceled, err)
 		_ = ab.Close()
@@ -301,7 +301,7 @@ func TestConnection_HandshakeCancel(t *testing.T) {
 		ab, ba = dialAccept(t, a, b)
 		timeoutCtx, cancel = context.WithTimeout(ctx, 200*time.Millisecond)
 		defer cancel()
-		_, _, err = ab.Handshake(timeoutCtx, types.NodeInfo{}, ed25519.GenPrivKey())
+		_, _, err = ab.Handshake(timeoutCtx, 0, types.NodeInfo{}, ed25519.GenPrivKey())
 		require.Error(t, err)
 		require.Equal(t, context.DeadlineExceeded, err)
 		_ = ab.Close()
@@ -630,13 +630,13 @@ func dialAcceptHandshake(t *testing.T, a, b p2p.Transport) (p2p.Connection, p2p.
 	go func() {
 		privKey := ed25519.GenPrivKey()
 		nodeInfo := types.NodeInfo{NodeID: types.NodeIDFromPubKey(privKey.PubKey())}
-		_, _, err := ba.Handshake(ctx, nodeInfo, privKey)
+		_, _, err := ba.Handshake(ctx, 0, nodeInfo, privKey)
 		errCh <- err
 	}()
 
 	privKey := ed25519.GenPrivKey()
 	nodeInfo := types.NodeInfo{NodeID: types.NodeIDFromPubKey(privKey.PubKey())}
-	_, _, err := ab.Handshake(ctx, nodeInfo, privKey)
+	_, _, err := ab.Handshake(ctx, 0, nodeInfo, privKey)
 	require.NoError(t, err)
 
 	timer := time.NewTimer(2 * time.Second)
diff --git a/internal/proxy/mocks/app_conn_consensus.go b/internal/proxy/mocks/app_conn_consensus.go
index fa93b0931..bff6bf6a4 100644
--- a/internal/proxy/mocks/app_conn_consensus.go
+++ b/internal/proxy/mocks/app_conn_consensus.go
@@ -150,3 +150,18 @@ func (_m *AppConnConsensus) InitChainSync(_a0 context.Context, _a1 types.Request
 func (_m *AppConnConsensus) SetResponseCallback(_a0 abciclient.Callback) {
 	_m.Called(_a0)
 }
+
+type NewAppConnConsensusT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewAppConnConsensus creates a new instance of AppConnConsensus. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewAppConnConsensus(t NewAppConnConsensusT) *AppConnConsensus {
+	mock := &AppConnConsensus{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/proxy/mocks/app_conn_mempool.go b/internal/proxy/mocks/app_conn_mempool.go
index 5429d8f90..a17deb7ca 100644
--- a/internal/proxy/mocks/app_conn_mempool.go
+++ b/internal/proxy/mocks/app_conn_mempool.go
@@ -118,3 +118,18 @@ func (_m *AppConnMempool) FlushSync(_a0 context.Context) error {
 func (_m *AppConnMempool) SetResponseCallback(_a0 abciclient.Callback) {
 	_m.Called(_a0)
 }
+
+type NewAppConnMempoolT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewAppConnMempool creates a new instance of AppConnMempool. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewAppConnMempool(t NewAppConnMempoolT) *AppConnMempool {
+	mock := &AppConnMempool{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/proxy/mocks/app_conn_query.go b/internal/proxy/mocks/app_conn_query.go
index 47ac5bef9..9bde883a0 100644
--- a/internal/proxy/mocks/app_conn_query.go
+++ b/internal/proxy/mocks/app_conn_query.go
@@ -97,3 +97,18 @@ func (_m *AppConnQuery) QuerySync(_a0 context.Context, _a1 types.RequestQuery) (
 
 	return r0, r1
 }
+
+type NewAppConnQueryT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewAppConnQuery creates a new instance of AppConnQuery. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewAppConnQuery(t NewAppConnQueryT) *AppConnQuery {
+	mock := &AppConnQuery{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/proxy/mocks/app_conn_snapshot.go b/internal/proxy/mocks/app_conn_snapshot.go
index 0b6f10ce1..36df29781 100644
--- a/internal/proxy/mocks/app_conn_snapshot.go
+++ b/internal/proxy/mocks/app_conn_snapshot.go
@@ -120,3 +120,18 @@ func (_m *AppConnSnapshot) OfferSnapshotSync(_a0 context.Context, _a1 types.Requ
 
 	return r0, r1
 }
+
+type NewAppConnSnapshotT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewAppConnSnapshot creates a new instance of AppConnSnapshot. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewAppConnSnapshot(t NewAppConnSnapshotT) *AppConnSnapshot {
+	mock := &AppConnSnapshot{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/state/indexer/mocks/event_sink.go b/internal/state/indexer/mocks/event_sink.go
index 98b32e935..ec5e279a8 100644
--- a/internal/state/indexer/mocks/event_sink.go
+++ b/internal/state/indexer/mocks/event_sink.go
@@ -165,3 +165,18 @@ func (_m *EventSink) Type() indexer.EventSinkType {
 
 	return r0
 }
+
+type NewEventSinkT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewEventSink creates a new instance of EventSink. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewEventSink(t NewEventSinkT) *EventSink {
+	mock := &EventSink{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/state/mocks/block_store.go b/internal/state/mocks/block_store.go
index 563183437..b4d0c7f99 100644
--- a/internal/state/mocks/block_store.go
+++ b/internal/state/mocks/block_store.go
@@ -208,3 +208,18 @@ func (_m *BlockStore) Size() int64 {
 
 	return r0
 }
+
+type NewBlockStoreT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewBlockStore creates a new instance of BlockStore. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewBlockStore(t NewBlockStoreT) *BlockStore {
+	mock := &BlockStore{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/state/mocks/evidence_pool.go b/internal/state/mocks/evidence_pool.go
index 8bf4a9b64..8bc208289 100644
--- a/internal/state/mocks/evidence_pool.go
+++ b/internal/state/mocks/evidence_pool.go
@@ -68,3 +68,18 @@ func (_m *EvidencePool) PendingEvidence(maxBytes int64) ([]types.Evidence, int64
 func (_m *EvidencePool) Update(_a0 state.State, _a1 types.EvidenceList) {
 	_m.Called(_a0, _a1)
 }
+
+type NewEvidencePoolT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewEvidencePool creates a new instance of EvidencePool. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewEvidencePool(t NewEvidencePoolT) *EvidencePool {
+	mock := &EvidencePool{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/state/mocks/store.go b/internal/state/mocks/store.go
index 02c69d3e0..d08ba4c9e 100644
--- a/internal/state/mocks/store.go
+++ b/internal/state/mocks/store.go
@@ -186,3 +186,18 @@ func (_m *Store) SaveValidatorSets(_a0 int64, _a1 int64, _a2 *types.ValidatorSet
 
 	return r0
 }
+
+type NewStoreT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewStore creates a new instance of Store. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewStore(t NewStoreT) *Store {
+	mock := &Store{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/internal/statesync/mocks/state_provider.go b/internal/statesync/mocks/state_provider.go
index b8d681631..17ddb54ac 100644
--- a/internal/statesync/mocks/state_provider.go
+++ b/internal/statesync/mocks/state_provider.go
@@ -82,3 +82,18 @@ func (_m *StateProvider) State(ctx context.Context, height uint64) (state.State,
 
 	return r0, r1
 }
+
+type NewStateProviderT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewStateProvider creates a new instance of StateProvider. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewStateProvider(t NewStateProviderT) *StateProvider {
+	mock := &StateProvider{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/light/client.go b/light/client.go
index 5bff76894..32d4c669b 100644
--- a/light/client.go
+++ b/light/client.go
@@ -1018,7 +1018,12 @@ func (c *Client) findNewPrimary(ctx context.Context, height int64, remove bool)
 
 	// process all the responses as they come in
 	for i := 0; i < cap(witnessResponsesC); i++ {
-		response := <-witnessResponsesC
+		var response witnessResponse
+		select {
+		case response = <-witnessResponsesC:
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
 		switch response.err {
 		// success! We have found a new primary
 		case nil:
@@ -1047,10 +1052,6 @@ func (c *Client) findNewPrimary(ctx context.Context, height int64, remove bool)
 			// return the light block that new primary responded with
 			return response.lb, nil
 
-		// catch canceled contexts or deadlines
-		case context.Canceled, context.DeadlineExceeded:
-			return nil, response.err
-
 		// process benign errors by logging them only
 		case provider.ErrNoResponse, provider.ErrLightBlockNotFound, provider.ErrHeightTooHigh:
 			lastError = response.err
diff --git a/light/provider/mocks/provider.go b/light/provider/mocks/provider.go
index aa36fa2d3..9a5789ce9 100644
--- a/light/provider/mocks/provider.go
+++ b/light/provider/mocks/provider.go
@@ -51,3 +51,18 @@ func (_m *Provider) ReportEvidence(_a0 context.Context, _a1 types.Evidence) erro
 
 	return r0
 }
+
+type NewProviderT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewProvider creates a new instance of Provider. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewProvider(t NewProviderT) *Provider {
+	mock := &Provider{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/light/rpc/mocks/light_client.go b/light/rpc/mocks/light_client.go
index cc32cf649..25e101c86 100644
--- a/light/rpc/mocks/light_client.go
+++ b/light/rpc/mocks/light_client.go
@@ -99,3 +99,18 @@ func (_m *LightClient) VerifyLightBlockAtHeight(ctx context.Context, height int6
 
 	return r0, r1
 }
+
+type NewLightClientT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewLightClient creates a new instance of LightClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewLightClient(t NewLightClientT) *LightClient {
+	mock := &LightClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/node/node.go b/node/node.go
index cd0f31396..e1de314d1 100644
--- a/node/node.go
+++ b/node/node.go
@@ -1246,7 +1246,9 @@ func createAndStartPrivValidatorGRPCClient(
 
 func getRouterConfig(conf *config.Config, proxyApp proxy.AppConns) p2p.RouterOptions {
 	opts := p2p.RouterOptions{
-		QueueType: conf.P2P.QueueType,
+		QueueType:        conf.P2P.QueueType,
+		HandshakeTimeout: conf.P2P.HandshakeTimeout,
+		DialTimeout:      conf.P2P.DialTimeout,
 	}
 
 	if conf.P2P.MaxNumInboundPeers > 0 {
diff --git a/rpc/client/mocks/client.go b/rpc/client/mocks/client.go
index 0a83ef201..0f1581f70 100644
--- a/rpc/client/mocks/client.go
+++ b/rpc/client/mocks/client.go
@@ -800,3 +800,18 @@ func (_m *Client) Validators(ctx context.Context, height *int64, page *int, perP
 
 	return r0, r1
 }
+
+type NewClientT interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewClient creates a new instance of Client. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewClient(t NewClientT) *Client {
+	mock := &Client{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/scripts/mockery_generate.sh b/scripts/mockery_generate.sh
index 382c277bb..2d6f40e63 100755
--- a/scripts/mockery_generate.sh
+++ b/scripts/mockery_generate.sh
@@ -1,3 +1,15 @@
 #!/bin/sh
+#
+# Invoke Mockery v2 to update generated mocks for the given type.
+#
+# This script runs a locally-installed "mockery" if available, otherwise it
+# runs the published Docker container. This legerdemain is so that the CI build
+# and a local build can work off the same script.
+#
+if ! which mockery ; then
+  mockery() {
+    docker run --rm -v "$PWD":/w --workdir=/w vektra/mockery:v2.12.3
+  }
+fi
 
-go run github.com/vektra/mockery/v2 --disable-version-string --case underscore --name $*
+mockery --disable-version-string --case underscore --name "$@"