diff --git a/peer/brontide_test.go b/peer/brontide_test.go index 3d8023b1a..10091b67a 100644 --- a/peer/brontide_test.go +++ b/peer/brontide_test.go @@ -6,6 +6,7 @@ import ( "testing" "time" + "github.com/btcsuite/btcd/btcec/v2" "github.com/btcsuite/btcd/btcutil" "github.com/btcsuite/btcd/chaincfg" "github.com/btcsuite/btcd/txscript" @@ -19,6 +20,9 @@ import ( "github.com/lightningnetwork/lnd/lnwallet" "github.com/lightningnetwork/lnd/lnwallet/chancloser" "github.com/lightningnetwork/lnd/lnwire" + "github.com/lightningnetwork/lnd/routing/route" + "github.com/lightningnetwork/lnd/tlv" + "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" ) @@ -1466,3 +1470,210 @@ func TestRemovePendingChannel(t *testing.T) { require.NoError(t, err) } + +// mockAuxTrafficShaper is a mock implementation of htlcswitch.AuxTrafficShaper +// for testing the createHtlcValidator function. +type mockAuxTrafficShaper struct { + mock.Mock +} + +// ShouldHandleTraffic returns the configured mock values. +func (m *mockAuxTrafficShaper) ShouldHandleTraffic( + cid lnwire.ShortChannelID, + fundingBlob, htlcBlob fn.Option[tlv.Blob]) (bool, error) { + + args := m.Called(cid, fundingBlob, htlcBlob) + return args.Bool(0), args.Error(1) +} + +// PaymentBandwidth returns the configured mock values. +func (m *mockAuxTrafficShaper) PaymentBandwidth(fundingBlob, htlcBlob, + commitmentBlob fn.Option[tlv.Blob], linkBandwidth, + htlcAmt lnwire.MilliSatoshi, htlcView lnwallet.AuxHtlcView, + peer route.Vertex) (lnwire.MilliSatoshi, error) { + + args := m.Called( + fundingBlob, htlcBlob, commitmentBlob, linkBandwidth, + htlcAmt, htlcView, peer, + ) + + bw, _ := args.Get(0).(lnwire.MilliSatoshi) + + return bw, args.Error(1) +} + +// ProduceHtlcExtraData is part of the AuxTrafficShaper interface. +func (m *mockAuxTrafficShaper) ProduceHtlcExtraData( + totalAmount lnwire.MilliSatoshi, + htlcCustomRecords lnwire.CustomRecords, + peer route.Vertex) (lnwire.MilliSatoshi, lnwire.CustomRecords, + error) { + + args := m.Called(totalAmount, htlcCustomRecords, peer) + + amt, _ := args.Get(0).(lnwire.MilliSatoshi) + records, _ := args.Get(1).(lnwire.CustomRecords) + + return amt, records, args.Error(2) +} + +// IsCustomHTLC is part of the AuxTrafficShaper interface. +func (m *mockAuxTrafficShaper) IsCustomHTLC( + htlcRecords lnwire.CustomRecords) bool { + + args := m.Called(htlcRecords) + return args.Bool(0) +} + +// Compile-time check that mockAuxTrafficShaper implements AuxTrafficShaper. +var _ htlcswitch.AuxTrafficShaper = (*mockAuxTrafficShaper)(nil) + +// TestCreateHtlcValidator tests that the HTLC validator created by +// createHtlcValidator respects the ShouldHandleTraffic check. When +// ShouldHandleTraffic returns false, the validator should return nil without +// calling PaymentBandwidth. +func TestCreateHtlcValidator(t *testing.T) { + t.Parallel() + + // Create a minimal Brontide with just the identity key set. + privKey, err := btcec.NewPrivateKey() + require.NoError(t, err) + + peer := &Brontide{ + cfg: Config{ + Addr: &lnwire.NetAddress{ + IdentityKey: privKey.PubKey(), + }, + }, + } + + // Create a mock channel with minimal required fields. + dbChan := &channeldb.OpenChannel{ + ShortChannelID: lnwire.NewShortChanIDFromInt(123), + } + + anyArg := mock.Anything + + testCases := []struct { + name string + setupMock func(*mockAuxTrafficShaper) + htlcAmount lnwire.MilliSatoshi + linkBw lnwire.MilliSatoshi + expectError bool + }{ + { + name: "non-custom channel skips check", + setupMock: func(m *mockAuxTrafficShaper) { + m.On( + "ShouldHandleTraffic", + anyArg, anyArg, anyArg, + ).Return(false, nil) + }, + htlcAmount: 1000, + linkBw: 5000, + expectError: false, + }, + { + name: "sufficient bandwidth", + setupMock: func(m *mockAuxTrafficShaper) { + m.On( + "ShouldHandleTraffic", + anyArg, anyArg, anyArg, + ).Return(true, nil) + m.On( + "PaymentBandwidth", + anyArg, anyArg, anyArg, + anyArg, anyArg, anyArg, + anyArg, + ).Return( + lnwire.MilliSatoshi(10000), + nil, + ) + }, + htlcAmount: 1000, + linkBw: 5000, + expectError: false, + }, + { + name: "insufficient bandwidth", + setupMock: func(m *mockAuxTrafficShaper) { + m.On( + "ShouldHandleTraffic", + anyArg, anyArg, anyArg, + ).Return(true, nil) + m.On( + "PaymentBandwidth", + anyArg, anyArg, anyArg, + anyArg, anyArg, anyArg, + anyArg, + ).Return( + lnwire.MilliSatoshi(500), + nil, + ) + }, + htlcAmount: 1000, + linkBw: 5000, + expectError: true, + }, + { + name: "ShouldHandleTraffic error", + setupMock: func(m *mockAuxTrafficShaper) { + m.On( + "ShouldHandleTraffic", + anyArg, anyArg, anyArg, + ).Return( + false, + fmt.Errorf("shaper error"), + ) + }, + htlcAmount: 1000, + linkBw: 5000, + expectError: true, + }, + { + name: "PaymentBandwidth error", + setupMock: func(m *mockAuxTrafficShaper) { + m.On( + "ShouldHandleTraffic", + anyArg, anyArg, anyArg, + ).Return(true, nil) + m.On( + "PaymentBandwidth", + anyArg, anyArg, anyArg, + anyArg, anyArg, anyArg, + anyArg, + ).Return( + lnwire.MilliSatoshi(0), + fmt.Errorf("bandwidth error"), + ) + }, + htlcAmount: 1000, + linkBw: 5000, + expectError: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + m := &mockAuxTrafficShaper{} + tc.setupMock(m) + + validator := peer.createHtlcValidator( + dbChan, m, + ) + + err := validator.ValidateHtlc( + tc.htlcAmount, tc.linkBw, + nil, lnwallet.AuxHtlcView{}, + ) + + if tc.expectError { + require.Error(t, err) + } else { + require.NoError(t, err) + } + + m.AssertExpectations(t) + }) + } +}