From 85373ef822645f66242c9b44dbd754247d0d9c63 Mon Sep 17 00:00:00 2001
From: Alex Chan <alexc@tailscale.com>
Date: Tue, 18 Nov 2025 09:44:12 +0000
Subject: [PATCH] tka: move RemoveAll() to CompactableChonk

I added a RemoveAll() method on tka.Chonk in #17946, but it's only used
in the node to purge local AUMs. We don't need it in the SQLite storage,
which currently implements tka.Chonk, so move it to CompactableChonk
instead.

Also add some automated tests, as a safety net.

Updates tailscale/corp#33599

Change-Id: I54de9ccf1d6a3d29b36a94eccb0ebd235acd4ebc
Signed-off-by: Alex Chan <alexc@tailscale.com>
---
 tka/tailchonk.go              |  8 +++---
 tstest/chonktest/chonktest.go | 47 +++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/tka/tailchonk.go b/tka/tailchonk.go
index 2dc03a6f6..0b7191747 100644
--- a/tka/tailchonk.go
+++ b/tka/tailchonk.go
@@ -58,10 +58,6 @@ type Chonk interface {
 	// as a hint to pick the correct chain in the event that the Chonk stores
 	// multiple distinct chains.
 	LastActiveAncestor() (*AUMHash, error)
-
-	// RemoveAll permanently and completely clears the TKA state. This should
-	// be called when the user disables Tailnet Lock.
-	RemoveAll() error
 }
 
 // CompactableChonk implementation are extensions of Chonk, which are
@@ -80,6 +76,10 @@ type CompactableChonk interface {
 	// PurgeAUMs permanently and irrevocably deletes the specified
 	// AUMs from storage.
 	PurgeAUMs(hashes []AUMHash) error
+
+	// RemoveAll permanently and completely clears the TKA state. This should
+	// be called when the user disables Tailnet Lock.
+	RemoveAll() error
 }
 
 // Mem implements in-memory storage of TKA state, suitable for
diff --git a/tstest/chonktest/chonktest.go b/tstest/chonktest/chonktest.go
index bfe394b28..404f1ec47 100644
--- a/tstest/chonktest/chonktest.go
+++ b/tstest/chonktest/chonktest.go
@@ -9,6 +9,7 @@ package chonktest
 import (
 	"bytes"
 	"encoding/binary"
+	"errors"
 	"math/rand"
 	"os"
 	"testing"
@@ -253,4 +254,50 @@ func RunCompactableChonkTests(t *testing.T, newChonk func(t *testing.T) tka.Comp
 			t.Fatalf("ChildAUMs() output differs (-want, +got):\n%s", diff)
 		}
 	})
+
+	t.Run("RemoveAll", func(t *testing.T) {
+		t.Parallel()
+		chonk := newChonk(t)
+		parentHash := randHash(t, 1)
+		data := []tka.AUM{
+			{
+				MessageKind: tka.AUMRemoveKey,
+				KeyID:       []byte{1, 2},
+				PrevAUMHash: parentHash[:],
+			},
+			{
+				MessageKind: tka.AUMRemoveKey,
+				KeyID:       []byte{3, 4},
+				PrevAUMHash: parentHash[:],
+			},
+		}
+
+		if err := chonk.CommitVerifiedAUMs(data); err != nil {
+			t.Fatalf("CommitVerifiedAUMs failed: %v", err)
+		}
+
+		// Check we can retrieve the AUMs we just stored
+		for _, want := range data {
+			got, err := chonk.AUM(want.Hash())
+			if err != nil {
+				t.Fatalf("could not get %s: %v", want.Hash(), err)
+			}
+			if diff := cmp.Diff(want, got); diff != "" {
+				t.Errorf("stored AUM %s differs (-want, +got):\n%s", want.Hash(), diff)
+			}
+		}
+
+		// Call RemoveAll() to drop all the AUM state
+		if err := chonk.RemoveAll(); err != nil {
+			t.Fatalf("RemoveAll failed: %v", err)
+		}
+
+		// Check we can no longer retrieve the previously-stored AUMs
+		for _, want := range data {
+			aum, err := chonk.AUM(want.Hash())
+			if !errors.Is(err, os.ErrNotExist) {
+				t.Fatalf("expected os.ErrNotExist for %s, instead got aum=%v, err=%v", want.Hash(), aum, err)
+			}
+		}
+	})
 }