Merge pull request #9980 from GeorgeTsagk/enhance-aux-modifier

AuxTrafficShaper methods use first hop pub key

Author: guggero

docs/release-notes/release-notes-0.19.2.md

@@ -57,7 +57,10 @@
 - [Improved](https://github.com/lightningnetwork/lnd/pull/9880) the connection
   restriction logic enforced by `accessman`. In addition, the restriction placed
   on outbound connections is now lifted.
-
+- [Enhanced](https://github.com/lightningnetwork/lnd/pull/9980) the aux traffic
+  shaper to now accept the first hop peer pub key as an argument. This can
+  affect the reported aux bandwidth and also the custom records that are
+  produced.
 ## RPC Updates
 
 ## lncli Updates

htlcswitch/interfaces.go

@@ -14,6 +14,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/record"
+	"github.com/lightningnetwork/lnd/routing/route"
 	"github.com/lightningnetwork/lnd/tlv"
 )
 
@@ -495,8 +496,9 @@ type AuxHtlcModifier interface {
 	// data blob of an HTLC, may produce a different blob or modify the
 	// amount of bitcoin this htlc should carry.
 	ProduceHtlcExtraData(totalAmount lnwire.MilliSatoshi,
-		htlcCustomRecords lnwire.CustomRecords) (lnwire.MilliSatoshi,
-		lnwire.CustomRecords, error)
+		htlcCustomRecords lnwire.CustomRecords,
+		peer route.Vertex) (lnwire.MilliSatoshi, lnwire.CustomRecords,
+		error)
 }
 
 // AuxTrafficShaper is an interface that allows the sender to determine if a
@@ -520,7 +522,8 @@ type AuxTrafficShaper interface {
 	PaymentBandwidth(fundingBlob, htlcBlob,
 		commitmentBlob fn.Option[tlv.Blob],
 		linkBandwidth, htlcAmt lnwire.MilliSatoshi,
-		htlcView lnwallet.AuxHtlcView) (lnwire.MilliSatoshi, error)
+		htlcView lnwallet.AuxHtlcView,
+		peer route.Vertex) (lnwire.MilliSatoshi, error)
 
 	// IsCustomHTLC returns true if the HTLC carries the set of relevant
 	// custom records to put it under the purview of the traffic shaper,

htlcswitch/link.go

@@ -31,6 +31,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/queue"
 	"github.com/lightningnetwork/lnd/record"
+	"github.com/lightningnetwork/lnd/routing/route"
 	"github.com/lightningnetwork/lnd/ticker"
 	"github.com/lightningnetwork/lnd/tlv"
 )
@@ -3503,11 +3504,19 @@ func (l *channelLink) AuxBandwidth(amount lnwire.MilliSatoshi,
 		})
 	}
 
+	peerBytes := l.cfg.Peer.PubKey()
+
+	peer, err := route.NewVertexFromBytes(peerBytes[:])
+	if err != nil {
+		return fn.Err[OptionalBandwidth](fmt.Errorf("failed to decode "+
+			"peer pub key: %v", err))
+	}
+
 	// Ask for a specific bandwidth to be used for the channel.
 	commitmentBlob := l.CommitmentCustomBlob()
 	auxBandwidth, err := ts.PaymentBandwidth(
 		fundingBlob, htlcBlob, commitmentBlob, l.Bandwidth(), amount,
-		l.channel.FetchLatestAuxHTLCView(),
+		l.channel.FetchLatestAuxHTLCView(), peer,
 	)
 	if err != nil {
 		return fn.Err[OptionalBandwidth](fmt.Errorf("failed to get "+

routing/bandwidth_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/lightningnetwork/lnd/htlcswitch"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/routing/route"
 	"github.com/lightningnetwork/lnd/tlv"
 	"github.com/stretchr/testify/require"
 )
@@ -152,7 +153,7 @@ func (*mockTrafficShaper) ShouldHandleTraffic(_ lnwire.ShortChannelID,
 // ShouldHandleTraffic method should be called first.
 func (*mockTrafficShaper) PaymentBandwidth(_, _, _ fn.Option[tlv.Blob],
 	linkBandwidth, _ lnwire.MilliSatoshi,
-	_ lnwallet.AuxHtlcView) (lnwire.MilliSatoshi, error) {
+	_ lnwallet.AuxHtlcView, _ route.Vertex) (lnwire.MilliSatoshi, error) {
 
 	return linkBandwidth, nil
 }
@@ -161,8 +162,8 @@ func (*mockTrafficShaper) PaymentBandwidth(_, _, _ fn.Option[tlv.Blob],
 // data blob of an HTLC, may produce a different blob or modify the
 // amount of bitcoin this htlc should carry.
 func (*mockTrafficShaper) ProduceHtlcExtraData(totalAmount lnwire.MilliSatoshi,
-	_ lnwire.CustomRecords) (lnwire.MilliSatoshi, lnwire.CustomRecords,
-	error) {
+	_ lnwire.CustomRecords, _ route.Vertex) (lnwire.MilliSatoshi,
+	lnwire.CustomRecords, error) {
 
 	return totalAmount, nil, nil
 }

routing/payment_lifecycle.go

@@ -724,6 +724,13 @@ func (p *paymentLifecycle) amendFirstHopData(rt *route.Route) error {
 	// value.
 	rt.FirstHopWireCustomRecords = p.firstHopCustomRecords
 
+	if len(rt.Hops) == 0 {
+		return fmt.Errorf("cannot amend first hop data, route length " +
+			"is zero")
+	}
+
+	firstHopPK := rt.Hops[0].PubKeyBytes
+
 	// extraDataRequest is a helper struct to pass the custom records and
 	// amount back from the traffic shaper.
 	type extraDataRequest struct {
@@ -740,6 +747,7 @@ func (p *paymentLifecycle) amendFirstHopData(rt *route.Route) error {
 		func(ts htlcswitch.AuxTrafficShaper) fn.Result[extraDataRequest] {
 			newAmt, newRecords, err := ts.ProduceHtlcExtraData(
 				rt.TotalAmount, p.firstHopCustomRecords,
+				firstHopPK,
 			)
 			if err != nil {
 				return fn.Err[extraDataRequest](err)

routing/payment_lifecycle_test.go

@@ -368,7 +368,11 @@ func TestRequestRouteSucceed(t *testing.T) {
 
 	// Create a mock payment session and a dummy route.
 	paySession := &mockPaymentSession{}
-	dummyRoute := &route.Route{}
+	dummyRoute := &route.Route{
+		Hops: []*route.Hop{
+			testHop,
+		},
+	}
 
 	// Mount the mocked payment session.
 	p.paySession = paySession