channeldb/invoice: persist optional AMP fields on InvoiceHTLC

channeldb/invoice_test.go

@@ -1133,6 +1133,70 @@ func TestCustomRecords(t *testing.T) {
 	)
 }
 
+// TestInvoiceHtlcAMPFields asserts that the set id and preimage fields are
+// properly recorded when updating an invoice.
+func TestInvoiceHtlcAMPFields(t *testing.T) {
+	t.Run("amp", func(t *testing.T) {
+		testInvoiceHtlcAMPFields(t, true)
+	})
+	t.Run("no amp", func(t *testing.T) {
+		testInvoiceHtlcAMPFields(t, false)
+	})
+}
+
+func testInvoiceHtlcAMPFields(t *testing.T, isAMP bool) {
+	db, cleanUp, err := MakeTestDB()
+	defer cleanUp()
+	require.Nil(t, err)
+
+	testInvoice, err := randInvoice(1000)
+	require.Nil(t, err)
+
+	payHash := testInvoice.Terms.PaymentPreimage.Hash()
+	_, err = db.AddInvoice(testInvoice, payHash)
+	require.Nil(t, err)
+
+	// Accept an htlc with custom records on this invoice.
+	key := CircuitKey{ChanID: lnwire.NewShortChanIDFromInt(1), HtlcID: 4}
+	records := make(map[uint64][]byte)
+
+	var ampData *InvoiceHtlcAMPData
+	if isAMP {
+		amp := record.NewAMP([32]byte{1}, [32]byte{2}, 3)
+		preimage := &lntypes.Preimage{4}
+
+		ampData = &InvoiceHtlcAMPData{
+			Record:   *amp,
+			Hash:     preimage.Hash(),
+			Preimage: preimage,
+		}
+	}
+
+	ref := InvoiceRefByHash(payHash)
+	_, err = db.UpdateInvoice(ref,
+		func(invoice *Invoice) (*InvoiceUpdateDesc, error) {
+			return &InvoiceUpdateDesc{
+				AddHtlcs: map[CircuitKey]*HtlcAcceptDesc{
+					key: {
+						Amt:           500,
+						AMP:           ampData,
+						CustomRecords: records,
+					},
+				},
+			}, nil
+		},
+	)
+	require.Nil(t, err)
+
+	// Retrieve the invoice from that database and verify that the AMP
+	// fields are as expected.
+	dbInvoice, err := db.LookupInvoice(ref)
+	require.Nil(t, err)
+
+	require.Equal(t, 1, len(dbInvoice.Htlcs))
+	require.Equal(t, ampData, dbInvoice.Htlcs[key].AMP)
+}
+
 // TestInvoiceRef asserts that the proper identifiers are returned from an
 // InvoiceRef depending on the constructor used.
 func TestInvoiceRef(t *testing.T) {

channeldb/invoices.go

@@ -135,6 +135,9 @@ const (
 	expiryHeightType tlv.Type = 13
 	htlcStateType    tlv.Type = 15
 	mppTotalAmtType  tlv.Type = 17
+	htlcAMPType      tlv.Type = 19
+	htlcHashType     tlv.Type = 21
+	htlcPreimageType tlv.Type = 23
 
 	// A set of tlv type definitions used to serialize invoice bodiees.
 	//
@@ -407,6 +410,60 @@ type InvoiceHTLC struct {
 	// CustomRecords contains the custom key/value pairs that accompanied
 	// the htlc.
 	CustomRecords record.CustomSet
+
+	// AMP encapsulates additional data relevant to AMP HTLCs. This includes
+	// the AMP onion record, in addition to the HTLC's payment hash and
+	// preimage since these are unique to each AMP HTLC, and not the invoice
+	// as a whole.
+	//
+	// NOTE: This value will only be set for AMP HTLCs.
+	AMP *InvoiceHtlcAMPData
+}
+
+// InvoiceHtlcAMPData is a struct hodling the additional metadata stored for
+// each received AMP HTLC. This includes the AMP onion record, in addition to
+// the HTLC's payment hash and preimage.
+type InvoiceHtlcAMPData struct {
+	// AMP is a copy of the AMP record presented in the onion payload
+	// containing the information necessary to correlate and settle a
+	// spontaneous HTLC set. Newly accepted legacy keysend payments will
+	// also have this field set as we automatically promote them into an AMP
+	// payment for internal processing.
+	Record record.AMP
+
+	// Hash is an HTLC-level payment hash that is stored only for AMP
+	// payments. This is done because an AMP HTLC will carry a different
+	// payment hash from the invoice it might be satisfying, so we track the
+	// payment hashes individually to able to compute whether or not the
+	// reconstructed preimage correctly matches the HTLC's hash.
+	Hash lntypes.Hash
+
+	// Preimage is an HTLC-level preimage that satisfies the AMP HTLC's
+	// Hash. The preimage will be be derived either from secret share
+	// reconstruction of the shares in the AMP payload.
+	//
+	// NOTE: Preimage will only be present once the HTLC is in
+	// HltcStateSetteled.
+	Preimage *lntypes.Preimage
+}
+
+// Copy returns a deep copy of the InvoiceHtlcAMPData.
+func (d *InvoiceHtlcAMPData) Copy() *InvoiceHtlcAMPData {
+	if d == nil {
+		return nil
+	}
+
+	var preimage *lntypes.Preimage
+	if d.Preimage != nil {
+		pimg := *d.Preimage
+		preimage = &pimg
+	}
+
+	return &InvoiceHtlcAMPData{
+		Record:   d.Record,
+		Hash:     d.Hash,
+		Preimage: preimage,
+	}
 }
 
 // HtlcAcceptDesc describes the details of a newly accepted htlc.
@@ -427,6 +484,14 @@ type HtlcAcceptDesc struct {
 	// CustomRecords contains the custom key/value pairs that accompanied
 	// the htlc.
 	CustomRecords record.CustomSet
+
+	// AMP encapsulates additional data relevant to AMP HTLCs. This includes
+	// the AMP onion record, in addition to the HTLC's payment hash and
+	// preimage since these are unique to each AMP HTLC, and not the invoice
+	// as a whole.
+	//
+	// NOTE: This value will only be set for AMP HTLCs.
+	AMP *InvoiceHtlcAMPData
 }
 
 // InvoiceUpdateDesc describes the changes that should be applied to the
@@ -1205,6 +1270,29 @@ func serializeHtlcs(w io.Writer, htlcs map[CircuitKey]*InvoiceHTLC) error {
 			tlv.MakePrimitiveRecord(mppTotalAmtType, &mppTotalAmt),
 		)
 
+		if htlc.AMP != nil {
+			setIDRecord := tlv.MakeDynamicRecord(
+				htlcAMPType, &htlc.AMP.Record,
+				htlc.AMP.Record.PayloadSize,
+				record.AMPEncoder, record.AMPDecoder,
+			)
+			records = append(records, setIDRecord)
+
+			hash32 := [32]byte(htlc.AMP.Hash)
+			hashRecord := tlv.MakePrimitiveRecord(
+				htlcHashType, &hash32,
+			)
+			records = append(records, hashRecord)
+
+			if htlc.AMP.Preimage != nil {
+				preimage32 := [32]byte(*htlc.AMP.Preimage)
+				preimageRecord := tlv.MakePrimitiveRecord(
+					htlcPreimageType, &preimage32,
+				)
+				records = append(records, preimageRecord)
+			}
+		}
+
 		// Convert the custom records to tlv.Record types that are ready
 		// for serialization.
 		customRecords := tlv.MapToRecords(htlc.CustomRecords)
@@ -1374,6 +1462,9 @@ func deserializeHtlcs(r io.Reader) (map[CircuitKey]*InvoiceHTLC, error) {
 			state                   uint8
 			acceptTime, resolveTime uint64
 			amt, mppTotalAmt        uint64
+			amp                     = &record.AMP{}
+			hash32                  = &[32]byte{}
+			preimage32              = &[32]byte{}
 		)
 		tlvStream, err := tlv.NewStream(
 			tlv.MakePrimitiveRecord(chanIDType, &chanID),
@@ -1387,6 +1478,12 @@ func deserializeHtlcs(r io.Reader) (map[CircuitKey]*InvoiceHTLC, error) {
 			tlv.MakePrimitiveRecord(expiryHeightType, &htlc.Expiry),
 			tlv.MakePrimitiveRecord(htlcStateType, &state),
 			tlv.MakePrimitiveRecord(mppTotalAmtType, &mppTotalAmt),
+			tlv.MakeDynamicRecord(
+				htlcAMPType, amp, amp.PayloadSize,
+				record.AMPEncoder, record.AMPDecoder,
+			),
+			tlv.MakePrimitiveRecord(htlcHashType, hash32),
+			tlv.MakePrimitiveRecord(htlcPreimageType, preimage32),
 		)
 		if err != nil {
 			return nil, err
@@ -1397,12 +1494,35 @@ func deserializeHtlcs(r io.Reader) (map[CircuitKey]*InvoiceHTLC, error) {
 			return nil, err
 		}
 
+		if _, ok := parsedTypes[htlcAMPType]; !ok {
+			amp = nil
+		}
+
+		var preimage *lntypes.Preimage
+		if _, ok := parsedTypes[htlcPreimageType]; ok {
+			pimg := lntypes.Preimage(*preimage32)
+			preimage = &pimg
+		}
+
+		var hash *lntypes.Hash
+		if _, ok := parsedTypes[htlcHashType]; ok {
+			h := lntypes.Hash(*hash32)
+			hash = &h
+		}
+
 		key.ChanID = lnwire.NewShortChanIDFromInt(chanID)
 		htlc.AcceptTime = time.Unix(0, int64(acceptTime))
 		htlc.ResolveTime = time.Unix(0, int64(resolveTime))
 		htlc.State = HtlcState(state)
 		htlc.Amt = lnwire.MilliSatoshi(amt)
 		htlc.MppTotalAmt = lnwire.MilliSatoshi(mppTotalAmt)
+		if amp != nil && hash != nil {
+			htlc.AMP = &InvoiceHtlcAMPData{
+				Record:   *amp,
+				Hash:     *hash,
+				Preimage: preimage,
+			}
+		}
 
 		// Reconstruct the custom records fields from the parsed types
 		// map return from the tlv parser.
@@ -1431,6 +1551,8 @@ func copyInvoiceHTLC(src *InvoiceHTLC) *InvoiceHTLC {
 		result.CustomRecords[k] = v
 	}
 
+	result.AMP = src.AMP.Copy()
+
 	return &result
 }
 
@@ -1531,6 +1653,7 @@ func (d *DB) updateInvoice(hash lntypes.Hash, invoices, settleIndex kvdb.RwBucke
 			AcceptTime:    now,
 			State:         HtlcStateAccepted,
 			CustomRecords: htlcUpdate.CustomRecords,
+			AMP:           htlcUpdate.AMP.Copy(),
 		}
 
 		invoice.Htlcs[key] = htlc