diff --git a/lnwallet/aux_test_utils.go b/lnwallet/aux_test_utils.go new file mode 100644 index 000000000..1e76a522e --- /dev/null +++ b/lnwallet/aux_test_utils.go @@ -0,0 +1,31 @@ +package lnwallet + +import ( + "github.com/lightningnetwork/lnd/lnwire" +) + +// NewTestAuxHtlcDescriptor creates an AuxHtlcDescriptor for testing purposes. +// This function allows tests to create descriptors with specific commit heights +// and entry types, which are normally unexported fields. +func NewTestAuxHtlcDescriptor(chanID lnwire.ChannelID, + rHash PaymentHash, timeout uint32, + amount lnwire.MilliSatoshi, htlcIndex, parentIndex uint64, + entryType uint8, customRecords lnwire.CustomRecords, + addHeightLocal, addHeightRemote, removeHeightLocal, + removeHeightRemote uint64) AuxHtlcDescriptor { + + return AuxHtlcDescriptor{ + ChanID: chanID, + RHash: rHash, + Timeout: timeout, + Amount: amount, + HtlcIndex: htlcIndex, + ParentIndex: parentIndex, + EntryType: updateType(entryType), + CustomRecords: customRecords, + addCommitHeightLocal: addHeightLocal, + addCommitHeightRemote: addHeightRemote, + removeCommitHeightLocal: removeHeightLocal, + removeCommitHeightRemote: removeHeightRemote, + } +} diff --git a/lnwallet/channel.go b/lnwallet/channel.go index fb1e3ac4b..02f5f9ccf 100644 --- a/lnwallet/channel.go +++ b/lnwallet/channel.go @@ -833,6 +833,20 @@ type LightningChannel struct { // is created. type ChannelOpt func(*channelOpts) +// AuxHtlcValidator is an interface for validating whether an HTLC can be added +// to a custom channel. It is called during HTLC validation with the current +// channel state and HTLC details. This allows external components (like the +// traffic shaper) to perform final validation checks against the most +// up-to-date channel state before the HTLC is committed. +type AuxHtlcValidator interface { + // ValidateHtlc checks whether the given HTLC can be added to the + // channel given the current link bandwidth, custom records, and HTLC + // view. + ValidateHtlc(amount, linkBandwidth lnwire.MilliSatoshi, + customRecords lnwire.CustomRecords, + view AuxHtlcView) error +} + // channelOpts is the set of options used to create a new channel. type channelOpts struct { localNonce *musig2.Nonces @@ -842,6 +856,10 @@ type channelOpts struct { auxSigner fn.Option[AuxSigner] auxResolver fn.Option[AuxContractResolver] + // auxHtlcValidator is an optional validator that performs custom + // validation on HTLCs before they are added to the channel state. + auxHtlcValidator fn.Option[AuxHtlcValidator] + skipNonceInit bool } @@ -894,6 +912,15 @@ func WithAuxResolver(resolver AuxContractResolver) ChannelOpt { } } +// WithAuxHtlcValidator is used to specify a custom HTLC validator for the +// channel. This validator will be called during HTLC addition to perform +// final validation checks against the most up-to-date channel state. +func WithAuxHtlcValidator(validator AuxHtlcValidator) ChannelOpt { + return func(o *channelOpts) { + o.auxHtlcValidator = fn.Some(validator) + } +} + // defaultChannelOpts returns the set of default options for a new channel. func defaultChannelOpts() *channelOpts { return &channelOpts{} @@ -2738,9 +2765,21 @@ func (lc *LightningChannel) FetchLatestAuxHTLCView() AuxHtlcView { lc.RLock() defer lc.RUnlock() - return newAuxHtlcView(lc.fetchHTLCView( - lc.updateLogs.Remote.logIndex, lc.updateLogs.Local.logIndex, - )) + nextHeight := lc.commitChains.Local.tip().height + 1 + + // We use the remote ACKed index from the last signed local commitment + // (tail) rather than the remote's latest log index. This ensures we + // only include remote HTLCs that have been locked into a signed + // commitment, giving the aux validator a stable, consistent view that + // matches the actual commitment state used for balance calculations. + remoteACKedIndex := lc.commitChains.Local.tail().messageIndices.Remote + view := lc.fetchHTLCView( + remoteACKedIndex, lc.updateLogs.Local.logIndex, + ) + + view.NextHeight = nextHeight + + return newAuxHtlcView(view) } // fetchHTLCView returns all the candidate HTLC updates which should be @@ -6065,6 +6104,52 @@ func (lc *LightningChannel) addHTLC(htlc *lnwire.UpdateAddHTLC, return 0, err } + // If an auxiliary HTLC validator is configured, call it now to perform + // custom validation checks against the current channel state. This is + // the final validation point before the HTLC is added to the update + // log, ensuring that the validator sees the most up-to-date state + // including all previously validated HTLCs in this batch. + // + // NOTE: This is called after the standard commitment sanity checks to + // ensure we only perform (potentially) expensive custom validation on + // HTLCs that have already passed the basic Lightning protocol + // constraints. + err := fn.MapOptionZ( + lc.opts.auxHtlcValidator, + func(validator AuxHtlcValidator) error { + // Fetch the current HTLC view which includes all + // pending HTLCs that haven't been committed yet. This + // provides the validator with the most accurate state. + commitChain := lc.commitChains.Local + remoteIndex := commitChain.tail().messageIndices.Remote + view := lc.fetchHTLCView( + remoteIndex, lc.updateLogs.Local.logIndex, + ) + + nextHeight := lc.commitChains.Local.tip().height + 1 + view.NextHeight = nextHeight + + lc.log.Debugf("Setting view nextheight=%v", nextHeight) + + auxView := newAuxHtlcView(view) + + // Get the current available balance for the link + // bandwidth check. This is needed for the balance + // validation in the traffic shaper. We use NoBuffer + // since the buffer check was already performed earlier, + // and assets don't pay on-chain fees. + linkBandwidth, _ := lc.availableBalance(NoBuffer) + + return validator.ValidateHtlc( + pd.Amount, linkBandwidth, pd.CustomRecords, + auxView, + ) + }, + ) + if err != nil { + return 0, fmt.Errorf("aux HTLC validation failed: %w", err) + } + lc.updateLogs.Local.appendHtlc(pd) return pd.HtlcIndex, nil @@ -6216,6 +6301,7 @@ func (lc *LightningChannel) htlcAddDescriptor(htlc *lnwire.UpdateAddHTLC, // remote commitments. func (lc *LightningChannel) validateAddHtlc(pd *paymentDescriptor, buffer BufferType) error { + // Make sure adding this HTLC won't violate any of the constraints we // must keep on the commitment transactions. remoteACKedIndex := lc.commitChains.Local.tail().messageIndices.Remote diff --git a/peer/brontide.go b/peer/brontide.go index 0f4d6e349..1aebdfbb5 100644 --- a/peer/brontide.go +++ b/peer/brontide.go @@ -52,6 +52,7 @@ import ( "github.com/lightningnetwork/lnd/pool" "github.com/lightningnetwork/lnd/protofsm" "github.com/lightningnetwork/lnd/queue" + "github.com/lightningnetwork/lnd/routing/route" "github.com/lightningnetwork/lnd/subscribe" "github.com/lightningnetwork/lnd/ticker" "github.com/lightningnetwork/lnd/tlv" @@ -1146,6 +1147,16 @@ func (p *Brontide) loadActiveChannels(chans []*channeldb.OpenChannel) ( }, ) + p.cfg.AuxTrafficShaper.WhenSome( + func(ts htlcswitch.AuxTrafficShaper) { + val := p.createHtlcValidator(dbChan, ts) + chanOpts = append( + chanOpts, + lnwallet.WithAuxHtlcValidator(val), + ) + }, + ) + lnChan, err := lnwallet.NewLightningChannel( p.cfg.Signer, dbChan, p.cfg.SigPool, chanOpts..., ) @@ -5259,6 +5270,15 @@ func (p *Brontide) addActiveChannel(c *lnpeer.NewChannel) error { chanOpts = append(chanOpts, lnwallet.WithAuxResolver(s)) }) + p.cfg.AuxTrafficShaper.WhenSome( + func(ts htlcswitch.AuxTrafficShaper) { + val := p.createHtlcValidator(c.OpenChannel, ts) + chanOpts = append( + chanOpts, lnwallet.WithAuxHtlcValidator(val), + ) + }, + ) + // If not already active, we'll add this channel to the set of active // channels, so we can look it up later easily according to its channel // ID. @@ -5465,6 +5485,96 @@ func (p *Brontide) scaleTimeout(timeout time.Duration) time.Duration { return timeout } +// auxHtlcValidator implements lnwallet.AuxHtlcValidator by checking HTLC +// bandwidth against the traffic shaper. +type auxHtlcValidator struct { + peer *Brontide + dbChan *channeldb.OpenChannel + ts htlcswitch.AuxTrafficShaper +} + +// ValidateHtlc performs final aux balance validation before an HTLC is added +// to the channel state. It calls into the traffic shaper's PaymentBandwidth +// method to check external balance against the most up-to-date channel state, +// preventing race conditions where multiple HTLCs could be approved based on +// stale bandwidth. +func (v *auxHtlcValidator) ValidateHtlc(amount, + linkBandwidth lnwire.MilliSatoshi, + customRecords lnwire.CustomRecords, + view lnwallet.AuxHtlcView) error { + + // Get the short channel ID for logging. + scid := v.dbChan.ShortChannelID + + // Extract the HTLC custom records to pass to the traffic shaper. + var htlcBlob fn.Option[tlv.Blob] + if len(customRecords) > 0 { + blob, err := customRecords.Serialize() + if err != nil { + return fmt.Errorf("unable to serialize "+ + "custom records: %w", err) + } + htlcBlob = fn.Some(blob) + } + + // Get the funding and commitment blobs for this channel. + fundingBlob := v.dbChan.CustomBlob + commitmentBlob := v.dbChan.LocalCommitment.CustomBlob + + // Check if this channel should be handled by the traffic shaper. If + // not, we skip the aux validation entirely and allow the HTLC to + // proceed through normal validation. + shouldHandle, err := v.ts.ShouldHandleTraffic( + scid, fundingBlob, htlcBlob, + ) + if err != nil { + return fmt.Errorf("traffic shaper failed to decide "+ + "whether to handle traffic: %w", err) + } + if !shouldHandle { + return nil + } + + peer := route.NewVertex(v.peer.IdentityKey()) + + // Call the traffic shaper's PaymentBandwidth method with the current + // state. This performs the same bandwidth checks as during + // pathfinding/forwarding, but against the absolute latest channel + // state. + // + // The linkBandwidth is provided by the channel and represents the + // current available balance, which is used by the traffic shaper to + // ensure we don't dip below channel reserves. + bandwidth, err := v.ts.PaymentBandwidth( + fundingBlob, htlcBlob, commitmentBlob, + linkBandwidth, amount, view, peer, + ) + if err != nil { + return fmt.Errorf("traffic shaper bandwidth check "+ + "failed: %w", err) + } + + if amount > bandwidth { + return fmt.Errorf("insufficient aux bandwidth: "+ + "need %v, have %v (scid=%v)", amount, + bandwidth, scid) + } + + return nil +} + +// createHtlcValidator creates an HTLC validator that performs final aux balance +// validation before HTLCs are added to the channel state. +func (p *Brontide) createHtlcValidator(dbChan *channeldb.OpenChannel, + ts htlcswitch.AuxTrafficShaper) lnwallet.AuxHtlcValidator { + + return &auxHtlcValidator{ + peer: p, + dbChan: dbChan, + ts: ts, + } +} + // CoopCloseUpdates is a struct used to communicate updates for an active close // to the caller. type CoopCloseUpdates struct { 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) + }) + } +}