From d6dde5a1aca795be5165e6add00ba814c68ef7dc Mon Sep 17 00:00:00 2001
From: Maisem Ali <maisem@tailscale.com>
Date: Thu, 4 Nov 2021 12:19:00 -0700
Subject: [PATCH] ipn/ipnlocal: handle key extensions after key has already
 expired

Signed-off-by: Maisem Ali <maisem@tailscale.com>
---
 ipn/ipnlocal/local.go                  | 27 +++++++++++++++++-
 ipn/ipnlocal/state_test.go             | 39 ++++++++++++++++++++++++++
 tstest/integration/integration_test.go |  1 -
 3 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/ipn/ipnlocal/local.go b/ipn/ipnlocal/local.go
index 712154944..8f766bdbb 100644
--- a/ipn/ipnlocal/local.go
+++ b/ipn/ipnlocal/local.go
@@ -122,6 +122,7 @@ type LocalBackend struct {
 	engineStatus     ipn.EngineStatus
 	endpoints        []tailcfg.Endpoint
 	blocked          bool
+	keyExpired       bool
 	authURL          string // cleared on Notify
 	authURLSticky    string // not cleared on Notify
 	interact         bool
@@ -465,8 +466,22 @@ func (b *LocalBackend) setClientStatus(st controlclient.Status) {
 
 	b.mu.Lock()
 	wasBlocked := b.blocked
+	keyExpiryExtended := false
+	if st.NetMap != nil {
+		wasExpired := b.keyExpired
+		isExpired := !st.NetMap.Expiry.IsZero() && st.NetMap.Expiry.Before(time.Now())
+		if wasExpired && !isExpired {
+			keyExpiryExtended = true
+		}
+		b.keyExpired = isExpired
+	}
 	b.mu.Unlock()
 
+	if keyExpiryExtended && wasBlocked {
+		// Key extended, unblock the engine
+		b.blockEngineUpdates(false)
+	}
+
 	if st.LoginFinished != nil && wasBlocked {
 		// Auth completed, unblock the engine
 		b.blockEngineUpdates(false)
@@ -1774,6 +1789,12 @@ func (b *LocalBackend) NetMap() *netmap.NetworkMap {
 	return b.netMap
 }
 
+func (b *LocalBackend) isEngineBlocked() bool {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.blocked
+}
+
 // blockEngineUpdate sets b.blocked to block, while holding b.mu. Its
 // indirect effect is to turn b.authReconfig() into a no-op if block
 // is true.
@@ -2396,6 +2417,7 @@ func (b *LocalBackend) nextState() ipn.State {
 		wantRunning = b.prefs.WantRunning
 		loggedOut   = b.prefs.LoggedOut
 		st          = b.engineStatus
+		keyExpired  = b.keyExpired
 	)
 	b.mu.Unlock()
 
@@ -2428,7 +2450,9 @@ func (b *LocalBackend) nextState() ipn.State {
 		}
 	case !wantRunning:
 		return ipn.Stopped
-	case !netMap.Expiry.IsZero() && time.Until(netMap.Expiry) <= 0:
+	case keyExpired:
+		// NetMap must be non-nil for us to get here.
+		// The node key expired, need to relogin.
 		return ipn.NeedsLogin
 	case netMap.MachineStatus != tailcfg.MachineAuthorized:
 		// TODO(crawshaw): handle tailcfg.MachineInvalid
@@ -2509,6 +2533,7 @@ func (b *LocalBackend) ResetForClientDisconnect() {
 	b.userID = ""
 	b.setNetMapLocked(nil)
 	b.prefs = new(ipn.Prefs)
+	b.keyExpired = false
 	b.authURL = ""
 	b.authURLSticky = ""
 	b.activeLogin = ""
diff --git a/ipn/ipnlocal/state_test.go b/ipn/ipnlocal/state_test.go
index 970f4c9a0..1ad950f95 100644
--- a/ipn/ipnlocal/state_test.go
+++ b/ipn/ipnlocal/state_test.go
@@ -867,6 +867,45 @@ func TestStateMachine(t *testing.T) {
 		// change either.
 		c.Assert(ipn.Starting, qt.Equals, b.State())
 	}
+	t.Logf("\n\nExpireKey")
+	notifies.expect(1)
+	cc.send(nil, "", false, &netmap.NetworkMap{
+		Expiry:        time.Now().Add(-time.Minute),
+		MachineStatus: tailcfg.MachineAuthorized,
+	})
+	{
+		nn := notifies.drain(1)
+		cc.assertCalls("unpause", "unpause")
+		c.Assert(nn[0].State, qt.IsNotNil)
+		c.Assert(ipn.NeedsLogin, qt.Equals, *nn[0].State)
+		c.Assert(ipn.NeedsLogin, qt.Equals, b.State())
+		c.Assert(b.isEngineBlocked(), qt.IsTrue)
+	}
+
+	t.Logf("\n\nExtendKey")
+	notifies.expect(1)
+	cc.send(nil, "", false, &netmap.NetworkMap{
+		Expiry:        time.Now().Add(time.Minute),
+		MachineStatus: tailcfg.MachineAuthorized,
+	})
+	{
+		nn := notifies.drain(1)
+		cc.assertCalls("unpause", "unpause", "unpause")
+		c.Assert(nn[0].State, qt.IsNotNil)
+		c.Assert(ipn.Starting, qt.Equals, *nn[0].State)
+		c.Assert(ipn.Starting, qt.Equals, b.State())
+		c.Assert(b.isEngineBlocked(), qt.IsFalse)
+	}
+	notifies.expect(1)
+	// Fake a DERP connection.
+	b.setWgengineStatus(&wgengine.Status{DERPs: 1}, nil)
+	{
+		nn := notifies.drain(1)
+		cc.assertCalls("unpause")
+		c.Assert(nn[0].State, qt.IsNotNil)
+		c.Assert(ipn.Running, qt.Equals, *nn[0].State)
+		c.Assert(ipn.Running, qt.Equals, b.State())
+	}
 }
 
 type testStateStorage struct {
diff --git a/tstest/integration/integration_test.go b/tstest/integration/integration_test.go
index 62f208ab1..d74e75183 100644
--- a/tstest/integration/integration_test.go
+++ b/tstest/integration/integration_test.go
@@ -85,7 +85,6 @@ func TestOneNodeUp_NoAuth(t *testing.T) {
 }
 
 func TestOneNodeExpiredKey(t *testing.T) {
-	t.Skip("Test to exercise a problem which is not fixed yet.")
 	t.Parallel()
 	bins := BuildTestBinaries(t)