multi: fixes to tlv traffic shaper

config_builder.go

@@ -152,12 +152,6 @@ type ImplementationCfg struct {
 	// AuxComponents is a set of auxiliary components that can be used by
 	// lnd for certain custom channel types.
 	AuxComponents
-
-	TlvTrafficShaper
-}
-
-type TlvTrafficShaper struct {
-	TrafficShaper fn.Option[routing.TlvTrafficShaper]
 }
 
 // AuxComponents is a set of auxiliary components that can be used by lnd for
@@ -180,6 +174,10 @@ type AuxComponents struct {
 	// AuxSigner is an optional signer that can be used to sign auxiliary
 	// leaves for certain custom channel types.
 	AuxSigner fn.Option[lnwallet.AuxSigner]
+
+	// TrafficShaper is an optional traffic shaper that can be used to
+	// control the outgoing channel of a payment.
+	TrafficShaper fn.Option[routing.TlvTrafficShaper]
 }
 
 // DefaultWalletImpl is the default implementation of our normal, btcwallet

htlcswitch/interfaces.go

@@ -272,9 +272,15 @@ type ChannelLink interface {
 	// have buffered messages.
 	AttachMailBox(MailBox)
 
-	// ChannelCustomBlob returns the custom blob of the channel that this
-	// link is associated with.
-	ChannelCustomBlob() fn.Option[tlv.Blob]
+	// FundingCustomBlob returns the custom funding blob of the channel that
+	// this link is associated with. The funding blob represents static
+	// information about the channel that was created at channel funding
+	// time.
+	FundingCustomBlob() fn.Option[tlv.Blob]
+
+	// CommitmentCustomBlob returns the custom blob of the current local
+	// commitment of the channel that this link is associated with.
+	CommitmentCustomBlob() fn.Option[tlv.Blob]
 
 	// Start/Stop are used to initiate the start/stop of the channel link
 	// functioning.

htlcswitch/link.go

@@ -3754,8 +3754,15 @@ func (l *channelLink) fail(linkErr LinkFailureError,
 	l.cfg.OnChannelFailure(l.ChanID(), l.ShortChanID(), linkErr)
 }
 
-// ChannelCustomBlob returns the custom blob of the channel that this link is
-// associated with.
-func (l *channelLink) ChannelCustomBlob() fn.Option[tlv.Blob] {
+// FundingCustomBlob returns the custom funding blob of the channel that this
+// link is associated with. The funding blob represents static information about
+// the channel that was created at channel funding time.
+func (l *channelLink) FundingCustomBlob() fn.Option[tlv.Blob] {
 	return l.channel.State().CustomBlob
 }
+
+// CommitmentCustomBlob returns the custom blob of the current local commitment
+// of the channel that this link is associated with.
+func (l *channelLink) CommitmentCustomBlob() fn.Option[tlv.Blob] {
+	return l.channel.LocalCommitmentBlob()
+}

htlcswitch/mock.go

@@ -948,6 +948,14 @@ func (f *mockChannelLink) OnCommitOnce(LinkDirection, func()) {
 	// TODO(proofofkeags): Implement
 }
 
+func (f *mockChannelLink) FundingCustomBlob() fn.Option[tlv.Blob] {
+	return fn.None[tlv.Blob]()
+}
+
+func (f *mockChannelLink) CommitmentCustomBlob() fn.Option[tlv.Blob] {
+	return fn.None[tlv.Blob]()
+}
+
 var _ ChannelLink = (*mockChannelLink)(nil)
 
 func newDB() (*channeldb.DB, func(), error) {

lnwallet/channel.go

@@ -9539,3 +9539,19 @@ func (lc *LightningChannel) MultiSigKeys() (keychain.KeyDescriptor,
 	return lc.channelState.LocalChanCfg.MultiSigKey,
 		lc.channelState.RemoteChanCfg.MultiSigKey
 }
+
+// LocalCommitmentBlob returns the custom blob of the local commitment.
+func (lc *LightningChannel) LocalCommitmentBlob() fn.Option[tlv.Blob] {
+	lc.RLock()
+	defer lc.RUnlock()
+
+	chanState := lc.channelState
+	localBalance := chanState.LocalCommitment.CustomBlob
+
+	return fn.MapOption(func(b tlv.Blob) tlv.Blob {
+		newBlob := make([]byte, len(b))
+		copy(newBlob, b)
+
+		return newBlob
+	})(localBalance)
+}

routing/bandwidth.go

@@ -1,6 +1,8 @@
 package routing
 
 import (
+	"fmt"
+
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/lightningnetwork/lnd/channeldb"
 	"github.com/lightningnetwork/lnd/fn"
@@ -29,16 +31,19 @@ type bandwidthHints interface {
 // payment should be carried by a channel based on the TLV records that may be
 // present in the `update_add_htlc` message or the channel commitment itself.
 type TlvTrafficShaper interface {
-	// ShouldCarryPayment returns true if the provided tlv records indicate
-	// that this channel may carry out the payment by utilizing external
-	// mechanisms.
-	ShouldCarryPayment(amt lnwire.MilliSatoshi, htlcTLV,
-		channelBlob fn.Option[tlv.Blob]) bool
-
 	// HandleTraffic is called in order to check if the channel identified
 	// by the provided channel ID may have external mechanisms that would
 	// allow it to carry out the payment.
-	HandleTraffic(cid lnwire.ShortChannelID) bool
+	HandleTraffic(cid lnwire.ShortChannelID,
+		fundingBlob fn.Option[tlv.Blob]) (bool, error)
+
+	// PaymentBandwidth returns the available bandwidth for a custom channel
+	// decided by the given channel aux blob and HTLC blob. A return value
+	// of 0 means there is no bandwidth available. To find out if a channel
+	// is a custom channel that should be handled by the traffic shaper, the
+	// HandleTraffic method should be called first.
+	PaymentBandwidth(htlcBlob,
+		commitmentBlob fn.Option[tlv.Blob]) (lnwire.MilliSatoshi, error)
 
 	AuxHtlcModifier
 }
@@ -47,10 +52,10 @@ type TlvTrafficShaper interface {
 // HTLC of a payment by changing the amount or the wire message tlv records.
 type AuxHtlcModifier interface {
 	// ProduceHtlcExtraData is a function that, based on the previous extra
-	// data blob of an htlc, may produce a different blob or modify the
+	// data blob of an HTLC, may produce a different blob or modify the
 	// amount of bitcoin this htlc should carry.
-	ProduceHtlcExtraData(htlcBlob tlv.Blob,
-		chanID uint64) (btcutil.Amount, tlv.Blob, error)
+	ProduceHtlcExtraData(totalAmount lnwire.MilliSatoshi,
+		htlcBlob tlv.Blob) (btcutil.Amount, tlv.Blob, error)
 }
 
 // getLinkQuery is the function signature used to lookup a link.
@@ -125,51 +130,57 @@ func (b *bandwidthManager) getBandwidth(cid lnwire.ShortChannelID,
 		return 0
 	}
 
-	res := fn.MapOptionZ(b.trafficShaper, func(ts TlvTrafficShaper) bool {
-		return ts.HandleTraffic(cid)
-	})
-
-	// If response is no and we do have a traffic shaper, we can't handle
-	// the traffic and return early.
-	if !res && b.trafficShaper.IsSome() {
-		log.Warnf("ShortChannelID=%v: can't handle traffic", cid)
-		return 0
-	}
-
-	channelBlob := link.ChannelCustomBlob()
-
-	// Run the wrapped traffic if it exists, otherwise return false.
-	res = fn.MapOptionZ(b.trafficShaper, func(ts TlvTrafficShaper) bool {
-		return ts.ShouldCarryPayment(amount, htlcBlob, channelBlob)
-	})
-
-	// If the traffic shaper indicates that this channel can route the
-	// payment, we immediatelly select this channel and return maximum
-	// bandwidth as response.
-	if res {
-		// If the amount is zero, but the traffic shaper signaled that
-		// the channel can carry the payment, we'll return the maximum
-		// amount. A zero amount is used when we try to figure out if
-		// enough balance exists for the payment to be carried out, but
-		// at that point we don't know the payment amount in order to
-		// return an exact value, so we signal a value that will
-		// certainly satisfy the payment amount.
-		if amount == 0 {
-			// We don't want to just return the max uint64 as this
-			// will overflow when further amounts are added
-			// together.
-			return lnwire.MaxMilliSatoshi / 2
+	var (
+		auxBandwidth           lnwire.MilliSatoshi
+		auxBandwidthDetermined bool
+	)
+	err = fn.MapOptionZ(b.trafficShaper, func(ts TlvTrafficShaper) error {
+		fundingBlob := link.FundingCustomBlob()
+		shouldHandle, err := ts.HandleTraffic(cid, fundingBlob)
+		if err != nil {
+			return fmt.Errorf("traffic shaper failed to decide "+
+				"whether to handle traffic: %w", err)
 		}
 
-		return amount
+		log.Debugf("ShortChannelID=%v: external traffic shaper is "+
+			"handling traffic: %v", cid, shouldHandle)
+
+		// If this channel isn't handled by the external traffic shaper,
+		// we'll return early.
+		if !shouldHandle {
+			return nil
+		}
+
+		// Ask for a specific bandwidth to be used for the channel.
+		commitmentBlob := link.CommitmentCustomBlob()
+		auxBandwidth, err = ts.PaymentBandwidth(
+			htlcBlob, commitmentBlob,
+		)
+		if err != nil {
+			return fmt.Errorf("failed to get bandwidth from "+
+				"external traffic shaper: %w", err)
+		}
+
+		log.Debugf("ShortChannelID=%v: external traffic shaper "+
+			"reported available bandwidth: %v", cid, auxBandwidth)
+
+		// TODO(guggero): Still need to ask the link if it can add an
+		// HTLC? To avoid running over the on-chain max htlc limit?
+
+		auxBandwidthDetermined = true
+		return nil
+	})
+	if err != nil {
+		log.Errorf("ShortChannelID=%v: failed to get bandwidth from "+
+			"external traffic shaper: %v", cid, err)
+		return 0
 	}
 
-	// If the traffic shaper is present and it returned false, we want to
-	// skip this channel.
-	if b.trafficShaper.IsSome() {
-		log.Warnf("ShortChannelID=%v: payment not allowed by traffic "+
-			"shaper", cid)
-		return 0
+	// If the external traffic shaper determined the bandwidth, we'll return
+	// that value, even if it is zero (which would mean no bandwidth is
+	// available on that channel).
+	if auxBandwidthDetermined {
+		return auxBandwidth
 	}
 
 	// If our link isn't currently in a state where it can add

routing/pathfind.go

@@ -465,7 +465,8 @@ type PathFindingConfig struct {
 // available balance.
 func getOutgoingBalance(node route.Vertex, outgoingChans map[uint64]struct{},
 	bandwidthHints bandwidthHints,
-	g routingGraph, htlcBlob fn.Option[tlv.Blob]) (lnwire.MilliSatoshi, lnwire.MilliSatoshi, error) {
+	g routingGraph, htlcBlob fn.Option[tlv.Blob]) (lnwire.MilliSatoshi,
+	lnwire.MilliSatoshi, error) {
 
 	var max, total lnwire.MilliSatoshi
 

routing/payment_lifecycle.go

@@ -6,7 +6,6 @@ import (
 	"time"
 
 	"github.com/btcsuite/btcd/btcec/v2"
-	"github.com/btcsuite/btcd/btcutil"
 	"github.com/davecgh/go-spew/spew"
 	sphinx "github.com/lightningnetwork/lightning-onion"
 	"github.com/lightningnetwork/lnd/channeldb"
@@ -691,22 +690,25 @@ func (p *paymentLifecycle) sendAttempt(
 
 	// If a hook exists that may affect our outgoing message, we call it now
 	// and apply its side effects to the UpdateAddHTLC message.
-	var err error
+	err := fn.MapOptionZ(
+		p.router.cfg.TrafficShaper,
+		func(ts TlvTrafficShaper) error {
+			newAmt, newData, err := ts.ProduceHtlcExtraData(
+				rt.TotalAmount, htlcAdd.ExtraData,
+			)
+			if err != nil {
+				return err
+			}
 
-	p.router.cfg.TrafficShaper.WhenSome(func(ts TlvTrafficShaper) {
-		var newAmt btcutil.Amount
-		var newData []byte
-
-		newAmt, newData, err = ts.ProduceHtlcExtraData(
-			htlcAdd.ExtraData, uint64(htlcAdd.Amount.ToSatoshis()),
-		)
-
-		htlcAdd.ExtraData = newData
-		htlcAdd.Amount = lnwire.MilliSatoshi(newAmt * 1000)
-	})
+			htlcAdd.ExtraData = newData
+			htlcAdd.Amount = lnwire.MilliSatoshi(newAmt * 1000)
 
+			return nil
+		},
+	)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("traffic shaper failed to produce "+
+			"extra data: %w", err)
 	}
 
 	// Generate the raw encoded sphinx packet to be included along