From 978a0a8de6719768af1f80cd13d23b28151dd21a Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Tue, 23 Apr 2019 20:06:04 -0700
Subject: [PATCH] watchtower/wtserver/server: permit unused session overrwite

---
 watchtower/wtdb/mock.go               |  3 +-
 watchtower/wtserver/create_session.go |  6 ++-
 watchtower/wtserver/server_test.go    | 77 ++++++++++++++++++++++++---
 3 files changed, 76 insertions(+), 10 deletions(-)

diff --git a/watchtower/wtdb/mock.go b/watchtower/wtdb/mock.go
index fa53d6cf..90230388 100644
--- a/watchtower/wtdb/mock.go
+++ b/watchtower/wtdb/mock.go
@@ -61,7 +61,8 @@ func (db *MockDB) InsertSessionInfo(info *SessionInfo) error {
 	db.mu.Lock()
 	defer db.mu.Unlock()
 
-	if _, ok := db.sessions[info.ID]; ok {
+	dbInfo, ok := db.sessions[info.ID]
+	if ok && dbInfo.LastApplied > 0 {
 		return ErrSessionAlreadyExists
 	}
 
diff --git a/watchtower/wtserver/create_session.go b/watchtower/wtserver/create_session.go
index 2f89b85a..3217d601 100644
--- a/watchtower/wtserver/create_session.go
+++ b/watchtower/wtserver/create_session.go
@@ -21,11 +21,15 @@ func (s *Server) handleCreateSession(peer Peer, id *wtdb.SessionID,
 	existingInfo, err := s.cfg.DB.GetSessionInfo(id)
 	switch {
 
+	// We already have a session, though it is currently unused. We'll allow
+	// the client to recommit the session if it wanted to change the policy.
+	case err == nil && existingInfo.LastApplied == 0:
+
 	// We already have a session corresponding to this session id, return an
 	// error signaling that it already exists in our database. We return the
 	// reward address to the client in case they were not able to process
 	// our reply earlier.
-	case err == nil:
+	case err == nil && existingInfo.LastApplied > 0:
 		log.Debugf("Already have session for %s", id)
 		return s.replyCreateSession(
 			peer, id, wtwire.CreateSessionCodeAlreadyExists,
diff --git a/watchtower/wtserver/server_test.go b/watchtower/wtserver/server_test.go
index d9e3bec5..6d99180f 100644
--- a/watchtower/wtserver/server_test.go
+++ b/watchtower/wtserver/server_test.go
@@ -29,6 +29,8 @@ var (
 	addrScript, _ = txscript.PayToAddrScript(addr)
 
 	testnetChainHash = *chaincfg.TestNet3Params.GenesisHash
+
+	rewardType = (blob.FlagCommitOutputs | blob.FlagReward).Type()
 )
 
 // randPubKey generates a new secp keypair, and returns the public key.
@@ -152,16 +154,17 @@ func TestServerOnlyAcceptOnePeer(t *testing.T) {
 }
 
 type createSessionTestCase struct {
-	name        string
-	initMsg     *wtwire.Init
-	createMsg   *wtwire.CreateSession
-	expReply    *wtwire.CreateSessionReply
-	expDupReply *wtwire.CreateSessionReply
+	name            string
+	initMsg         *wtwire.Init
+	createMsg       *wtwire.CreateSession
+	expReply        *wtwire.CreateSessionReply
+	expDupReply     *wtwire.CreateSessionReply
+	sendStateUpdate bool
 }
 
 var createSessionTests = []createSessionTestCase{
 	{
-		name: "reject duplicate session create",
+		name: "duplicate session create",
 		initMsg: wtwire.NewInitMessage(
 			lnwire.NewRawFeatureVector(),
 			testnetChainHash,
@@ -173,12 +176,58 @@ var createSessionTests = []createSessionTestCase{
 			RewardRate:   0,
 			SweepFeeRate: 1,
 		},
+		expReply: &wtwire.CreateSessionReply{
+			Code: wtwire.CodeOK,
+			Data: []byte{},
+		},
+		expDupReply: &wtwire.CreateSessionReply{
+			Code: wtwire.CodeOK,
+			Data: []byte{},
+		},
+	},
+	{
+		name: "duplicate session create after use",
+		initMsg: wtwire.NewInitMessage(
+			lnwire.NewRawFeatureVector(),
+			testnetChainHash,
+		),
+		createMsg: &wtwire.CreateSession{
+			BlobType:     blob.TypeDefault,
+			MaxUpdates:   1000,
+			RewardBase:   0,
+			RewardRate:   0,
+			SweepFeeRate: 1,
+		},
+		expReply: &wtwire.CreateSessionReply{
+			Code: wtwire.CodeOK,
+			Data: []byte{},
+		},
+		expDupReply: &wtwire.CreateSessionReply{
+			Code:        wtwire.CreateSessionCodeAlreadyExists,
+			LastApplied: 1,
+			Data:        []byte{},
+		},
+		sendStateUpdate: true,
+	},
+	{
+		name: "duplicate session create reward",
+		initMsg: wtwire.NewInitMessage(
+			lnwire.NewRawFeatureVector(),
+			testnetChainHash,
+		),
+		createMsg: &wtwire.CreateSession{
+			BlobType:     rewardType,
+			MaxUpdates:   1000,
+			RewardBase:   0,
+			RewardRate:   0,
+			SweepFeeRate: 1,
+		},
 		expReply: &wtwire.CreateSessionReply{
 			Code: wtwire.CodeOK,
 			Data: addrScript,
 		},
 		expDupReply: &wtwire.CreateSessionReply{
-			Code: wtwire.CreateSessionCodeAlreadyExists,
+			Code: wtwire.CodeOK,
 			Data: addrScript,
 		},
 	},
@@ -251,6 +300,18 @@ func testServerCreateSession(t *testing.T, i int, test createSessionTestCase) {
 		return
 	}
 
+	if test.sendStateUpdate {
+		peer = wtmock.NewMockPeer(localPub, peerPub, nil, 0)
+		connect(t, s, peer, test.initMsg, timeoutDuration)
+		update := &wtwire.StateUpdate{
+			SeqNum:     1,
+			IsComplete: 1,
+		}
+		sendMsg(t, update, peer, timeoutDuration)
+
+		assertConnClosed(t, peer, 2*timeoutDuration)
+	}
+
 	// Simulate a peer with the same session id connection to the server
 	// again.
 	peer = wtmock.NewMockPeer(localPub, peerPub, nil, 0)
@@ -705,7 +766,7 @@ func TestServerDeleteSession(t *testing.T) {
 			send: createSession,
 			recv: &wtwire.CreateSessionReply{
 				Code: wtwire.CodeOK,
-				Data: addrScript,
+				Data: []byte{},
 			},
 			assert: func(t *testing.T) {
 				// Both peers should have sessions.