From 0767a2cbc00da248407870267987619da17e12a6 Mon Sep 17 00:00:00 2001
From: Alex Sears <me@alexsears.com>
Date: Fri, 26 Jan 2024 15:21:43 -0500
Subject: [PATCH] routerrpc: unmarshal AMP record in route hop marshalling

---
 itest/lnd_amp_test.go             | 10 ++++++++++
 lnrpc/routerrpc/router_backend.go | 13 +++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/itest/lnd_amp_test.go b/itest/lnd_amp_test.go
index a76aedd99..bbf5cb543 100644
--- a/itest/lnd_amp_test.go
+++ b/itest/lnd_amp_test.go
@@ -410,6 +410,16 @@ func testSendPaymentAMP(ht *lntest.HarnessTest) {
 		if htlc.Status == lnrpc.HTLCAttempt_SUCCEEDED {
 			succeeded++
 		}
+
+		// When an AMP record is expected, it will only be seen on the
+		// last hop of a route. So we make sure it isn't set on any of
+		// the hops except the last one
+		for _, hop := range htlc.Route.Hops[:len(htlc.Route.Hops)-1] {
+			require.Nil(ht, hop.AmpRecord)
+		}
+
+		lastHop := htlc.Route.Hops[len(htlc.Route.Hops)-1]
+		require.NotNil(ht, lastHop.AmpRecord)
 	}
 
 	const minExpectedShards = 3
diff --git a/lnrpc/routerrpc/router_backend.go b/lnrpc/routerrpc/router_backend.go
index 49d080a32..640333734 100644
--- a/lnrpc/routerrpc/router_backend.go
+++ b/lnrpc/routerrpc/router_backend.go
@@ -613,6 +613,18 @@ func (r *RouterBackend) MarshallRoute(route *route.Route) (*lnrpc.Route, error)
 			}
 		}
 
+		var amp *lnrpc.AMPRecord
+		if hop.AMP != nil {
+			rootShare := hop.AMP.RootShare()
+			setID := hop.AMP.SetID()
+
+			amp = &lnrpc.AMPRecord{
+				RootShare:  rootShare[:],
+				SetId:      setID[:],
+				ChildIndex: hop.AMP.ChildIndex(),
+			}
+		}
+
 		resp.Hops[i] = &lnrpc.Hop{
 			ChanId:           hop.ChannelID,
 			ChanCapacity:     int64(chanCapacity),
@@ -627,6 +639,7 @@ func (r *RouterBackend) MarshallRoute(route *route.Route) (*lnrpc.Route, error)
 			CustomRecords: hop.CustomRecords,
 			TlvPayload:    !hop.LegacyPayload,
 			MppRecord:     mpp,
+			AmpRecord:     amp,
 			Metadata:      hop.Metadata,
 			EncryptedData: hop.EncryptedData,
 			TotalAmtMsat:  uint64(hop.TotalAmtMsat),