tka: move RemoveAll() to CompactableChonk

I added a RemoveAll() method on tka.Chonk in tailscale#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 <[email protected]>
(cherry picked from commit c17ba64)

Author: alexwlchan

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

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)
+			}
+		}
+	})
 }