Merge pull request #10434 from GeorgeTsagk/extra-aux-band-check

Add `AuxHtlcValidator`

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,
+	}
+}

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

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 {

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)
+		})
+	}
+}