From 9537026df96db604f9dd955a4b199d152bd07bd2 Mon Sep 17 00:00:00 2001 From: Elle Mouton Date: Tue, 22 Oct 2024 14:15:35 +0200 Subject: [PATCH] channeldb: implement a multi-source AddrSource In this commit, we implement a version of the AddrSource interface which merges the results of a set of AddrSource implementations. This will later be used to merge the results of the channel db and graph db. --- channeldb/addr_source.go | 53 +++++++++++++++ channeldb/addr_source_test.go | 119 ++++++++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+) create mode 100644 channeldb/addr_source_test.go diff --git a/channeldb/addr_source.go b/channeldb/addr_source.go index 7fc0f8134..ef06fbe1f 100644 --- a/channeldb/addr_source.go +++ b/channeldb/addr_source.go @@ -13,3 +13,56 @@ type AddrSource interface { // key. AddrsForNode(nodePub *btcec.PublicKey) ([]net.Addr, error) } + +// multiAddrSource is an implementation of AddrSource which gathers all the +// known addresses for a given node from multiple backends and de-duplicates the +// results. +type multiAddrSource struct { + sources []AddrSource +} + +// NewMultiAddrSource constructs a new AddrSource which will query all the +// provided sources for a node's addresses and will then de-duplicate the +// results. +func NewMultiAddrSource(sources ...AddrSource) AddrSource { + return &multiAddrSource{ + sources: sources, + } +} + +// AddrsForNode returns all known addresses for the target node public key. It +// queries all the address sources provided and de-duplicates the results. +// +// NOTE: this implements the AddrSource interface. +func (c *multiAddrSource) AddrsForNode(nodePub *btcec.PublicKey) ([]net.Addr, + error) { + + // The multiple address sources will likely contain duplicate addresses, + // so we use a map here to de-dup them. + dedupedAddrs := make(map[string]net.Addr) + + // 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 { + addrs, err := src.AddrsForNode(nodePub) + if err != nil { + return nil, err + } + + for _, addr := range addrs { + dedupedAddrs[addr.String()] = addr + } + } + + // Convert the map into a list we can return. + addrs := make([]net.Addr, 0, len(dedupedAddrs)) + for _, addr := range dedupedAddrs { + addrs = append(addrs, addr) + } + + return addrs, nil +} + +// A compile-time check to ensure that multiAddrSource implements the AddrSource +// interface. +var _ AddrSource = (*multiAddrSource)(nil) diff --git a/channeldb/addr_source_test.go b/channeldb/addr_source_test.go new file mode 100644 index 000000000..3d7d71ac9 --- /dev/null +++ b/channeldb/addr_source_test.go @@ -0,0 +1,119 @@ +package channeldb + +import ( + "net" + "testing" + + "github.com/btcsuite/btcd/btcec/v2" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +var ( + addr1 = &net.TCPAddr{IP: (net.IP)([]byte{0x1}), Port: 1} + addr2 = &net.TCPAddr{IP: (net.IP)([]byte{0x2}), Port: 2} + addr3 = &net.TCPAddr{IP: (net.IP)([]byte{0x3}), Port: 3} +) + +// TestMultiAddrSource tests that the multiAddrSource correctly merges and +// deduplicates the results of a set of AddrSource implementations. +func TestMultiAddrSource(t *testing.T) { + t.Parallel() + + var pk1 = newTestPubKey(t) + + t.Run("both sources have results", func(t *testing.T) { + t.Parallel() + + var ( + src1 = newMockAddrSource(t) + src2 = newMockAddrSource(t) + ) + t.Cleanup(func() { + src1.AssertExpectations(t) + src2.AssertExpectations(t) + }) + + // Let source 1 know of 2 addresses (addr 1 and 2) for node 1. + src1.On("AddrsForNode", pk1).Return( + []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( + []net.Addr{addr2, addr3}, nil, + ).Once() + + // Create a multi-addr source that consists of both source 1 + // and 2. + multiSrc := NewMultiAddrSource(src1, src2) + + // Query it for the addresses known for node 1. The results + // should contain addr 1, 2 and 3. + addrs, err := multiSrc.AddrsForNode(pk1) + require.NoError(t, err) + require.ElementsMatch(t, addrs, []net.Addr{addr1, addr2, addr3}) + }) + + t.Run("only once source has results", func(t *testing.T) { + t.Parallel() + + var ( + src1 = newMockAddrSource(t) + src2 = newMockAddrSource(t) + ) + t.Cleanup(func() { + src1.AssertExpectations(t) + src2.AssertExpectations(t) + }) + + // Let source 1 know of address 1 for node 1. + src1.On("AddrsForNode", pk1).Return( + []net.Addr{addr1}, nil, + ).Once() + src2.On("AddrsForNode", pk1).Return(nil, nil).Once() + + // Create a multi-addr source that consists of both source 1 + // and 2. + multiSrc := NewMultiAddrSource(src1, src2) + + // Query it for the addresses known for node 1. The results + // should contain addr 1. + addrs, err := multiSrc.AddrsForNode(pk1) + require.NoError(t, err) + require.ElementsMatch(t, addrs, []net.Addr{addr1}) + }) + +} + +type mockAddrSource struct { + t *testing.T + mock.Mock +} + +var _ AddrSource = (*mockAddrSource)(nil) + +func newMockAddrSource(t *testing.T) *mockAddrSource { + return &mockAddrSource{t: t} +} + +func (m *mockAddrSource) AddrsForNode(pub *btcec.PublicKey) ([]net.Addr, + error) { + + args := m.Called(pub) + if args.Get(0) == nil { + return nil, args.Error(1) + } + + addrs, ok := args.Get(0).([]net.Addr) + require.True(m.t, ok) + + return addrs, args.Error(1) +} + +func newTestPubKey(t *testing.T) *btcec.PublicKey { + priv, err := btcec.NewPrivateKey() + require.NoError(t, err) + + return priv.PubKey() +}