lnd/watchtower/wtclient/candidate_iterator_test.go

package wtclient

import (
	"encoding/binary"
	"math/rand"
	"net"
	"reflect"
	"testing"
	"time"

	"github.com/btcsuite/btcd/btcec"
	"github.com/davecgh/go-spew/spew"
	"github.com/lightningnetwork/lnd/watchtower/wtdb"
)

func init() {
	rand.Seed(time.Now().Unix())
}

func randAddr(t *testing.T) net.Addr {
	var ip [4]byte
	if _, err := rand.Read(ip[:]); err != nil {
		t.Fatal(err)
	}
	var port [2]byte
	if _, err := rand.Read(port[:]); err != nil {
		t.Fatal(err)

	}
	return &net.TCPAddr{
		IP:   net.IP(ip[:]),
		Port: int(binary.BigEndian.Uint16(port[:])),
	}
}

func randTower(t *testing.T) *wtdb.Tower {
	priv, err := btcec.NewPrivateKey(btcec.S256())
	if err != nil {
		t.Fatalf("unable to create private key: %v", err)
	}
	pubKey := priv.PubKey()
	pubKey.Curve = nil
	return &wtdb.Tower{
		ID:          wtdb.TowerID(rand.Uint64()),
		IdentityKey: pubKey,
		Addresses:   []net.Addr{randAddr(t)},
	}
}

func copyTower(tower *wtdb.Tower) *wtdb.Tower {
	t := &wtdb.Tower{
		ID:          tower.ID,
		IdentityKey: tower.IdentityKey,
		Addresses:   make([]net.Addr, len(tower.Addresses)),
	}
	copy(t.Addresses, tower.Addresses)
	return t
}

func assertActiveCandidate(t *testing.T, i TowerCandidateIterator,
	c *wtdb.Tower, active bool) {

	isCandidate := i.IsActive(c.ID)
	if isCandidate && !active {
		t.Fatalf("expected tower %v to no longer be an active candidate",
			c.ID)
	}
	if !isCandidate && active {
		t.Fatalf("expected tower %v to be an active candidate", c.ID)
	}
}

func assertNextCandidate(t *testing.T, i TowerCandidateIterator, c *wtdb.Tower) {
	t.Helper()

	tower, err := i.Next()
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(tower, c) {
		t.Fatalf("expected tower: %v\ngot: %v", spew.Sdump(c),
			spew.Sdump(tower))
	}
}

// TestTowerCandidateIterator asserts the internal state of a
// TowerCandidateIterator after a series of updates to its candidates.
func TestTowerCandidateIterator(t *testing.T) {
	t.Parallel()

	// We'll start our test by creating an iterator of four candidate
	// towers. We'll use copies of these towers within the iterator to
	// ensure the iterator properly updates the state of its candidates.
	const numTowers = 4
	towers := make([]*wtdb.Tower, 0, numTowers)
	for i := 0; i < numTowers; i++ {
		towers = append(towers, randTower(t))
	}
	towerCopies := make([]*wtdb.Tower, 0, numTowers)
	for _, tower := range towers {
		towerCopies = append(towerCopies, copyTower(tower))
	}
	towerIterator := newTowerListIterator(towerCopies...)

	// We should expect to see all of our candidates in the order that they
	// were added.
	for _, expTower := range towers {
		tower, err := towerIterator.Next()
		if err != nil {
			t.Fatal(err)
		}
		if !reflect.DeepEqual(tower, expTower) {
			t.Fatalf("expected tower: %v\ngot: %v",
				spew.Sdump(expTower), spew.Sdump(tower))
		}
	}

	if _, err := towerIterator.Next(); err != ErrTowerCandidatesExhausted {
		t.Fatalf("expected ErrTowerCandidatesExhausted, got %v", err)
	}
	towerIterator.Reset()

	// We'll then attempt to test the RemoveCandidate behavior of the
	// iterator. We'll remove the address of the first tower, which should
	// result in it not having any addresses left, but still being an active
	// candidate.
	firstTower := towers[0]
	firstTowerAddr := firstTower.Addresses[0]
	firstTower.RemoveAddress(firstTowerAddr)
	towerIterator.RemoveCandidate(firstTower.ID, firstTowerAddr)
	assertActiveCandidate(t, towerIterator, firstTower, true)
	assertNextCandidate(t, towerIterator, firstTower)

	// We'll then remove the second tower completely from the iterator by
	// not providing the optional address. Since it's been removed, we
	// should expect to see the third tower next.
	secondTower, thirdTower := towers[1], towers[2]
	towerIterator.RemoveCandidate(secondTower.ID, nil)
	assertActiveCandidate(t, towerIterator, secondTower, false)
	assertNextCandidate(t, towerIterator, thirdTower)

	// We'll then update the fourth candidate with a new address. A
	// duplicate shouldn't be added since it already exists within the
	// iterator, but the new address should be.
	fourthTower := towers[3]
	assertActiveCandidate(t, towerIterator, fourthTower, true)
	fourthTower.AddAddress(randAddr(t))
	towerIterator.AddCandidate(fourthTower)
	assertNextCandidate(t, towerIterator, fourthTower)

	// Finally, we'll attempt to add a new candidate to the end of the
	// iterator. Since it didn't already exist and we've reached the end, it
	// should be available as the next candidate.
	towerIterator.AddCandidate(secondTower)
	assertActiveCandidate(t, towerIterator, secondTower, true)
	assertNextCandidate(t, towerIterator, secondTower)
}