routing/route: add AMP record to payload size calcs

routing/route/route.go

@@ -34,6 +34,10 @@ var (
 	// record to an intermediate hop, only final hops can receive MPP
 	// records.
 	ErrIntermediateMPPHop = errors.New("cannot send MPP to intermediate")
+
+	// ErrAMPMissingMPP is returned when the caller tries to attach an AMP
+	// record but no MPP record is presented for the final hop.
+	ErrAMPMissingMPP = errors.New("cannot send AMP without MPP record")
 )
 
 // Vertex is a simple alias for the serialization of a compressed Bitcoin
@@ -111,6 +115,10 @@ type Hop struct {
 	// only be set for the final hop.
 	MPP *record.MPP
 
+	// AMP encapsulates the data required for option_amp. This field should
+	// only be set for the final hop.
+	AMP *record.AMP
+
 	// CustomRecords if non-nil are a set of additional TLV records that
 	// should be included in the forwarding instructions for this node.
 	CustomRecords record.CustomSet
@@ -168,6 +176,18 @@ func (h *Hop) PackHopPayload(w io.Writer, nextChanID uint64) error {
 		}
 	}
 
+	// If an AMP record is destined for this hop, ensure that we only ever
+	// attach it if we also have an MPP record. We can infer that this is
+	// already a final hop if MPP is non-nil otherwise we would have exited
+	// above.
+	if h.AMP != nil {
+		if h.MPP != nil {
+			records = append(records, h.AMP.Record())
+		} else {
+			return ErrAMPMissingMPP
+		}
+	}
+
 	// Append any custom types destined for this hop.
 	tlvRecords := tlv.MapToRecords(h.CustomRecords)
 	records = append(records, tlvRecords...)
@@ -217,6 +237,11 @@ func (h *Hop) PayloadSize(nextChanID uint64) uint64 {
 		addRecord(record.MPPOnionType, h.MPP.PayloadSize())
 	}
 
+	// Add amp if present.
+	if h.AMP != nil {
+		addRecord(record.AMPOnionType, h.AMP.PayloadSize())
+	}
+
 	// Add custom records.
 	for k, v := range h.CustomRecords {
 		addRecord(tlv.Type(k), uint64(len(v)))

routing/route/route_test.go

@@ -71,8 +71,8 @@ var (
 	testAddr = [32]byte{0x01, 0x02}
 )
 
-// TestMPPHop asserts that a Hop will encode a non-nil to final nodes, and fail
+// TestMPPHop asserts that a Hop will encode a non-nil MPP to final nodes, and
-// when trying to send to intermediaries.
+// fail when trying to send to intermediaries.
 func TestMPPHop(t *testing.T) {
 	t.Parallel()
 
@@ -123,6 +123,7 @@ func TestPayloadSize(t *testing.T) {
 			AmtToForward:     1200,
 			OutgoingTimeLock: 700000,
 			MPP:              record.NewMPP(500, [32]byte{}),
+			AMP:              record.NewAMP([32]byte{}, [32]byte{}, 8),
 			CustomRecords: map[uint64][]byte{
 				100000:  {1, 2, 3},
 				1000000: {4, 5},