From bfe6262b2960442cc0e3cf6f4379571e6ba006a1 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Tue, 12 Nov 2024 08:11:46 +0200
Subject: [PATCH] graph+channeldb: add AddrSource interface to GraphSource

---
 chanbackup/backup.go          |  7 ++++++-
 chanbackup/backup_test.go     |  5 +++--
 channel_notifier.go           |  5 ++++-
 channeldb/addr_source.go      | 10 ++++++----
 channeldb/addr_source_test.go | 30 +++++++++++++++++-------------
 channeldb/db.go               |  5 ++++-
 channeldb/db_test.go          |  5 ++++-
 graph/db/graph.go             |  5 +++--
 graph/sources/chan_graph.go   | 13 +++++++++++++
 graph/sources/interfaces.go   |  2 ++
 server.go                     |  2 +-
 11 files changed, 63 insertions(+), 26 deletions(-)

diff --git a/chanbackup/backup.go b/chanbackup/backup.go
index 5853b37e4..8248e547c 100644
--- a/chanbackup/backup.go
+++ b/chanbackup/backup.go
@@ -1,6 +1,7 @@
 package chanbackup
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/btcsuite/btcd/wire"
@@ -27,12 +28,16 @@ type LiveChannelSource interface {
 func assembleChanBackup(addrSource channeldb.AddrSource,
 	openChan *channeldb.OpenChannel) (*Single, error) {
 
+	ctx := context.TODO()
+
 	log.Debugf("Crafting backup for ChannelPoint(%v)",
 		openChan.FundingOutpoint)
 
 	// First, we'll query the channel source to obtain all the addresses
 	// that are associated with the peer for this channel.
-	known, nodeAddrs, err := addrSource.AddrsForNode(openChan.IdentityPub)
+	known, nodeAddrs, err := addrSource.AddrsForNode(
+		ctx, openChan.IdentityPub,
+	)
 	if err != nil {
 		return nil, err
 	}
diff --git a/chanbackup/backup_test.go b/chanbackup/backup_test.go
index 46ccf4c24..c0d80c2ef 100644
--- a/chanbackup/backup_test.go
+++ b/chanbackup/backup_test.go
@@ -1,6 +1,7 @@
 package chanbackup
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"testing"
@@ -61,8 +62,8 @@ func (m *mockChannelSource) addAddrsForNode(nodePub *btcec.PublicKey, addrs []ne
 	m.addrs[nodeKey] = addrs
 }
 
-func (m *mockChannelSource) AddrsForNode(nodePub *btcec.PublicKey) (bool,
-	[]net.Addr, error) {
+func (m *mockChannelSource) AddrsForNode(_ context.Context,
+	nodePub *btcec.PublicKey) (bool, []net.Addr, error) {
 
 	if m.failQuery {
 		return false, nil, fmt.Errorf("fail")
diff --git a/channel_notifier.go b/channel_notifier.go
index 88a05ac4c..62d02e5ac 100644
--- a/channel_notifier.go
+++ b/channel_notifier.go
@@ -1,6 +1,7 @@
 package lnd
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/btcsuite/btcd/wire"
@@ -34,6 +35,8 @@ type channelNotifier struct {
 func (c *channelNotifier) SubscribeChans(startingChans map[wire.OutPoint]struct{}) (
 	*chanbackup.ChannelSubscription, error) {
 
+	ctx := context.TODO()
+
 	ltndLog.Infof("Channel backup proxy channel notifier starting")
 
 	// TODO(roasbeef): read existing set of chans and diff
@@ -46,7 +49,7 @@ func (c *channelNotifier) SubscribeChans(startingChans map[wire.OutPoint]struct{
 	// confirmed channels.
 	sendChanOpenUpdate := func(newOrPendingChan *channeldb.OpenChannel) {
 		_, nodeAddrs, err := c.addrs.AddrsForNode(
-			newOrPendingChan.IdentityPub,
+			ctx, newOrPendingChan.IdentityPub,
 		)
 		if err != nil {
 			pub := newOrPendingChan.IdentityPub
diff --git a/channeldb/addr_source.go b/channeldb/addr_source.go
index de933ed49..f6b176090 100644
--- a/channeldb/addr_source.go
+++ b/channeldb/addr_source.go
@@ -1,6 +1,7 @@
 package channeldb
 
 import (
+	"context"
 	"errors"
 	"net"
 
@@ -13,7 +14,8 @@ type AddrSource interface {
 	// AddrsForNode returns all known addresses for the target node public
 	// key. The returned boolean must indicate if the given node is unknown
 	// to the backing source.
-	AddrsForNode(nodePub *btcec.PublicKey) (bool, []net.Addr, error)
+	AddrsForNode(ctx context.Context, nodePub *btcec.PublicKey) (bool,
+		[]net.Addr, error)
 }
 
 // multiAddrSource is an implementation of AddrSource which gathers all the
@@ -38,8 +40,8 @@ func NewMultiAddrSource(sources ...AddrSource) AddrSource {
 // node.
 //
 // NOTE: this implements the AddrSource interface.
-func (c *multiAddrSource) AddrsForNode(nodePub *btcec.PublicKey) (bool,
-	[]net.Addr, error) {
+func (c *multiAddrSource) AddrsForNode(ctx context.Context,
+	nodePub *btcec.PublicKey) (bool, []net.Addr, error) {
 
 	if len(c.sources) == 0 {
 		return false, nil, errors.New("no address sources")
@@ -55,7 +57,7 @@ func (c *multiAddrSource) AddrsForNode(nodePub *btcec.PublicKey) (bool,
 	// Iterate over all the address sources and query each one for the
 	// addresses it has for the node in question.
 	for _, src := range c.sources {
-		isKnown, addrs, err := src.AddrsForNode(nodePub)
+		isKnown, addrs, err := src.AddrsForNode(ctx, nodePub)
 		if err != nil {
 			return false, nil, err
 		}
diff --git a/channeldb/addr_source_test.go b/channeldb/addr_source_test.go
index 85ee30bf5..19f46aebe 100644
--- a/channeldb/addr_source_test.go
+++ b/channeldb/addr_source_test.go
@@ -1,6 +1,7 @@
 package channeldb
 
 import (
+	"context"
 	"net"
 	"testing"
 
@@ -20,7 +21,10 @@ var (
 func TestMultiAddrSource(t *testing.T) {
 	t.Parallel()
 
-	var pk1 = newTestPubKey(t)
+	var (
+		ctx = context.Background()
+		pk1 = newTestPubKey(t)
+	)
 
 	t.Run("both sources have results", func(t *testing.T) {
 		t.Parallel()
@@ -35,12 +39,12 @@ func TestMultiAddrSource(t *testing.T) {
 		})
 
 		// Let source 1 know of 2 addresses (addr 1 and 2) for node 1.
-		src1.On("AddrsForNode", pk1).Return(
+		src1.On("AddrsForNode", ctx, pk1).Return(
 			true, []net.Addr{addr1, addr2}, nil,
 		).Once()
 
 		// Let source 2 know of 2 addresses (addr 2 and 3) for node 1.
-		src2.On("AddrsForNode", pk1).Return(
+		src2.On("AddrsForNode", ctx, pk1).Return(
 			true, []net.Addr{addr2, addr3}, nil,
 			[]net.Addr{addr2, addr3}, nil,
 		).Once()
@@ -51,7 +55,7 @@ func TestMultiAddrSource(t *testing.T) {
 
 		// Query it for the addresses known for node 1. The results
 		// should contain addr 1, 2 and 3.
-		known, addrs, err := multiSrc.AddrsForNode(pk1)
+		known, addrs, err := multiSrc.AddrsForNode(ctx, pk1)
 		require.NoError(t, err)
 		require.True(t, known)
 		require.ElementsMatch(t, addrs, []net.Addr{addr1, addr2, addr3})
@@ -70,10 +74,10 @@ func TestMultiAddrSource(t *testing.T) {
 		})
 
 		// Let source 1 know of address 1 for node 1.
-		src1.On("AddrsForNode", pk1).Return(
+		src1.On("AddrsForNode", ctx, pk1).Return(
 			true, []net.Addr{addr1}, nil,
 		).Once()
-		src2.On("AddrsForNode", pk1).Return(false, nil, nil).Once()
+		src2.On("AddrsForNode", ctx, pk1).Return(false, nil, nil).Once()
 
 		// Create a multi-addr source that consists of both source 1
 		// and 2.
@@ -81,7 +85,7 @@ func TestMultiAddrSource(t *testing.T) {
 
 		// Query it for the addresses known for node 1. The results
 		// should contain addr 1.
-		known, addrs, err := multiSrc.AddrsForNode(pk1)
+		known, addrs, err := multiSrc.AddrsForNode(ctx, pk1)
 		require.NoError(t, err)
 		require.True(t, known)
 		require.ElementsMatch(t, addrs, []net.Addr{addr1})
@@ -103,13 +107,13 @@ func TestMultiAddrSource(t *testing.T) {
 		// and 2. Neither source known of node 1.
 		multiSrc := NewMultiAddrSource(src1, src2)
 
-		src1.On("AddrsForNode", pk1).Return(false, nil, nil).Once()
-		src2.On("AddrsForNode", pk1).Return(false, nil, nil).Once()
+		src1.On("AddrsForNode", ctx, pk1).Return(false, nil, nil).Once()
+		src2.On("AddrsForNode", ctx, pk1).Return(false, nil, nil).Once()
 
 		// Query it for the addresses known for node 1. It should return
 		// false to indicate that the node is unknown to all backing
 		// sources.
-		known, addrs, err := multiSrc.AddrsForNode(pk1)
+		known, addrs, err := multiSrc.AddrsForNode(ctx, pk1)
 		require.NoError(t, err)
 		require.False(t, known)
 		require.Empty(t, addrs)
@@ -127,10 +131,10 @@ func newMockAddrSource(t *testing.T) *mockAddrSource {
 	return &mockAddrSource{t: t}
 }
 
-func (m *mockAddrSource) AddrsForNode(pub *btcec.PublicKey) (bool, []net.Addr,
-	error) {
+func (m *mockAddrSource) AddrsForNode(ctx context.Context,
+	pub *btcec.PublicKey) (bool, []net.Addr, error) {
 
-	args := m.Called(pub)
+	args := m.Called(ctx, pub)
 	if args.Get(1) == nil {
 		return args.Bool(0), nil, args.Error(2)
 	}
diff --git a/channeldb/db.go b/channeldb/db.go
index bf7909ba5..f884cd4a6 100644
--- a/channeldb/db.go
+++ b/channeldb/db.go
@@ -2,6 +2,7 @@ package channeldb
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"fmt"
 	"net"
@@ -1344,7 +1345,9 @@ func (c *ChannelStateDB) RestoreChannelShells(channelShells ...*ChannelShell) er
 // unknown to the channel DB or not.
 //
 // NOTE: this is part of the AddrSource interface.
-func (d *DB) AddrsForNode(nodePub *btcec.PublicKey) (bool, []net.Addr, error) {
+func (d *DB) AddrsForNode(_ context.Context, nodePub *btcec.PublicKey) (bool,
+	[]net.Addr, error) {
+
 	linkNode, err := d.channelStateDB.linkNodeDB.FetchLinkNode(nodePub)
 	// Only if the error is something other than ErrNodeNotFound do we
 	// return it.
diff --git a/channeldb/db_test.go b/channeldb/db_test.go
index 9fed9934b..740e29338 100644
--- a/channeldb/db_test.go
+++ b/channeldb/db_test.go
@@ -1,6 +1,7 @@
 package channeldb
 
 import (
+	"context"
 	"image/color"
 	"math"
 	"math/rand"
@@ -182,6 +183,8 @@ func TestFetchClosedChannelForID(t *testing.T) {
 func TestMultiSourceAddrsForNode(t *testing.T) {
 	t.Parallel()
 
+	ctx := context.Background()
+
 	fullDB, err := MakeTestDB(t)
 	require.NoError(t, err, "unable to make test database")
 
@@ -212,7 +215,7 @@ func TestMultiSourceAddrsForNode(t *testing.T) {
 
 	// Now that we've created a link node, as well as a vertex for the
 	// node, we'll query for all its addresses.
-	known, nodeAddrs, err := addrSource.AddrsForNode(nodePub)
+	known, nodeAddrs, err := addrSource.AddrsForNode(ctx, nodePub)
 	require.NoError(t, err, "unable to obtain node addrs")
 	require.True(t, known)
 
diff --git a/graph/db/graph.go b/graph/db/graph.go
index dcad56f3b..9f240de98 100644
--- a/graph/db/graph.go
+++ b/graph/db/graph.go
@@ -2,6 +2,7 @@ package graphdb
 
 import (
 	"bytes"
+	"context"
 	"crypto/sha256"
 	"encoding/binary"
 	"errors"
@@ -418,8 +419,8 @@ func (c *ChannelGraph) NewPathFindTx() (kvdb.RTx, error) {
 // unknown to the graph DB or not.
 //
 // NOTE: this is part of the channeldb.AddrSource interface.
-func (c *ChannelGraph) AddrsForNode(nodePub *btcec.PublicKey) (bool, []net.Addr,
-	error) {
+func (c *ChannelGraph) AddrsForNode(_ context.Context,
+	nodePub *btcec.PublicKey) (bool, []net.Addr, error) {
 
 	pubKey, err := route.NewVertexFromBytes(nodePub.SerializeCompressed())
 	if err != nil {
diff --git a/graph/sources/chan_graph.go b/graph/sources/chan_graph.go
index 57d7c2098..03910f792 100644
--- a/graph/sources/chan_graph.go
+++ b/graph/sources/chan_graph.go
@@ -3,7 +3,9 @@ package sources
 import (
 	"context"
 	"fmt"
+	"net"
 
+	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/wire"
 	graphdb "github.com/lightningnetwork/lnd/graph/db"
 	"github.com/lightningnetwork/lnd/graph/db/models"
@@ -117,6 +119,17 @@ func (s *DBSource) FetchChannelEdgesByOutpoint(_ context.Context,
 	return s.db.FetchChannelEdgesByOutpoint(point)
 }
 
+// AddrsForNode returns all known addresses for the target node public key. The
+// returned boolean indicatex if the given node is unknown to the backing
+// source.
+//
+// NOTE: this is part of the channeldb.AddrSource interface.
+func (s *DBSource) AddrsForNode(ctx context.Context,
+	nodePub *btcec.PublicKey) (bool, []net.Addr, error) {
+
+	return s.db.AddrsForNode(ctx, nodePub)
+}
+
 // kvdbRTx is an implementation of graphdb.RTx backed by a KVDB database read
 // transaction.
 type kvdbRTx struct {
diff --git a/graph/sources/interfaces.go b/graph/sources/interfaces.go
index 0753eec3b..df860a7e3 100644
--- a/graph/sources/interfaces.go
+++ b/graph/sources/interfaces.go
@@ -1,6 +1,7 @@
 package sources
 
 import (
+	"github.com/lightningnetwork/lnd/channeldb"
 	"github.com/lightningnetwork/lnd/graph/session"
 	"github.com/lightningnetwork/lnd/lnrpc/invoicesrpc"
 	"github.com/lightningnetwork/lnd/netann"
@@ -12,4 +13,5 @@ type GraphSource interface {
 	session.ReadOnlyGraph
 	invoicesrpc.GraphSource
 	netann.ChannelGraph
+	channeldb.AddrSource
 }
diff --git a/server.go b/server.go
index 5d5a45d33..796f5a7fb 100644
--- a/server.go
+++ b/server.go
@@ -612,7 +612,7 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 		HtlcInterceptor:             invoiceHtlcModifier,
 	}
 
-	addrSource := channeldb.NewMultiAddrSource(dbs.ChanStateDB, dbs.GraphDB)
+	addrSource := channeldb.NewMultiAddrSource(dbs.ChanStateDB, graphSource)
 
 	s := &server{
 		cfg:            cfg,