lnd: improve brontide mock

lnd/brontide.go

@@ -10,11 +10,16 @@ import (
 	"github.com/btcsuite/btcd/btcec/v2/ecdsa"
 	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/btcsuite/btcd/connmgr"
+	"github.com/btcsuite/btcd/wire"
 	"github.com/lightningnetwork/lnd/aliasmgr"
 	"github.com/lightningnetwork/lnd/brontide"
+	"github.com/lightningnetwork/lnd/chainntnfs"
 	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/channelnotifier"
 	"github.com/lightningnetwork/lnd/discovery"
 	"github.com/lightningnetwork/lnd/feature"
+	"github.com/lightningnetwork/lnd/fn/v2"
+	"github.com/lightningnetwork/lnd/graph/db/models"
 	"github.com/lightningnetwork/lnd/htlcswitch"
 	"github.com/lightningnetwork/lnd/htlcswitch/hodl"
 	"github.com/lightningnetwork/lnd/keychain"
@@ -24,7 +29,9 @@ import (
 	"github.com/lightningnetwork/lnd/lntest/mock"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
+	"github.com/lightningnetwork/lnd/lnwallet/chancloser"
 	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/msgmux"
 	"github.com/lightningnetwork/lnd/netann"
 	"github.com/lightningnetwork/lnd/peer"
 	"github.com/lightningnetwork/lnd/pool"
@@ -175,11 +182,47 @@ func ConnectPeer(conn *brontide.Conn, connReq *connmgr.ConnReq,
 		PubKey:     identityECDH.PubKey(),
 	})
 
+	chanStatusMgr, err := netann.NewChanStatusManager(&netann.
+		ChanStatusConfig{
+		ChanStatusSampleInterval: 30 * time.Second,
+		ChanDisableTimeout:       2 * time.Minute,
+		DB:                       channelDB.ChannelStateDB(),
+		IsChannelActive: func(lnwire.ChannelID) bool {
+			return true
+		},
+		ApplyChannelUpdate: func(*lnwire.ChannelUpdate1,
+			*wire.OutPoint, bool) error {
+
+			return nil
+		},
+	})
+
+	channelNotifier := channelnotifier.New(channelDB.ChannelStateDB())
+	interceptableSwitchNotifier := &mock.ChainNotifier{
+		EpochChan: make(chan *chainntnfs.BlockEpoch, 1),
+	}
+	interceptableSwitchNotifier.EpochChan <- &chainntnfs.BlockEpoch{
+		Height: 1,
+	}
+	interceptableSwitch, err := htlcswitch.NewInterceptableSwitch(
+		&htlcswitch.InterceptableSwitchConfig{
+			CltvRejectDelta:    13,
+			CltvInterceptDelta: 16,
+			Notifier:           interceptableSwitchNotifier,
+		},
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to create interceptable "+
+			"switch: %w", err)
+	}
+
 	pCfg := peer.Config{
-		Conn:                    conn,
-		ConnReq:                 connReq,
+		Conn:    conn,
+		ConnReq: connReq,
+		PubKeyBytes: [33]byte(
+			identityECDH.PubKey().SerializeCompressed(),
+		),
 		Addr:                    peerAddr,
-		Inbound:                 false,
 		Features:                initFeatures,
 		LegacyFeatures:          legacyFeatures,
 		OutgoingCltvRejectDelta: lncfg.DefaultOutgoingCltvRejectDelta,
@@ -187,9 +230,30 @@ func ConnectPeer(conn *brontide.Conn, connReq *connmgr.ConnReq,
 		ErrorBuffer:             errBuffer,
 		WritePool:               writePool,
 		ReadPool:                readPool,
+		Switch:                  &mockMessageSwitch{},
+		InterceptSwitch:         interceptableSwitch,
 		ChannelDB:               channelDB.ChannelStateDB(),
+		ChainArb:                nil,
 		AuthGossiper:            gossiper,
-		ChainNotifier:           &mock.ChainNotifier{},
+		ChanStatusMgr:           chanStatusMgr,
+		ChainIO:                 &mock.ChainIO{},
+		FeeEstimator:            nil,
+		Signer:                  nil,
+		SigPool:                 nil,
+		Wallet: &lnwallet.LightningWallet{
+			WalletController: &mock.WalletController{},
+		},
+		ChainNotifier: &mock.ChainNotifier{},
+		BestBlockView: chainntnfs.NewBestBlockTracker(
+			&mock.ChainNotifier{},
+		),
+		RoutingPolicy:   models.ForwardingPolicy{},
+		Sphinx:          nil,
+		WitnessBeacon:   nil,
+		Invoices:        nil,
+		ChannelNotifier: channelNotifier,
+		HtlcNotifier:    nil,
+		TowerClient:     nil,
 		DisconnectPeer: func(key *btcec.PublicKey) error {
 			fmt.Printf("Peer %x disconnected\n",
 				key.SerializeCompressed())
@@ -201,23 +265,20 @@ func ConnectPeer(conn *brontide.Conn, connReq *connmgr.ConnReq,
 			return lnwire.NodeAnnouncement{},
 				errors.New("unimplemented")
 		},
-
-		PongBuf: pongBuf,
-
 		PrunePersistentPeerConnection: func(_ [33]byte) {},
-
 		FetchLastChanUpdate: func(_ lnwire.ShortChannelID) (
 			*lnwire.ChannelUpdate1, error) {
 
 			return nil, errors.New("unimplemented")
 		},
-
+		FundingManager:          nil,
 		Hodl:                    &hodl.Config{},
 		UnsafeReplay:            false,
 		MaxOutgoingCltvExpiry:   htlcswitch.DefaultMaxOutgoingCltvExpiry,
 		MaxChannelFeeAllocation: htlcswitch.DefaultMaxLinkFeeAllocation,
-		CoopCloseTargetConfs:    defaultCoopCloseTargetConfs,
 		MaxAnchorsCommitFeeRate: commitFee.FeePerKWeight(),
+		CoopCloseTargetConfs:    defaultCoopCloseTargetConfs,
+		ServerPubKey:            [33]byte{},
 		ChannelCommitInterval:   defaultChannelCommitInterval,
 		PendingCommitInterval:   defaultPendingCommitInterval,
 		ChannelCommitBatchSize:  defaultChannelCommitBatchSize,
@@ -241,7 +302,19 @@ func ConnectPeer(conn *brontide.Conn, connReq *connmgr.ConnReq,
 
 			return nil
 		},
-		Quit: make(chan struct{}),
+		AuxLeafStore:              fn.None[lnwallet.AuxLeafStore](),
+		AuxSigner:                 fn.None[lnwallet.AuxSigner](),
+		AuxResolver:               fn.None[lnwallet.AuxContractResolver](),
+		AuxTrafficShaper:          fn.None[htlcswitch.AuxTrafficShaper](),
+		PongBuf:                   pongBuf,
+		DisallowRouteBlinding:     false,
+		DisallowQuiescence:        false,
+		MaxFeeExposure:            0,
+		MsgRouter:                 fn.None[msgmux.Router](),
+		AuxChanCloser:             fn.None[chancloser.AuxChanCloser](),
+		ShouldFwdExpEndorsement:   nil,
+		NoDisconnectOnPongFailure: false,
+		Quit:                      make(chan struct{}),
 	}
 
 	copy(pCfg.PubKeyBytes[:], peerAddr.IdentityKey.SerializeCompressed())

lnd/mock.go

@@ -0,0 +1,230 @@
+package lnd
+
+import (
+	"net"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/lightningnetwork/lnd/fn/v2"
+	"github.com/lightningnetwork/lnd/htlcswitch"
+	"github.com/lightningnetwork/lnd/lntypes"
+	"github.com/lightningnetwork/lnd/lnwallet"
+	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/peer"
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	timeout = time.Second * 5
+)
+
+// mockMessageSwitch is a mock implementation of the messageSwitch interface
+// used for testing without relying on a *htlcswitch.Switch in unit tests.
+type mockMessageSwitch struct {
+	links []htlcswitch.ChannelUpdateHandler
+}
+
+// BestHeight currently returns a dummy value.
+func (m *mockMessageSwitch) BestHeight() uint32 {
+	return 0
+}
+
+// CircuitModifier currently returns a dummy value.
+func (m *mockMessageSwitch) CircuitModifier() htlcswitch.CircuitModifier {
+	return nil
+}
+
+// RemoveLink currently does nothing.
+func (m *mockMessageSwitch) RemoveLink(cid lnwire.ChannelID) {}
+
+// CreateAndAddLink currently returns a dummy value.
+func (m *mockMessageSwitch) CreateAndAddLink(cfg htlcswitch.ChannelLinkConfig,
+	lnChan *lnwallet.LightningChannel) error {
+
+	return nil
+}
+
+// GetLinksByInterface returns the active links.
+func (m *mockMessageSwitch) GetLinksByInterface(pub [33]byte) (
+	[]htlcswitch.ChannelUpdateHandler, error) {
+
+	return m.links, nil
+}
+
+// mockUpdateHandler is a mock implementation of the ChannelUpdateHandler
+// interface. It is used in mockMessageSwitch's GetLinksByInterface method.
+type mockUpdateHandler struct {
+	cid                  lnwire.ChannelID
+	isOutgoingAddBlocked atomic.Bool
+	isIncomingAddBlocked atomic.Bool
+}
+
+// newMockUpdateHandler creates a new mockUpdateHandler.
+func newMockUpdateHandler(cid lnwire.ChannelID) *mockUpdateHandler {
+	return &mockUpdateHandler{
+		cid: cid,
+	}
+}
+
+// HandleChannelUpdate currently does nothing.
+func (m *mockUpdateHandler) HandleChannelUpdate(msg lnwire.Message) {}
+
+// ChanID returns the mockUpdateHandler's cid.
+func (m *mockUpdateHandler) ChanID() lnwire.ChannelID { return m.cid }
+
+// Bandwidth currently returns a dummy value.
+func (m *mockUpdateHandler) Bandwidth() lnwire.MilliSatoshi { return 0 }
+
+// EligibleToForward currently returns a dummy value.
+func (m *mockUpdateHandler) EligibleToForward() bool { return false }
+
+// MayAddOutgoingHtlc currently returns nil.
+func (m *mockUpdateHandler) MayAddOutgoingHtlc(lnwire.MilliSatoshi) error { return nil }
+
+type mockMessageConn struct {
+	t *testing.T
+
+	// MessageConn embeds our interface so that the mock does not need to
+	// implement every function. The mock will panic if an unspecified function
+	// is called.
+	peer.MessageConn
+
+	// writtenMessages is a channel that our mock pushes written messages into.
+	writtenMessages chan []byte
+
+	readMessages   chan []byte
+	curReadMessage []byte
+
+	// writeRaceDetectingCounter is incremented on any function call
+	// associated with writing to the connection. The race detector will
+	// trigger on this counter if a data race exists.
+	writeRaceDetectingCounter int
+
+	// readRaceDetectingCounter is incremented on any function call
+	// associated with reading from the connection. The race detector will
+	// trigger on this counter if a data race exists.
+	readRaceDetectingCounter int
+}
+
+func (m *mockUpdateHandler) EnableAdds(dir htlcswitch.LinkDirection) bool {
+	if dir == htlcswitch.Outgoing {
+		return m.isOutgoingAddBlocked.Swap(false)
+	}
+
+	return m.isIncomingAddBlocked.Swap(false)
+}
+
+func (m *mockUpdateHandler) DisableAdds(dir htlcswitch.LinkDirection) bool {
+	if dir == htlcswitch.Outgoing {
+		return !m.isOutgoingAddBlocked.Swap(true)
+	}
+
+	return !m.isIncomingAddBlocked.Swap(true)
+}
+
+func (m *mockUpdateHandler) IsFlushing(dir htlcswitch.LinkDirection) bool {
+	switch dir {
+	case htlcswitch.Outgoing:
+		return m.isOutgoingAddBlocked.Load()
+	case htlcswitch.Incoming:
+		return m.isIncomingAddBlocked.Load()
+	}
+
+	return false
+}
+
+func (m *mockUpdateHandler) OnFlushedOnce(hook func()) {
+	hook()
+}
+func (m *mockUpdateHandler) OnCommitOnce(
+	_ htlcswitch.LinkDirection, hook func(),
+) {
+
+	hook()
+}
+func (m *mockUpdateHandler) InitStfu() <-chan fn.Result[lntypes.ChannelParty] {
+	// TODO(proofofkeags): Implement
+	c := make(chan fn.Result[lntypes.ChannelParty], 1)
+
+	c <- fn.Errf[lntypes.ChannelParty]("InitStfu not yet implemented")
+
+	return c
+}
+
+func newMockConn(t *testing.T, expectedMessages int) *mockMessageConn {
+	return &mockMessageConn{
+		t:               t,
+		writtenMessages: make(chan []byte, expectedMessages),
+		readMessages:    make(chan []byte, 1),
+	}
+}
+
+// SetWriteDeadline mocks setting write deadline for our conn.
+func (m *mockMessageConn) SetWriteDeadline(time.Time) error {
+	m.writeRaceDetectingCounter++
+	return nil
+}
+
+// Flush mocks a message conn flush.
+func (m *mockMessageConn) Flush() (int, error) {
+	m.writeRaceDetectingCounter++
+	return 0, nil
+}
+
+// WriteMessage mocks sending of a message on our connection. It will push
+// the bytes sent into the mock's writtenMessages channel.
+func (m *mockMessageConn) WriteMessage(msg []byte) error {
+	m.writeRaceDetectingCounter++
+
+	msgCopy := make([]byte, len(msg))
+	copy(msgCopy, msg)
+
+	select {
+	case m.writtenMessages <- msgCopy:
+	case <-time.After(timeout):
+		m.t.Fatalf("timeout sending message: %v", msgCopy)
+	}
+
+	return nil
+}
+
+// assertWrite asserts that our mock as had WriteMessage called with the byte
+// slice we expect.
+func (m *mockMessageConn) assertWrite(expected []byte) {
+	select {
+	case actual := <-m.writtenMessages:
+		require.Equal(m.t, expected, actual)
+
+	case <-time.After(timeout):
+		m.t.Fatalf("timeout waiting for write: %v", expected)
+	}
+}
+
+func (m *mockMessageConn) SetReadDeadline(t time.Time) error {
+	m.readRaceDetectingCounter++
+	return nil
+}
+
+func (m *mockMessageConn) ReadNextHeader() (uint32, error) {
+	m.readRaceDetectingCounter++
+	m.curReadMessage = <-m.readMessages
+	return uint32(len(m.curReadMessage)), nil
+}
+
+func (m *mockMessageConn) ReadNextBody(buf []byte) ([]byte, error) {
+	m.readRaceDetectingCounter++
+	return m.curReadMessage, nil
+}
+
+func (m *mockMessageConn) RemoteAddr() net.Addr {
+	return nil
+}
+
+func (m *mockMessageConn) LocalAddr() net.Addr {
+	return nil
+}
+
+func (m *mockMessageConn) Close() error {
+	return nil
+}