From 359b2049b674049ac4f91a1cdb52fb26814a6dab Mon Sep 17 00:00:00 2001
From: Wilmer Paulino <wilmer.paulino@gmail.com>
Date: Wed, 3 Jul 2019 19:41:39 -0700
Subject: [PATCH] watchtower: load persisted towers upon TowerClient creation

We do this as a convenience for WatchtowerClient users so that they do
not need to re-add towers upon restarts. We ensure not to re-add towers
that have been previously removed by determining whether it has any
lingering active sessions.
---
 watchtower/wtclient/client.go | 71 +++++++++++++++++++----------------
 watchtower/wtdb/tower.go      | 11 ++++++
 2 files changed, 50 insertions(+), 32 deletions(-)

diff --git a/watchtower/wtclient/client.go b/watchtower/wtclient/client.go
index be199437..58666e0d 100644
--- a/watchtower/wtclient/client.go
+++ b/watchtower/wtclient/client.go
@@ -273,40 +273,53 @@ func New(config *Config) (*TowerClient, error) {
 		return nil, err
 	}
 
-	log.Infof("Using private watchtower %s, offering policy %s",
-		cfg.PrivateTower, cfg.Policy)
-
-	candidateTowers := newTowerListIterator(tower)
-
-	// Next, load all active sessions from the db into the client. We will
-	// use any of these session if their policies match the current policy
-	// of the client, otherwise they will be ignored and new sessions will
-	// be requested.
+	// Next, load all candidate sessions and towers from the database into
+	// the client. We will use any of these session if their policies match
+	// the current policy of the client, otherwise they will be ignored and
+	// new sessions will be requested.
 	sessions, err := cfg.DB.ListClientSessions(nil)
 	if err != nil {
 		return nil, err
 	}
 
-	// Reload any towers from disk using the tower IDs contained in each
-	// candidate session. We will also rederive any session keys needed to
-	// be able to communicate with the towers and authenticate session
-	// requests. This prevents us from having to store the private keys on
-	// disk.
+	candidateSessions := make(map[wtdb.SessionID]*wtdb.ClientSession)
+	sessionTowers := make(map[wtdb.TowerID]*wtdb.Tower)
 	for _, s := range sessions {
-		tower, err := cfg.DB.LoadTowerByID(s.TowerID)
-		if err != nil {
-			return nil, err
+		// Candidate sessions must be in an active state.
+		if s.Status != wtdb.CSessionActive {
+			continue
 		}
 
-		sessionPriv, err := DeriveSessionKey(
-			cfg.SecretKeyRing, s.KeyIndex,
-		)
-		if err != nil {
-			return nil, err
+		// Reload the tower from disk using the tower ID contained in
+		// each candidate session. We will also rederive any session
+		// keys needed to be able to communicate with the towers and
+		// authenticate session requests. This prevents us from having
+		// to store the private keys on disk.
+		tower, ok := sessionTowers[s.TowerID]
+		if !ok {
+			var err error
+			tower, err = cfg.DB.LoadTowerByID(s.TowerID)
+			if err != nil {
+				return nil, err
+			}
 		}
-
 		s.Tower = tower
-		s.SessionPrivKey = sessionPriv
+
+		sessionKey, err := DeriveSessionKey(cfg.SecretKeyRing, s.KeyIndex)
+		if err != nil {
+			return nil, err
+		}
+		s.SessionPrivKey = sessionKey
+
+		candidateSessions[s.ID] = s
+		sessionTowers[tower.ID] = tower
+	}
+
+	var candidateTowers []*wtdb.Tower
+	for _, tower := range sessionTowers {
+		log.Infof("Using private watchtower %s, offering policy %s",
+			tower, cfg.Policy)
+		candidateTowers = append(candidateTowers, tower)
 	}
 
 	// Load the sweep pkscripts that have been generated for all previously
@@ -319,8 +332,8 @@ func New(config *Config) (*TowerClient, error) {
 	c := &TowerClient{
 		cfg:               cfg,
 		pipeline:          newTaskPipeline(),
-		candidateTowers:   candidateTowers,
-		candidateSessions: sessions,
+		candidateTowers:   newTowerListIterator(candidateTowers...),
+		candidateSessions: candidateSessions,
 		activeSessions:    make(sessionQueueSet),
 		summaries:         chanSummaries,
 		statTicker:        time.NewTicker(DefaultStatInterval),
@@ -612,12 +625,6 @@ func (c *TowerClient) nextSessionQueue() *sessionQueue {
 			continue
 		}
 
-		// Skip any sessions that are still active, but are not for the
-		// users currently configured tower.
-		if !c.candidateTowers.IsActive(sessionInfo.TowerID) {
-			continue
-		}
-
 		candidateSession = sessionInfo
 		break
 	}
diff --git a/watchtower/wtdb/tower.go b/watchtower/wtdb/tower.go
index fc4a0404..1924d9d5 100644
--- a/watchtower/wtdb/tower.go
+++ b/watchtower/wtdb/tower.go
@@ -1,6 +1,8 @@
 package wtdb
 
 import (
+	"encoding/hex"
+	"fmt"
 	"io"
 	"net"
 
@@ -92,6 +94,15 @@ func (t *Tower) LNAddrs() []*lnwire.NetAddress {
 	return addrs
 }
 
+// String returns a user-friendly identifier of the tower.
+func (t *Tower) String() string {
+	pubKey := hex.EncodeToString(t.IdentityKey.SerializeCompressed())
+	if len(t.Addresses) == 0 {
+		return pubKey
+	}
+	return fmt.Sprintf("%v@%v", pubKey, t.Addresses[0])
+}
+
 // Encode writes the Tower to the passed io.Writer. The TowerID is not
 // serialized, since it acts as the key.
 func (t *Tower) Encode(w io.Writer) error {