routing: use first hop records on path finding

routing/pathfind.go

@@ -466,6 +466,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 lnwire.CustomRecords
 }
 
 // PathFindingConfig defines global parameters that control the trade-off in
@@ -524,7 +528,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
 	}
@@ -1446,3 +1459,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,11 +11,13 @@ 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"
 	"github.com/lightningnetwork/lnd/routing/route"
 	"github.com/lightningnetwork/lnd/routing/shards"
+	"github.com/lightningnetwork/lnd/tlv"
 )
 
 // ErrPaymentLifecycleExiting is used when waiting for htlc attempt result, but
@@ -275,6 +277,13 @@ lifecycle:
 
 		log.Tracef("Found route: %s", spew.Sdump(rt.Hops))
 
+		// Allow the traffic shaper to add custom records to the
+		// outgoing HTLC and also adjust the amount if needed.
+		err = p.amendFirstHopData(rt)
+		if err != nil {
+			return exitWithErr(err)
+		}
+
 		// We found a route to try, create a new HTLC attempt to try.
 		attempt, err := p.registerAttempt(rt, ps.RemainingAmt)
 		if err != nil {
@@ -668,8 +677,10 @@ func (p *paymentLifecycle) createNewPaymentAttempt(rt *route.Route,
 func (p *paymentLifecycle) sendAttempt(
 	attempt *channeldb.HTLCAttempt) (*attemptResult, error) {
 
-	log.Debugf("Sending HTLC attempt(id=%v, amt=%v) for payment %v",
-		attempt.AttemptID, attempt.Route.TotalAmount, p.identifier)
+	log.Debugf("Sending HTLC attempt(id=%v, total_amt=%v, first_hop_amt=%d"+
+		") for payment %v", attempt.AttemptID,
+		attempt.Route.TotalAmount, attempt.Route.FirstHopAmount.Val,
+		p.identifier)
 
 	rt := attempt.Route
 
@@ -680,10 +691,10 @@ func (p *paymentLifecycle) sendAttempt(
 	// this packet will be used to route the payment through the network,
 	// starting with the first-hop.
 	htlcAdd := &lnwire.UpdateAddHTLC{
-		Amount:        rt.TotalAmount,
+		Amount:        rt.FirstHopAmount.Val.Int(),
 		Expiry:        rt.TotalTimeLock,
 		PaymentHash:   *attempt.Hash,
-		CustomRecords: p.firstHopCustomRecords,
+		CustomRecords: rt.FirstHopWireCustomRecords,
 	}
 
 	// Generate the raw encoded sphinx packet to be included along
@@ -722,6 +733,75 @@ func (p *paymentLifecycle) sendAttempt(
 	}, nil
 }
 
+// amendFirstHopData 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) amendFirstHopData(rt *route.Route) error {
+	// The first hop amount on the route is the full route amount if not
+	// overwritten by the traffic shaper. So we set the initial value now
+	// and potentially overwrite it later.
+	rt.FirstHopAmount = tlv.NewRecordT[tlv.TlvType0](
+		tlv.NewBigSizeT(rt.TotalAmount),
+	)
+
+	// By default, we set the first hop custom records to the initial
+	// value requested by the RPC. The traffic shaper may overwrite this
+	// value.
+	rt.FirstHopWireCustomRecords = p.firstHopCustomRecords
+
+	// 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, p.firstHopCustomRecords,
+			)
+			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) {
+		rt.FirstHopWireCustomRecords = records
+	})
+	result.amount.WhenSome(func(amount lnwire.MilliSatoshi) {
+		rt.FirstHopAmount = tlv.NewRecordT[tlv.TlvType0](
+			tlv.NewBigSizeT(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

@@ -278,6 +278,7 @@ func (p *paymentSession) RequestRoute(maxAmt, feeLimit lnwire.MilliSatoshi,
 		PaymentAddr:           p.payment.PaymentAddr,
 		Amp:                   p.payment.amp,
 		Metadata:              p.payment.Metadata,
+		FirstHopCustomRecords: firstHopCustomRecords,
 	}
 
 	finalHtlcExpiry := int32(height) + int32(finalCltvDelta)

routing/router.go

@@ -1188,6 +1188,13 @@ func (r *ChannelRouter) sendToRoute(htlcHash lntypes.Hash, rt *route.Route,
 		firstHopCustomRecords,
 	)
 
+	// Allow the traffic shaper to add custom records to the outgoing HTLC
+	// and also adjust the amount if needed.
+	err = p.amendFirstHopData(rt)
+	if err != nil {
+		return nil, err
+	}
+
 	// We found a route to try, create a new HTLC attempt to try.
 	//
 	// NOTE: we use zero `remainingAmt` here to simulate the same effect of