From 20dea0d6bc163e0718326054e4001aaf02f9a51e Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Fri, 26 Jun 2020 02:17:18 -0700
Subject: [PATCH] invoices/invoiceregistry: explicitly set blank pay addr for
 keysend

---
 invoices/invoiceregistry.go      | 10 +++++
 invoices/invoiceregistry_test.go | 64 +++++++++++++++++++++++---------
 2 files changed, 57 insertions(+), 17 deletions(-)

diff --git a/invoices/invoiceregistry.go b/invoices/invoiceregistry.go
index abb9f779..e2af66a9 100644
--- a/invoices/invoiceregistry.go
+++ b/invoices/invoiceregistry.go
@@ -676,6 +676,15 @@ func (i *InvoiceRegistry) processKeySend(ctx invoiceUpdateCtx) error {
 		return errors.New("final expiry too soon")
 	}
 
+	// The invoice database indexes all invoices by payment address, however
+	// legacy keysend payment do not have one. In order to avoid a new
+	// payment type on-disk wrt. to indexing, we'll continue to insert a
+	// blank payment address which is special cased in the insertion logic
+	// to not be indexed. In the future, once AMP is merged, this should be
+	// replaced by generating a random payment address on the behalf of the
+	// sender.
+	payAddr := channeldb.BlankPayAddr
+
 	// Create placeholder invoice.
 	invoice := &channeldb.Invoice{
 		CreationDate: i.cfg.Clock.Now(),
@@ -683,6 +692,7 @@ func (i *InvoiceRegistry) processKeySend(ctx invoiceUpdateCtx) error {
 			FinalCltvDelta:  finalCltvDelta,
 			Value:           amt,
 			PaymentPreimage: &preimage,
+			PaymentAddr:     payAddr,
 			Features:        features,
 		},
 	}
diff --git a/invoices/invoiceregistry_test.go b/invoices/invoiceregistry_test.go
index fa08d5db..a9131f8a 100644
--- a/invoices/invoiceregistry_test.go
+++ b/invoices/invoiceregistry_test.go
@@ -725,29 +725,59 @@ func testKeySend(t *testing.T, keySendEnabled bool) {
 		return
 	}
 
-	// Otherwise we expect no error and a settle resolution for the htlc.
-	settleResolution, ok := resolution.(*HtlcSettleResolution)
-	if !ok {
-		t.Fatalf("expected settle resolution, got: %T",
-			resolution)
+	checkResolution := func(res HtlcResolution, pimg lntypes.Preimage) {
+		// Otherwise we expect no error and a settle res for the htlc.
+		settleResolution, ok := res.(*HtlcSettleResolution)
+		assert.True(t, ok)
+		assert.Equal(t, settleResolution.Preimage, pimg)
 	}
-	if settleResolution.Preimage != preimage {
-		t.Fatalf("expected settle with matching preimage")
+	checkSubscription := func() {
+		// We expect a new invoice notification to be sent out.
+		newInvoice := <-allSubscriptions.NewInvoices
+		assert.Equal(t, newInvoice.State, channeldb.ContractOpen)
+
+		// We expect a settled notification to be sent out.
+		settledInvoice := <-allSubscriptions.SettledInvoices
+		assert.Equal(t, settledInvoice.State, channeldb.ContractSettled)
 	}
 
-	// We expect a new invoice notification to be sent out.
-	newInvoice := <-allSubscriptions.NewInvoices
-	if newInvoice.State != channeldb.ContractOpen {
-		t.Fatalf("expected state ContractOpen, but got %v",
-			newInvoice.State)
+	checkResolution(resolution, preimage)
+	checkSubscription()
+
+	// Replay the same keysend payment. We expect an identical resolution,
+	// but no event should be generated.
+	resolution, err = ctx.registry.NotifyExitHopHtlc(
+		hash, amt, expiry,
+		testCurrentHeight, getCircuitKey(10), hodlChan, keySendPayload,
+	)
+	assert.Nil(t, err)
+	checkResolution(resolution, preimage)
+
+	select {
+	case <-allSubscriptions.NewInvoices:
+		t.Fatalf("replayed keysend should not generate event")
+	case <-time.After(time.Second):
 	}
 
-	// We expect a settled notification to be sent out.
-	settledInvoice := <-allSubscriptions.SettledInvoices
-	if settledInvoice.State != channeldb.ContractSettled {
-		t.Fatalf("expected state ContractOpen, but got %v",
-			settledInvoice.State)
+	// Finally, test that we can properly fulfill a second keysend payment
+	// with a unique preiamge.
+	preimage2 := lntypes.Preimage{1, 2, 3, 4}
+	hash2 := preimage2.Hash()
+
+	keySendPayload2 := &mockPayload{
+		customRecords: map[uint64][]byte{
+			record.KeySendType: preimage2[:],
+		},
 	}
+
+	resolution, err = ctx.registry.NotifyExitHopHtlc(
+		hash2, amt, expiry,
+		testCurrentHeight, getCircuitKey(20), hodlChan, keySendPayload2,
+	)
+	assert.Nil(t, err)
+
+	checkResolution(resolution, preimage2)
+	checkSubscription()
 }
 
 // TestMppPayment tests settling of an invoice with multiple partial payments.