From 0209e6feb8246b5a61776580c24f5f9541dbbb39 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Thu, 9 Feb 2023 17:05:23 +0200
Subject: [PATCH] watchtower/wtclient: add Copy method to AddressIterator

---
 watchtower/wtclient/addr_iterator.go          | 33 +++++++++++++++++++
 watchtower/wtclient/addr_iterator_test.go     |  5 +++
 .../wtclient/candidate_iterator_test.go       |  6 +---
 3 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/watchtower/wtclient/addr_iterator.go b/watchtower/wtclient/addr_iterator.go
index 87065c011..cb16d335a 100644
--- a/watchtower/wtclient/addr_iterator.go
+++ b/watchtower/wtclient/addr_iterator.go
@@ -69,6 +69,12 @@ type AddressIterator interface {
 	// Reset clears the iterators state, and makes the address at the front
 	// of the list the next item to be returned.
 	Reset()
+
+	// Copy constructs a new AddressIterator that has the same addresses
+	// as this iterator.
+	//
+	// NOTE that the address locks are not expected to be copied.
+	Copy() AddressIterator
 }
 
 // A compile-time check to ensure that addressIterator implements the
@@ -324,6 +330,33 @@ func (a *addressIterator) GetAll() []net.Addr {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
+	return a.getAllUnsafe()
+}
+
+// Copy constructs a new AddressIterator that has the same addresses
+// as this iterator.
+//
+// NOTE that the address locks will not be copied.
+func (a *addressIterator) Copy() AddressIterator {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	addrs := a.getAllUnsafe()
+
+	// Since newAddressIterator will only ever return an error if it is
+	// initialised with zero addresses, we can ignore the error here since
+	// we are initialising it with the set of addresses of this
+	// addressIterator which is by definition a non-empty list.
+	iter, _ := newAddressIterator(addrs...)
+
+	return iter
+}
+
+// getAllUnsafe returns a copy of all the addresses in the iterator.
+//
+// NOTE: this method is not thread safe and so must only be called once the
+// addressIterator mutex is already being held.
+func (a *addressIterator) getAllUnsafe() []net.Addr {
 	var addrs []net.Addr
 	cursor := a.addrList.Front()
 
diff --git a/watchtower/wtclient/addr_iterator_test.go b/watchtower/wtclient/addr_iterator_test.go
index d3674d985..89a35c4cb 100644
--- a/watchtower/wtclient/addr_iterator_test.go
+++ b/watchtower/wtclient/addr_iterator_test.go
@@ -97,6 +97,11 @@ func TestAddrIterator(t *testing.T) {
 	addrList := iter.GetAll()
 	require.ElementsMatch(t, addrList, []net.Addr{addr1, addr2, addr3})
 
+	// Also check that an iterator constructed via the Copy method, also
+	// contains all the expected addresses.
+	newIterAddrs := iter.Copy().GetAll()
+	require.ElementsMatch(t, newIterAddrs, []net.Addr{addr1, addr2, addr3})
+
 	// Let's now remove addr3.
 	err = iter.Remove(addr3)
 	require.NoError(t, err)
diff --git a/watchtower/wtclient/candidate_iterator_test.go b/watchtower/wtclient/candidate_iterator_test.go
index 70dfb7505..b4df80f4d 100644
--- a/watchtower/wtclient/candidate_iterator_test.go
+++ b/watchtower/wtclient/candidate_iterator_test.go
@@ -52,14 +52,10 @@ func randTower(t *testing.T) *Tower {
 func copyTower(t *testing.T, tower *Tower) *Tower {
 	t.Helper()
 
-	addrs := tower.Addresses.GetAll()
-	addrIterator, err := newAddressIterator(addrs...)
-	require.NoError(t, err)
-
 	return &Tower{
 		ID:          tower.ID,
 		IdentityKey: tower.IdentityKey,
-		Addresses:   addrIterator,
+		Addresses:   tower.Addresses.Copy(),
 	}
 }