From ed1fae6c73639a0a1d25498c4dac7afda2230259 Mon Sep 17 00:00:00 2001
From: Tom DNetto <tom@tailscale.com>
Date: Tue, 15 Nov 2022 15:03:11 -0800
Subject: [PATCH] ipn/ipnlocal: always tx TKA sync after enablement

By always firing off a sync after enablement, the control plane should know the node's TKA head
at all times.

Signed-off-by: Tom DNetto <tom@tailscale.com>
---
 ipn/ipnlocal/network-lock.go      |  7 ++++++-
 ipn/ipnlocal/network-lock_test.go | 26 ++++++++++++++++++++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/ipn/ipnlocal/network-lock.go b/ipn/ipnlocal/network-lock.go
index 9823367da..afd328a98 100644
--- a/ipn/ipnlocal/network-lock.go
+++ b/ipn/ipnlocal/network-lock.go
@@ -115,6 +115,7 @@ func (b *LocalBackend) tkaSyncIfNeeded(nm *netmap.NetworkMap, prefs ipn.PrefsVie
 
 	isEnabled := b.tka != nil
 	wantEnabled := nm.TKAEnabled
+	didJustEnable := false
 	if isEnabled != wantEnabled {
 		var ourHead tka.AUMHash
 		if b.tka != nil {
@@ -135,6 +136,7 @@ func (b *LocalBackend) tkaSyncIfNeeded(nm *netmap.NetworkMap, prefs ipn.PrefsVie
 				return fmt.Errorf("bootstrap: %w", err)
 			}
 			isEnabled = true
+			didJustEnable = true
 		} else if !wantEnabled && isEnabled {
 			if err := b.tkaApplyDisablementLocked(bs.DisablementSecret); err != nil {
 				// We log here instead of returning an error (which itself would be
@@ -149,7 +151,10 @@ func (b *LocalBackend) tkaSyncIfNeeded(nm *netmap.NetworkMap, prefs ipn.PrefsVie
 		}
 	}
 
-	if isEnabled && b.tka.authority.Head() != nm.TKAHead {
+	// We always transmit the sync RPCs if TKA was just enabled.
+	// This informs the control plane that our TKA state is now
+	// initialized to the transmitted TKA head hash.
+	if isEnabled && (b.tka.authority.Head() != nm.TKAHead || didJustEnable) {
 		if err := b.tkaSyncLocked(ourNodeKey); err != nil {
 			return fmt.Errorf("tka sync: %w", err)
 		}
diff --git a/ipn/ipnlocal/network-lock_test.go b/ipn/ipnlocal/network-lock_test.go
index 6d34b16a8..1044e2d6d 100644
--- a/ipn/ipnlocal/network-lock_test.go
+++ b/ipn/ipnlocal/network-lock_test.go
@@ -108,8 +108,30 @@ func TestTKAEnablementFlow(t *testing.T) {
 				t.Fatal(err)
 			}
 
-		case "/machine/tka/sync/offer", "/machine/tka/sync/send":
-			t.Error("node attempted to sync, but should have been up to date")
+		// Sync offer/send endpoints are hit even though the node is up-to-date,
+		// so we implement enough of a fake that the client doesn't explode.
+		case "/machine/tka/sync/offer":
+			head, err := a1.Head().MarshalText()
+			if err != nil {
+				t.Fatal(err)
+			}
+			w.WriteHeader(200)
+			if err := json.NewEncoder(w).Encode(tailcfg.TKASyncOfferResponse{
+				Head: string(head),
+			}); err != nil {
+				t.Fatal(err)
+			}
+		case "/machine/tka/sync/send":
+			head, err := a1.Head().MarshalText()
+			if err != nil {
+				t.Fatal(err)
+			}
+			w.WriteHeader(200)
+			if err := json.NewEncoder(w).Encode(tailcfg.TKASyncSendResponse{
+				Head: string(head),
+			}); err != nil {
+				t.Fatal(err)
+			}
 
 		default:
 			t.Errorf("unhandled endpoint path: %v", r.URL.Path)