routing: use first hop records on path finding

routing/pathfind.go

@@ -467,6 +467,10 @@ type RestrictParams struct {
 	// BlindedPaymentPathSet is necessary to determine the hop size of the
 	// last/exit hop.
 	BlindedPaymentPathSet *BlindedPaymentPathSet
+
+	// FirstHopCustomRecords includes any records that should be included in
+	// the update_add_htlc message towards our peer.
+	FirstHopCustomRecords record.CustomSet
 }
 
 // PathFindingConfig defines global parameters that control the trade-off in
@@ -526,7 +530,16 @@ func getOutgoingBalance(node route.Vertex, outgoingChans map[uint64]struct{},
 			max = bandwidth
 		}
 
-		total += bandwidth
+		var overflow bool
+		total, overflow = overflowSafeAdd(total, bandwidth)
+		if overflow {
+			// If the current total and the bandwidth would
+			// overflow the maximum value, we set the total to the
+			// maximum value. Which is more milli-satoshis than are
+			// in existence anyway, so the actual value is
+			// irrelevant.
+			total = lnwire.MilliSatoshi(math.MaxUint64)
+		}
 
 		return nil
 	}
@@ -612,12 +625,30 @@ func findPath(g *graphParams, r *RestrictParams, cfg *PathFindingConfig,
 		}
 	}
 
+	// Assemble any custom data we want to send to the first hop only.
+	var firstHopData fn.Option[tlv.Blob]
+	if len(r.FirstHopCustomRecords) > 0 {
+		firstHopRecords := lnwire.CustomRecords(r.FirstHopCustomRecords)
+		if err := firstHopRecords.Validate(); err != nil {
+			return nil, 0, fmt.Errorf("invalid first hop custom "+
+				"records: %w", err)
+		}
+
+		firstHopBlob, err := firstHopRecords.Serialize()
+		if err != nil {
+			return nil, 0, fmt.Errorf("unable to serialize first "+
+				"hop custom records: %w", err)
+		}
+
+		firstHopData = fn.Some(firstHopBlob)
+	}
+
 	// If we are routing from ourselves, check that we have enough local
 	// balance available.
 	if source == self {
 		max, total, err := getOutgoingBalance(
 			self, outgoingChanMap, g.bandwidthHints, g.graph,
-			fn.None[tlv.Blob](),
+			firstHopData,
 		)
 		if err != nil {
 			return nil, 0, err
@@ -1054,7 +1085,7 @@ func findPath(g *graphParams, r *RestrictParams, cfg *PathFindingConfig,
 
 			edge := edgeUnifier.getEdge(
 				netAmountReceived, g.bandwidthHints,
-				partialPath.outboundFee, fn.None[tlv.Blob](),
+				partialPath.outboundFee, firstHopData,
 			)
 
 			if edge == nil {
@@ -1449,3 +1480,15 @@ func lastHopPayloadSize(r *RestrictParams, finalHtlcExpiry int32,
 	// The final hop does not have a short chanID set.
 	return finalHop.PayloadSize(0)
 }
+
+// overflowSafeAdd adds two MilliSatoshi values and returns the result. If an
+// overflow could occur, zero is returned instead and the boolean is set to
+// true.
+func overflowSafeAdd(x, y lnwire.MilliSatoshi) (lnwire.MilliSatoshi, bool) {
+	if y > math.MaxUint64-x {
+		// Overflow would occur, return 0 and set overflow flag.
+		return 0, true
+	}
+
+	return x + y, false
+}

routing/payment_lifecycle.go

@@ -11,6 +11,7 @@ import (
 	sphinx "github.com/lightningnetwork/lightning-onion"
 	"github.com/lightningnetwork/lnd/channeldb"
 	"github.com/lightningnetwork/lnd/channeldb/models"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/htlcswitch"
 	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwire"
@@ -366,7 +367,8 @@ func (p *paymentLifecycle) requestRoute(
 	// Query our payment session to construct a route.
 	rt, err := p.paySession.RequestRoute(
 		ps.RemainingAmt, remainingFees,
-		uint32(ps.NumAttemptsInFlight), uint32(p.currentHeight), nil,
+		uint32(ps.NumAttemptsInFlight), uint32(p.currentHeight),
+		p.firstHopCustomRecords,
 	)
 
 	// Exit early if there's no error.
@@ -686,6 +688,17 @@ func (p *paymentLifecycle) sendAttempt(
 		CustomRecords: lnwire.CustomRecords(p.firstHopCustomRecords),
 	}
 
+	// Allow the traffic shaper to add custom records to the outgoing HTLC
+	// and also adjust the amount if needed.
+	err := p.amendHtlcCustomRecords(rt, htlcAdd)
+	if err != nil {
+		log.Errorf("Failed to amend custom records: attempt=%d in "+
+			"payment=%v, err:%v", attempt.AttemptID,
+			p.identifier, err)
+
+		return p.failAttempt(attempt.AttemptID, err)
+	}
+
 	// Generate the raw encoded sphinx packet to be included along
 	// with the htlcAdd message that we send directly to the
 	// switch.
@@ -722,6 +735,63 @@ func (p *paymentLifecycle) sendAttempt(
 	}, nil
 }
 
+// amendHtlcCustomRecords is a function that calls the traffic shaper to allow
+// it to add custom records to the outgoing HTLC and also adjust the amount if
+// needed.
+func (p *paymentLifecycle) amendHtlcCustomRecords(rt route.Route,
+	htlcAdd *lnwire.UpdateAddHTLC) error {
+
+	// extraDataRequest is a helper struct to pass the custom records and
+	// amount back from the traffic shaper.
+	type extraDataRequest struct {
+		customRecords fn.Option[lnwire.CustomRecords]
+
+		amount fn.Option[lnwire.MilliSatoshi]
+	}
+
+	// If a hook exists that may affect our outgoing message, we call it now
+	// and apply its side effects to the UpdateAddHTLC message.
+	result, err := fn.MapOptionZ(
+		p.router.cfg.TrafficShaper,
+		func(ts TlvTrafficShaper) fn.Result[extraDataRequest] {
+			newAmt, newRecords, err := ts.ProduceHtlcExtraData(
+				rt.TotalAmount, htlcAdd.CustomRecords,
+			)
+			if err != nil {
+				return fn.Err[extraDataRequest](err)
+			}
+
+			// Make sure we only received valid records.
+			if err := newRecords.Validate(); err != nil {
+				return fn.Err[extraDataRequest](err)
+			}
+
+			log.Debugf("TLV traffic shaper returned custom "+
+				"records %v and amount %d msat for HTLC",
+				spew.Sdump(newRecords), newAmt)
+
+			return fn.Ok(extraDataRequest{
+				customRecords: fn.Some(newRecords),
+				amount:        fn.Some(newAmt),
+			})
+		},
+	).Unpack()
+	if err != nil {
+		return fmt.Errorf("traffic shaper failed to produce extra "+
+			"data: %w", err)
+	}
+
+	// Apply the side effects to the UpdateAddHTLC message.
+	result.customRecords.WhenSome(func(records lnwire.CustomRecords) {
+		htlcAdd.CustomRecords = records
+	})
+	result.amount.WhenSome(func(amount lnwire.MilliSatoshi) {
+		htlcAdd.Amount = amount
+	})
+
+	return nil
+}
+
 // failAttemptAndPayment fails both the payment and its attempt via the
 // router's control tower, which marks the payment as failed in db.
 func (p *paymentLifecycle) failPaymentAndAttempt(

routing/payment_session.go

@@ -268,16 +268,17 @@ func (p *paymentSession) RequestRoute(maxAmt, feeLimit lnwire.MilliSatoshi,
 	// to our destination, respecting the recommendations from
 	// MissionControl.
 	restrictions := &RestrictParams{
-		ProbabilitySource:  p.missionControl.GetProbability,
-		FeeLimit:           feeLimit,
-		OutgoingChannelIDs: p.payment.OutgoingChannelIDs,
-		LastHop:            p.payment.LastHop,
-		CltvLimit:          cltvLimit,
-		DestCustomRecords:  p.payment.DestCustomRecords,
-		DestFeatures:       p.payment.DestFeatures,
-		PaymentAddr:        p.payment.PaymentAddr,
-		Amp:                p.payment.amp,
-		Metadata:           p.payment.Metadata,
+		ProbabilitySource:     p.missionControl.GetProbability,
+		FeeLimit:              feeLimit,
+		OutgoingChannelIDs:    p.payment.OutgoingChannelIDs,
+		LastHop:               p.payment.LastHop,
+		CltvLimit:             cltvLimit,
+		DestCustomRecords:     p.payment.DestCustomRecords,
+		DestFeatures:          p.payment.DestFeatures,
+		PaymentAddr:           p.payment.PaymentAddr,
+		Amp:                   p.payment.amp,
+		Metadata:              p.payment.Metadata,
+		FirstHopCustomRecords: firstHopTLVs,
 	}
 
 	finalHtlcExpiry := int32(height) + int32(finalCltvDelta)