watchtower/wtclient: add Copy method to AddressIterator

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()
 

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)

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(),
 	}
 }