From 8b0cc487f0a4c3244afde62adb62e49eaee9e9b1 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Fri, 15 Mar 2019 02:33:05 -0700
Subject: [PATCH] watchtower/wtdb+wtserver: allow retransmission of last update

---
 watchtower/wtdb/session_info.go    |  2 +-
 watchtower/wtserver/server_test.go | 36 ++++++++++++++++++++++++++++--
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/watchtower/wtdb/session_info.go b/watchtower/wtdb/session_info.go
index f1b2e2a8..1c7b7f0f 100644
--- a/watchtower/wtdb/session_info.go
+++ b/watchtower/wtdb/session_info.go
@@ -82,7 +82,7 @@ func (s *SessionInfo) AcceptUpdateSequence(seqNum, lastApplied uint16) error {
 		return ErrSessionConsumed
 
 	// Client update does not match our expected next seqnum.
-	case seqNum != s.LastApplied+1:
+	case seqNum != s.LastApplied && seqNum != s.LastApplied+1:
 		return ErrUpdateOutOfOrder
 	}
 
diff --git a/watchtower/wtserver/server_test.go b/watchtower/wtserver/server_test.go
index 4a6ee27e..6df7cd58 100644
--- a/watchtower/wtserver/server_test.go
+++ b/watchtower/wtserver/server_test.go
@@ -418,8 +418,8 @@ var stateUpdateTests = []stateUpdateTestCase{
 			{Code: wtwire.CodeOK, LastApplied: 4},
 		},
 	},
-	// Valid update sequence with disconnection, ensure resumes resume.
-	// Client doesn't echo last applied until last message.
+	// Valid update sequence with disconnection, resume next update. Client
+	// doesn't echo last applied until last message.
 	{
 		name: "resume after disconnect lagging lastapplied",
 		initMsg: wtwire.NewInitMessage(
@@ -448,6 +448,38 @@ var stateUpdateTests = []stateUpdateTestCase{
 			{Code: wtwire.CodeOK, LastApplied: 4},
 		},
 	},
+	// Valid update sequence with disconnection, resume last update.  Client
+	// doesn't echo last applied until last message.
+	{
+		name: "resume after disconnect lagging lastapplied",
+		initMsg: wtwire.NewInitMessage(
+			lnwire.NewRawFeatureVector(),
+			testnetChainHash,
+		),
+		createMsg: &wtwire.CreateSession{
+			BlobType:     blob.TypeDefault,
+			MaxUpdates:   4,
+			RewardBase:   0,
+			RewardRate:   0,
+			SweepFeeRate: 1,
+		},
+		updates: []*wtwire.StateUpdate{
+			{SeqNum: 1, LastApplied: 0},
+			{SeqNum: 2, LastApplied: 0},
+			nil, // Wait for read timeout to drop conn, then reconnect.
+			{SeqNum: 2, LastApplied: 0},
+			{SeqNum: 3, LastApplied: 0},
+			{SeqNum: 4, LastApplied: 3},
+		},
+		replies: []*wtwire.StateUpdateReply{
+			{Code: wtwire.CodeOK, LastApplied: 1},
+			{Code: wtwire.CodeOK, LastApplied: 2},
+			nil,
+			{Code: wtwire.CodeOK, LastApplied: 2},
+			{Code: wtwire.CodeOK, LastApplied: 3},
+			{Code: wtwire.CodeOK, LastApplied: 4},
+		},
+	},
 	// Send update with sequence number that exceeds MaxUpdates.
 	{
 		name: "seqnum exceed maxupdates",