From 6780f74c8730e54b012ce0e5ccdd811848e9b868 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Wed, 24 Mar 2021 19:51:03 -0700
Subject: [PATCH] channeldb/invoices: set AMP HTLC preimages when settling
 invoice

---
 channeldb/invoice_test.go | 99 ++++++++++++++++++++++++++++++++++++++-
 channeldb/invoices.go     | 31 +++++++++++-
 2 files changed, 128 insertions(+), 2 deletions(-)

diff --git a/channeldb/invoice_test.go b/channeldb/invoice_test.go
index 0b27d8df..510e0802 100644
--- a/channeldb/invoice_test.go
+++ b/channeldb/invoice_test.go
@@ -17,7 +17,15 @@ import (
 
 var (
 	emptyFeatures = lnwire.NewFeatureVector(nil, lnwire.Features)
-	testNow       = time.Unix(1, 0)
+	ampFeatures   = lnwire.NewFeatureVector(
+		lnwire.NewRawFeatureVector(
+			lnwire.TLVOnionPayloadOptional,
+			lnwire.PaymentAddrOptional,
+			lnwire.AMPRequired,
+		),
+		lnwire.Features,
+	)
+	testNow = time.Unix(1, 0)
 )
 
 func randInvoice(value lnwire.MilliSatoshi) (*Invoice, error) {
@@ -1571,6 +1579,95 @@ func TestUnexpectedInvoicePreimage(t *testing.T) {
 	require.Error(t, ErrUnexpectedInvoicePreimage, err)
 }
 
+type updateHTLCPreimageTestCase struct {
+	name               string
+	settleSamePreimage bool
+	expError           error
+}
+
+// TestUpdateHTLCPreimages asserts various properties of setting HTLC-level
+// preimages on invoice state transitions.
+func TestUpdateHTLCPreimages(t *testing.T) {
+	t.Parallel()
+
+	tests := []updateHTLCPreimageTestCase{
+		{
+			name:               "same preimage on settle",
+			settleSamePreimage: true,
+			expError:           nil,
+		},
+		{
+			name:               "diff preimage on settle",
+			settleSamePreimage: false,
+			expError:           ErrHTLCPreimageAlreadyExists,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			testUpdateHTLCPreimages(t, test)
+		})
+	}
+}
+
+func testUpdateHTLCPreimages(t *testing.T, test updateHTLCPreimageTestCase) {
+	db, cleanup, err := MakeTestDB()
+	defer cleanup()
+	require.NoError(t, err, "unable to make test db")
+
+	// We'll start out by creating an invoice and writing it to the DB.
+	amt := lnwire.NewMSatFromSatoshis(1000)
+	invoice, err := randInvoice(amt)
+	require.Nil(t, err)
+
+	preimage := *invoice.Terms.PaymentPreimage
+	payHash := preimage.Hash()
+
+	// Set AMP-specific features so that we can settle with HTLC-level
+	// preimages.
+	invoice.Terms.Features = ampFeatures
+
+	_, err = db.AddInvoice(invoice, payHash)
+	require.Nil(t, err)
+
+	setID := &[32]byte{1}
+
+	// Update the invoice with an accepted HTLC that also accepts the
+	// invoice.
+	ref := InvoiceRefByAddr(invoice.Terms.PaymentAddr)
+	dbInvoice, err := db.UpdateInvoice(ref, updateAcceptAMPHtlc(0, amt, setID, true))
+	require.Nil(t, err)
+
+	htlcPreimages := make(map[CircuitKey]lntypes.Preimage)
+	for key := range dbInvoice.Htlcs {
+		// Set the either the same preimage used to accept above, or a
+		// blank preimage depending on the test case.
+		var pre lntypes.Preimage
+		if test.settleSamePreimage {
+			pre = preimage
+		}
+		htlcPreimages[key] = pre
+	}
+
+	updateInvoice := func(invoice *Invoice) (*InvoiceUpdateDesc, error) {
+		update := &InvoiceUpdateDesc{
+			State: &InvoiceStateUpdateDesc{
+				Preimage:      nil,
+				NewState:      ContractSettled,
+				HTLCPreimages: htlcPreimages,
+				SetID:         setID,
+			},
+		}
+
+		return update, nil
+	}
+
+	// Now settle the HTLC set and assert the resulting error.
+	_, err = db.UpdateInvoice(ref, updateInvoice)
+	require.Equal(t, test.expError, err)
+}
+
 // TestDeleteInvoices tests that deleting a list of invoices will succeed
 // if all delete references are valid, or will fail otherwise.
 func TestDeleteInvoices(t *testing.T) {
diff --git a/channeldb/invoices.go b/channeldb/invoices.go
index a4ddd950..6d6c4d5c 100644
--- a/channeldb/invoices.go
+++ b/channeldb/invoices.go
@@ -131,6 +131,12 @@ var (
 	ErrUnexpectedInvoicePreimage = errors.New(
 		"unexpected invoice preimage provided on settle",
 	)
+
+	// ErrHTLCPreimageAlreadyExists is returned when trying to set an
+	// htlc-level preimage but one is already known.
+	ErrHTLCPreimageAlreadyExists = errors.New(
+		"htlc-level preimage already exists",
+	)
 )
 
 // ErrDuplicateSetID is an error returned when attempting to adding an AMP HTLC
@@ -646,6 +652,11 @@ type InvoiceStateUpdateDesc struct {
 	// Preimage must be set to the preimage when NewState is settled.
 	Preimage *lntypes.Preimage
 
+	// HTLCPreimages set the HTLC-level preimages stored for AMP HTLCs.
+	// These are only learned when settling the invoice as a whole. Must be
+	// set when settling an invoice with non-nil SetID.
+	HTLCPreimages map[CircuitKey]lntypes.Preimage
+
 	// SetID identifies a specific set of HTLCs destined for the same
 	// invoice as part of a larger AMP payment. This value will be nil for
 	// legacy or MPP payments.
@@ -1914,7 +1925,25 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 	// the process by updating the state transitions for individual HTLCs
 	// and recalculate the total amount paid to the invoice.
 	var amtPaid lnwire.MilliSatoshi
-	for _, htlc := range invoice.Htlcs {
+	for key, htlc := range invoice.Htlcs {
+		// Set the HTLC preimage for any AMP HTLCs.
+		if setID != nil {
+			preimage, ok := update.State.HTLCPreimages[key]
+			switch {
+
+			// If we don't already have a preiamge for this HTLC, we
+			// can set it now.
+			case ok && htlc.AMP.Preimage == nil:
+				htlc.AMP.Preimage = &preimage
+
+			// Otherwise, prevent over-writing an existing preimage.
+			// Ignore the case where the preimage is identical.
+			case ok && *htlc.AMP.Preimage != preimage:
+				return nil, ErrHTLCPreimageAlreadyExists
+
+			}
+		}
+
 		// The invoice state may have changed and this could have
 		// implications for the states of the individual htlcs. Align
 		// the htlc state with the current invoice state.