From f1b7953465421b52a05efd25653794b5b060bd18 Mon Sep 17 00:00:00 2001
From: Alex Akselrod <alex@akselrod.org>
Date: Wed, 21 Aug 2024 17:33:37 -0700
Subject: [PATCH] invoices/sqldb: query by ChanID when updating AMP invoice
 preimage

---
 .golangci.yml                       | 2 ++
 go.mod                              | 3 +++
 go.sum                              | 2 --
 invoices/sql_store.go               | 3 +++
 sqldb/sqlc/amp_invoices.sql.go      | 6 ++++--
 sqldb/sqlc/queries/amp_invoices.sql | 4 ++--
 6 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/.golangci.yml b/.golangci.yml
index 536c1b655..8114945c6 100644
--- a/.golangci.yml
+++ b/.golangci.yml
@@ -93,11 +93,13 @@ linters-settings:
       - 'errors.Wrap'
   
   gomoddirectives:
+    replace-local: true
     replace-allow-list:
       # See go.mod for the explanation why these are needed.
       - github.com/ulikunitz/xz
       - github.com/gogo/protobuf
       - google.golang.org/protobuf
+      - github.com/lightningnetwork/lnd/sqldb
 
 
 linters:
diff --git a/go.mod b/go.mod
index bbe3c8845..3e5b9e8f2 100644
--- a/go.mod
+++ b/go.mod
@@ -204,6 +204,9 @@ replace github.com/gogo/protobuf => github.com/gogo/protobuf v1.3.2
 // allows us to specify that as an option.
 replace google.golang.org/protobuf => github.com/lightninglabs/protobuf-go-hex-display v1.30.0-hex-display
 
+// Temporary replace until the next version of sqldb is taged.
+replace github.com/lightningnetwork/lnd/sqldb => ./sqldb
+
 // If you change this please also update .github/pull_request_template.md,
 // docs/INSTALL.md and GO_IMAGE in lnrpc/gen_protos_docker.sh.
 go 1.22.6
diff --git a/go.sum b/go.sum
index a51e291b8..7dd8449cd 100644
--- a/go.sum
+++ b/go.sum
@@ -458,8 +458,6 @@ github.com/lightningnetwork/lnd/kvdb v1.4.10 h1:vK89IVv1oVH9ubQWU+EmoCQFeVRaC8kf
 github.com/lightningnetwork/lnd/kvdb v1.4.10/go.mod h1:J2diNABOoII9UrMnxXS5w7vZwP7CA1CStrl8MnIrb3A=
 github.com/lightningnetwork/lnd/queue v1.1.1 h1:99ovBlpM9B0FRCGYJo6RSFDlt8/vOkQQZznVb18iNMI=
 github.com/lightningnetwork/lnd/queue v1.1.1/go.mod h1:7A6nC1Qrm32FHuhx/mi1cieAiBZo5O6l8IBIoQxvkz4=
-github.com/lightningnetwork/lnd/sqldb v1.0.3 h1:zLfAwOvM+6+3+hahYO9Q3h8pVV0TghAR7iJ5YMLCd3I=
-github.com/lightningnetwork/lnd/sqldb v1.0.3/go.mod h1:4cQOkdymlZ1znnjuRNvMoatQGJkRneTj2CoPSPaQhWo=
 github.com/lightningnetwork/lnd/ticker v1.1.1 h1:J/b6N2hibFtC7JLV77ULQp++QLtCwT6ijJlbdiZFbSM=
 github.com/lightningnetwork/lnd/ticker v1.1.1/go.mod h1:waPTRAAcwtu7Ji3+3k+u/xH5GHovTsCoSVpho0KDvdA=
 github.com/lightningnetwork/lnd/tlv v1.2.6 h1:icvQG2yDr6k3ZuZzfRdG3EJp6pHurcuh3R6dg0gv/Mw=
diff --git a/invoices/sql_store.go b/invoices/sql_store.go
index 4b488715b..c6a58c77b 100644
--- a/invoices/sql_store.go
+++ b/invoices/sql_store.go
@@ -1116,6 +1116,9 @@ func (s *sqlInvoiceUpdater) AddAmpHtlcPreimage(setID [32]byte,
 			SetID:     setID[:],
 			HtlcID:    int64(circuitKey.HtlcID),
 			Preimage:  preimage[:],
+			ChanID: strconv.FormatUint(
+				circuitKey.ChanID.ToUint64(), 10,
+			),
 		},
 	)
 	if err != nil {
diff --git a/sqldb/sqlc/amp_invoices.sql.go b/sqldb/sqlc/amp_invoices.sql.go
index 3fcfe4b27..e47b1c803 100644
--- a/sqldb/sqlc/amp_invoices.sql.go
+++ b/sqldb/sqlc/amp_invoices.sql.go
@@ -268,15 +268,16 @@ func (q *Queries) InsertAMPSubInvoiceHTLC(ctx context.Context, arg InsertAMPSubI
 
 const updateAMPSubInvoiceHTLCPreimage = `-- name: UpdateAMPSubInvoiceHTLCPreimage :execresult
 UPDATE amp_sub_invoice_htlcs AS a
-SET preimage = $4
+SET preimage = $5
 WHERE a.invoice_id = $1 AND a.set_id = $2 AND a.htlc_id = (
-    SELECT id FROM invoice_htlcs AS i WHERE i.htlc_id = $3
+    SELECT id FROM invoice_htlcs AS i WHERE i.chan_id = $3 AND i.htlc_id = $4
 )
 `
 
 type UpdateAMPSubInvoiceHTLCPreimageParams struct {
 	InvoiceID int64
 	SetID     []byte
+	ChanID    string
 	HtlcID    int64
 	Preimage  []byte
 }
@@ -285,6 +286,7 @@ func (q *Queries) UpdateAMPSubInvoiceHTLCPreimage(ctx context.Context, arg Updat
 	return q.db.ExecContext(ctx, updateAMPSubInvoiceHTLCPreimage,
 		arg.InvoiceID,
 		arg.SetID,
+		arg.ChanID,
 		arg.HtlcID,
 		arg.Preimage,
 	)
diff --git a/sqldb/sqlc/queries/amp_invoices.sql b/sqldb/sqlc/queries/amp_invoices.sql
index 3b6ee76ac..1fad75e0d 100644
--- a/sqldb/sqlc/queries/amp_invoices.sql
+++ b/sqldb/sqlc/queries/amp_invoices.sql
@@ -61,7 +61,7 @@ WHERE (
 
 -- name: UpdateAMPSubInvoiceHTLCPreimage :execresult
 UPDATE amp_sub_invoice_htlcs AS a
-SET preimage = $4
+SET preimage = $5
 WHERE a.invoice_id = $1 AND a.set_id = $2 AND a.htlc_id = (
-    SELECT id FROM invoice_htlcs AS i WHERE i.htlc_id = $3
+    SELECT id FROM invoice_htlcs AS i WHERE i.chan_id = $3 AND i.htlc_id = $4
 );