From b1701cf232ec3350cde625fb3261af14f9f90ea9 Mon Sep 17 00:00:00 2001
From: George Tsagkarelis <george.tsagkarelis@gmail.com>
Date: Mon, 8 Dec 2025 12:38:18 +0100
Subject: [PATCH] peer: set and use aux htlc validator

When instantiating the lightning channel we now pass in the created HTLC
validator. This validator simply performs a bandwidth check and errors
out if that is insufficient.
---
 peer/brontide.go | 110 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 110 insertions(+)

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 {