From 7965cb08db9ece400c575729494ea37fa3924167 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Wed, 18 Dec 2019 23:56:23 -0800
Subject: [PATCH] routing/pathfind: consolidate population of custom records

---
 routing/pathfind.go | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/routing/pathfind.go b/routing/pathfind.go
index 7788dc83f..297392a0e 100644
--- a/routing/pathfind.go
+++ b/routing/pathfind.go
@@ -151,6 +151,7 @@ func newRoute(sourceVertex route.Vertex,
 			fee              lnwire.MilliSatoshi
 			outgoingTimeLock uint32
 			tlvPayload       bool
+			customRecords    record.CustomSet
 		)
 
 		// Define a helper function that checks this edge's feature
@@ -182,6 +183,14 @@ func newRoute(sourceVertex route.Vertex,
 			// last link in the route.
 			totalTimeLock += uint32(finalHop.cltvDelta)
 			outgoingTimeLock = totalTimeLock
+
+			// Attach any custom records to the final hop if the
+			// receiver supports TLV.
+			if !tlvPayload && finalHop.records != nil {
+				return nil, errors.New("cannot attach " +
+					"custom records")
+			}
+			customRecords = finalHop.records
 		} else {
 			// The amount that the current hop needs to forward is
 			// equal to the incoming amount of the next hop.
@@ -210,12 +219,7 @@ func newRoute(sourceVertex route.Vertex,
 			AmtToForward:     amtToForward,
 			OutgoingTimeLock: outgoingTimeLock,
 			LegacyPayload:    !tlvPayload,
-		}
-
-		// If this is the last hop, then we'll populate any TLV records
-		// destined for it.
-		if i == len(pathEdges)-1 && len(finalHop.records) != 0 {
-			currentHop.CustomRecords = finalHop.records
+			CustomRecords:    customRecords,
 		}
 
 		hops = append([]*route.Hop{currentHop}, hops...)