From be6698447e546d1a19be1721ea0da0b9591b4ae1 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Wed, 24 Mar 2021 19:48:23 -0700
Subject: [PATCH] channeldb/invoice: add InvoiceRefByAddr

---
 channeldb/invoice_test.go |  7 +++++++
 channeldb/invoices.go     | 12 ++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/channeldb/invoice_test.go b/channeldb/invoice_test.go
index 3b0416d1..7132ba27 100644
--- a/channeldb/invoice_test.go
+++ b/channeldb/invoice_test.go
@@ -1225,6 +1225,13 @@ func TestInvoiceRef(t *testing.T) {
 	require.Equal(t, lntypes.Hash{}, refBySetID.PayHash())
 	require.Equal(t, (*[32]byte)(nil), refBySetID.PayAddr())
 	require.Equal(t, &setID, refBySetID.SetID())
+
+	// An InvoiceRef by pay addr should only return a pay addr, but nil for
+	// pay hash and set id.
+	refByAddr := InvoiceRefByAddr(payAddr)
+	require.Equal(t, (*lntypes.Hash)(nil), refByAddr.PayHash())
+	require.Equal(t, &payAddr, refByAddr.PayAddr())
+	require.Equal(t, (*[32]byte)(nil), refByAddr.SetID())
 }
 
 // TestHTLCSet asserts that HTLCSet returns the proper set of accepted HTLCs
diff --git a/channeldb/invoices.go b/channeldb/invoices.go
index 7598575a..444602ae 100644
--- a/channeldb/invoices.go
+++ b/channeldb/invoices.go
@@ -236,6 +236,14 @@ func InvoiceRefByHashAndAddr(payHash lntypes.Hash,
 	}
 }
 
+// InvoiceRefByAddr creates an InvoiceRef that queries the payment addr index
+// for an invoice with the provided payment address.
+func InvoiceRefByAddr(addr [32]byte) InvoiceRef {
+	return InvoiceRef{
+		payAddr: &addr,
+	}
+}
+
 // InvoiceRefBySetID creates an InvoiceRef that queries the set id index for an
 // invoice with the provided setID. If the invoice is not found, the query will
 // not fallback to payHash or payAddr.
@@ -903,6 +911,10 @@ func fetchInvoiceNumByRef(invoiceIndex, payAddrIndex, setIDIndex kvdb.RBucket,
 
 		return invoiceNumByAddr, nil
 
+	// Return invoices by payment addr only.
+	case invoiceNumByAddr != nil:
+		return invoiceNumByAddr, nil
+
 	// If we were only able to reference the invoice by hash, return the
 	// corresponding invoice number. This can happen when no payment address
 	// was provided, or if it didn't match anything in our records.