wtclient: test case re-add removed tower

watchtower/wtclient/client_test.go

@@ -373,6 +373,7 @@ type testHarness struct {
 	clientDB   *wtmock.ClientDB
 	clientCfg  *wtclient.Config
 	client     wtclient.Client
+	serverAddr *lnwire.NetAddress
 	serverDB   *wtmock.TowerDB
 	serverCfg  *wtserver.Config
 	server     *wtserver.Server
@@ -474,6 +475,7 @@ func newHarness(t *testing.T, cfg harnessCfg) *testHarness {
 		clientDB:   clientDB,
 		clientCfg:  clientCfg,
 		client:     client,
+		serverAddr: towerAddr,
 		serverDB:   serverDB,
 		serverCfg:  serverCfg,
 		server:     server,
@@ -782,6 +784,25 @@ func (h *testHarness) assertUpdatesForPolicy(hints []blob.BreachHint,
 	}
 }
 
+// addTower adds a tower found at `addr` to the client.
+func (h *testHarness) addTower(addr *lnwire.NetAddress) {
+	h.t.Helper()
+
+	if err := h.client.AddTower(addr); err != nil {
+		h.t.Fatalf("unable to add tower: %v", err)
+	}
+}
+
+// removeTower removes a tower from the client. If `addr` is specified, then the
+// only said address is removed from the tower.
+func (h *testHarness) removeTower(pubKey *btcec.PublicKey, addr net.Addr) {
+	h.t.Helper()
+
+	if err := h.client.RemoveTower(pubKey, addr); err != nil {
+		h.t.Fatalf("unable to remove tower: %v", err)
+	}
+}
+
 const (
 	localBalance  = lnwire.MilliSatoshi(100000000)
 	remoteBalance = lnwire.MilliSatoshi(200000000)
@@ -1396,6 +1417,60 @@ var clientTests = []clientTest{
 			h.waitServerUpdates(hints, 5*time.Second)
 		},
 	},
+	{
+		// Asserts that the client can continue making backups to a
+		// tower that's been re-added after it's been removed.
+		name: "re-add removed tower",
+		cfg: harnessCfg{
+			localBalance:  localBalance,
+			remoteBalance: remoteBalance,
+			policy: wtpolicy.Policy{
+				TxPolicy: wtpolicy.TxPolicy{
+					BlobType:     blob.TypeAltruistCommit,
+					SweepFeeRate: wtpolicy.DefaultSweepFeeRate,
+				},
+				MaxUpdates: 5,
+			},
+		},
+		fn: func(h *testHarness) {
+			const (
+				chanID     = 0
+				numUpdates = 4
+			)
+
+			// Create four channel updates and only back up the
+			// first two.
+			hints := h.advanceChannelN(chanID, numUpdates)
+			h.backupStates(chanID, 0, numUpdates/2, nil)
+			h.waitServerUpdates(hints[:numUpdates/2], 5*time.Second)
+
+			// Fully remove the tower, causing its existing sessions
+			// to be marked inactive.
+			h.removeTower(h.serverAddr.IdentityKey, nil)
+
+			// Back up the remaining states. Since the tower has
+			// been removed, it shouldn't receive any updates.
+			h.backupStates(chanID, numUpdates/2, numUpdates, nil)
+			h.waitServerUpdates(nil, time.Second)
+
+			// Re-add the tower. We prevent the tower from acking
+			// session creation to ensure the inactive sessions are
+			// not used.
+			h.server.Stop()
+			h.serverCfg.NoAckCreateSession = true
+			h.startServer()
+			h.addTower(h.serverAddr)
+			h.waitServerUpdates(nil, time.Second)
+
+			// Finally, allow the tower to ack session creation,
+			// allowing the state updates to be sent through the new
+			// session.
+			h.server.Stop()
+			h.serverCfg.NoAckCreateSession = false
+			h.startServer()
+			h.waitServerUpdates(hints[numUpdates/2:], 5*time.Second)
+		},
+	},
 }
 
 // TestClient executes the client test suite, asserting the ability to backup