Merge pull request #7556 from ellemouton/addrIterFix

wtclient: address iterator fix

docs/release-notes/release-notes-0.16.1.md

@@ -14,6 +14,8 @@
 
 ## Watchtowers
 
+* [Fix Address iterator 
+  panic](https://github.com/lightningnetwork/lnd/pull/7556)
 * [Allow caller to filter sessions at the time of reading them from 
   disk](https://github.com/lightningnetwork/lnd/pull/7059)
 * [Clean up sessions once all channels for which they have updates for are

watchtower/wtclient/addr_iterator.go

@@ -162,7 +162,15 @@ func (a *addressIterator) next(lock bool) (net.Addr, error) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
-	// Set the next candidate to the subsequent element.
+	// In-case currentTopAddr is nil (meaning that Reset has not yet been
+	// called), return an error indicating this.
+	if a.currentTopAddr == nil {
+		return nil, ErrAddressesExhausted
+	}
+
+	// Set the next candidate to the subsequent element. If we are at the
+	// end of the address list, this could mean setting currentTopAddr to
+	// nil.
 	a.currentTopAddr = a.currentTopAddr.Next()
 
 	for a.currentTopAddr != nil {
@@ -172,9 +180,26 @@ func (a *addressIterator) next(lock bool) (net.Addr, error) {
 		// Check whether this address is still considered a candidate.
 		// If it's not, we'll proceed to the next.
 		candidate, ok := a.candidates[addrID]
+
+		// If the address cannot be found in the candidate set, then
+		// this must mean that the Remove method was called for the
+		// address. The Remove method would have checked that the
+		// address is not the last one in the iterator and that it has
+		// no locks on it. It is therefor safe to remove.
 		if !ok {
+			// Grab the next address candidate. This might be nil
+			// if the iterator is on the last item in the list.
 			nextCandidate := a.currentTopAddr.Next()
+
+			// Remove the address from the list that is no longer
+			// in the candidate set.
 			a.addrList.Remove(a.currentTopAddr)
+
+			// Set the current top to the next candidate. This might
+			// mean setting it to nil if the iterator is on its last
+			// item in which case the loop will be exited and an
+			// ErrAddressesExhausted exhausted error will be
+			// returned.
 			a.currentTopAddr = nextCandidate
 			continue
 		}
@@ -230,9 +255,25 @@ func (a *addressIterator) peek(lock bool) net.Addr {
 
 		addrID := a.currentTopAddr.Value.(string)
 		candidate, ok := a.candidates[addrID]
+
+		// If the address cannot be found in the candidate set, then
+		// this must mean that the Remove method was called for the
+		// address. The Remove method would have checked that the
+		// address is not the last one in the iterator and that it has
+		// no locks on it. It is therefor safe to remove.
 		if !ok {
+			// Grab the next address candidate. This might be nil
+			// if the iterator is on the last item in the list.
 			nextCandidate := a.currentTopAddr.Next()
+
+			// Remove the address from the list that is no longer
+			// in the candidate set.
 			a.addrList.Remove(a.currentTopAddr)
+
+			// Set the current top to the next candidate. This might
+			// mean setting it to nil if the iterator is on its last
+			// item but this will be reset at the top of the for
+			// loop.
 			a.currentTopAddr = nextCandidate
 			continue
 		}

watchtower/wtclient/addr_iterator_test.go

@@ -10,30 +10,47 @@ import (
 
 // TestAddrIterator tests the behaviour of the addressIterator.
 func TestAddrIterator(t *testing.T) {
-	// Assert that an iterator can't be initialised with an empty address
+	t.Parallel()
-	// list.
-	_, err := newAddressIterator()
-	require.ErrorContains(t, err, "must have at least one address")
 
+	// Prep some addresses to use in the subtests.
 	addr1, err := net.ResolveTCPAddr("tcp", "1.2.3.4:8000")
 	require.NoError(t, err)
 
+	addr2, err := net.ResolveTCPAddr("tcp", "1.2.3.4:8001")
+	require.NoError(t, err)
+
+	addr3, err := net.ResolveTCPAddr("tcp", "1.2.3.4:8002")
+	require.NoError(t, err)
+
+	t.Run("always at least one address", func(t *testing.T) {
+		t.Parallel()
+
+		// Assert that an iterator can't be initialised with an empty
+		// address list.
+		_, err := newAddressIterator()
+		require.ErrorContains(t, err, "must have at least one address")
+
 		// Initialise the iterator with addr1.
 		iter, err := newAddressIterator(addr1)
 		require.NoError(t, err)
 
-	// Attempting to remove addr1 should fail now since it is the only
+		// Attempting to remove addr1 should fail now since it is the
-	// address in the iterator.
+		// only address in the iterator.
 		iter.Add(addr1)
 		err = iter.Remove(addr1)
 		require.ErrorIs(t, err, wtdb.ErrLastTowerAddr)
 
-	// Adding a duplicate of addr1 and then calling Remove should still
+		// Adding a duplicate of addr1 and then calling Remove should
-	// return an error.
+		// still return an error.
 		err = iter.Remove(addr1)
 		require.ErrorIs(t, err, wtdb.ErrLastTowerAddr)
+	})
 
-	addr2, err := net.ResolveTCPAddr("tcp", "1.2.3.4:8001")
+	t.Run("address iteration", func(t *testing.T) {
+		t.Parallel()
+
+		// Initialise the iterator with addr1.
+		iter, err := newAddressIterator(addr1)
 		require.NoError(t, err)
 
 		// Add addr2 to the iterator.
@@ -66,7 +83,8 @@ func TestAddrIterator(t *testing.T) {
 		a1 = iter.Peek()
 		require.Equal(t, addr1, a1)
 
-	// Wind the list to the end again so that we can test the Reset func.
+		// Wind the list to the end again so that we can test the Reset
+		// func.
 		_, err = iter.Next()
 		require.NoError(t, err)
 
@@ -80,11 +98,8 @@ func TestAddrIterator(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, addr2, a2)
 
-	addr3, err := net.ResolveTCPAddr("tcp", "1.2.3.4:8002")
+		// Add addr3 now to ensure that the iteration works even if we
-	require.NoError(t, err)
+		// are midway through the queue.
-
-	// Add addr3 now to ensure that the iteration works even if we are
-	// midway through the queue.
 		iter.Add(addr3)
 
 		// Now Next should return addr 3.
@@ -95,12 +110,16 @@ func TestAddrIterator(t *testing.T) {
 		// Quickly test that GetAll correctly returns a copy of all the
 		// addresses in the iterator.
 		addrList := iter.GetAll()
-	require.ElementsMatch(t, addrList, []net.Addr{addr1, addr2, addr3})
+		require.ElementsMatch(
+			t, addrList, []net.Addr{addr1, addr2, addr3},
+		)
 
-	// Also check that an iterator constructed via the Copy method, also
+		// Also check that an iterator constructed via the Copy method,
-	// contains all the expected addresses.
+		// also contains all the expected addresses.
 		newIterAddrs := iter.Copy().GetAll()
-	require.ElementsMatch(t, newIterAddrs, []net.Addr{addr1, addr2, addr3})
+		require.ElementsMatch(
+			t, newIterAddrs, []net.Addr{addr1, addr2, addr3},
+		)
 
 		// Let's now remove addr3.
 		err = iter.Remove(addr3)
@@ -109,11 +128,18 @@ func TestAddrIterator(t *testing.T) {
 		// Since addr3 is gone, Peek should return addr1.
 		a1 = iter.Peek()
 		require.Equal(t, addr1, a1)
+	})
 
-	// Lastly, we will test the "locking" of addresses.
+	t.Run("address locking", func(t *testing.T) {
+		t.Parallel()
 
-	// First we test the locking of an address via the PeekAndLock function.
+		// Initialise the iterator with addr1 and addr2.
-	a1 = iter.PeekAndLock()
+		iter, err := newAddressIterator(addr1, addr2)
+		require.NoError(t, err)
+
+		// First we test the locking of an address via the PeekAndLock
+		// function.
+		a1 := iter.PeekAndLock()
 		require.Equal(t, addr1, a1)
 		require.True(t, iter.HasLocked())
 
@@ -125,16 +151,16 @@ func TestAddrIterator(t *testing.T) {
 		iter.ReleaseLock(addr1)
 		require.False(t, iter.HasLocked())
 
-	// Since the lock has been released, we should now be able to remove
+		// Since the lock has been released, we should now be able to
-	// addr1.
+		// remove addr1.
 		err = iter.Remove(addr1)
 		require.NoError(t, err)
 
-	// Now we test the locking of an address via the NextAndLock function.
+		// Now we test the locking of an address via the NextAndLock
-	// To do this, we first re-add addr3.
+		// function. To do this, we first re-add addr3.
 		iter.Add(addr3)
 
-	a2, err = iter.NextAndLock()
+		a2, err := iter.NextAndLock()
 		require.NoError(t, err)
 		require.Equal(t, addr2, a2)
 		require.True(t, iter.HasLocked())
@@ -147,23 +173,22 @@ func TestAddrIterator(t *testing.T) {
 		iter.ReleaseLock(addr2)
 		require.False(t, iter.HasLocked())
 
-	// Since the lock has been released, we should now be able to remove
+		// Since the lock has been released, we should now be able to
-	// addr1.
+		// remove addr1.
 		err = iter.Remove(addr2)
 		require.NoError(t, err)
 
 		// Only addr3 should still be left in the iterator.
-	addrList = iter.GetAll()
+		addrList := iter.GetAll()
 		require.Len(t, addrList, 1)
 		require.Contains(t, addrList, addr3)
 
-	// Ensure that HasLocked acts correctly in the case where more than one
+		// Ensure that HasLocked acts correctly in the case where more
-	// address is being locked and unlock as well as the case where the same
+		// than one address is being locked and unlock as well as the
-	// address is locked more than once.
+		// case where the same address is locked more than once.
-
 		require.False(t, iter.HasLocked())
 
-	a3 = iter.PeekAndLock()
+		a3 := iter.PeekAndLock()
 		require.Equal(t, addr3, a3)
 		require.True(t, iter.HasLocked())
 
@@ -177,17 +202,35 @@ func TestAddrIterator(t *testing.T) {
 		require.Equal(t, addr2, a2)
 		require.True(t, iter.HasLocked())
 
-	// Now release addr2 and asset that HasLock is still true.
+		// Now release addr2 and assert that HasLock is still true.
 		iter.ReleaseLock(addr2)
 		require.True(t, iter.HasLocked())
 
-	// Releasing one of the locks on addr3 now should still result in
+		// Releasing one of the locks on addr3 now should still result
-	// HasLocked returning true.
+		// in HasLocked returning true.
 		iter.ReleaseLock(addr3)
 		require.True(t, iter.HasLocked())
 
-	// Releasing it again should now result in should still result in
+		// Releasing it again should now result in should still result
-	// HasLocked returning false.
+		// in HasLocked returning false.
 		iter.ReleaseLock(addr3)
 		require.False(t, iter.HasLocked())
+	})
+
+	t.Run("calling Next twice without Reset is safe", func(t *testing.T) {
+		t.Parallel()
+
+		// Initialise the iterator with addr1.
+		iter, err := newAddressIterator(addr1)
+		require.NoError(t, err)
+
+		a1 := iter.Peek()
+		require.Equal(t, addr1, a1)
+
+		_, err = iter.Next()
+		require.ErrorIs(t, err, ErrAddressesExhausted)
+
+		_, err = iter.Next()
+		require.ErrorIs(t, err, ErrAddressesExhausted)
+	})
 }

watchtower/wtclient/candidate_iterator.go

@@ -145,6 +145,7 @@ func (t *towerListIterator) AddCandidate(candidate *Tower) {
 		for {
 			next, err := candidate.Addresses.Next()
 			if err != nil {
+				candidate.Addresses.Reset()
 				break
 			}