ipn/ipnlocal, tka: compact TKA state after every sync

Previously a TKA compaction would only run when a node starts, which means a long-running node could use unbounded storage as it accumulates ever-increasing amounts of TKA state. This patch changes TKA so it runs a compaction after every sync.

Updates tailscale/corp#33537

Change-Id: I91df887ea0c5a5b00cb6caced85aeffa2a4b24ee
Signed-off-by: Alex Chan <[email protected]>

Author: alexwlchan

ipn/ipnlocal/network-lock.go

@@ -360,6 +360,13 @@ func (b *LocalBackend) tkaSyncIfNeeded(nm *netmap.NetworkMap, prefs ipn.PrefsVie
 		if err := b.tkaSyncLocked(ourNodeKey); err != nil {
 			return fmt.Errorf("tka sync: %w", err)
 		}
+		// Try to compact the TKA state, to avoid unbounded storage on nodes.
+		//
+		// We run this on every sync so that clients compact consistently. In many
+		// cases this will be a no-op.
+		if err := b.tka.authority.Compact(b.tka.storage, tkaCompactionDefaults); err != nil {
+			return fmt.Errorf("tka compact: %w", err)
+		}
 	}
 
 	return nil
@@ -508,7 +515,7 @@ func (b *LocalBackend) tkaBootstrapFromGenesisLocked(g tkatype.MarshaledAUM, per
 	if root == "" {
 		b.health.SetUnhealthy(noNetworkLockStateDirWarnable, nil)
 		b.logf("network-lock using in-memory storage; no state directory")
-		storage = &tka.Mem{}
+		storage = tka.ChonkMem()
 	} else {
 		chonkDir := b.chonkPathLocked()
 		chonk, err := tka.ChonkDir(chonkDir)
@@ -686,7 +693,7 @@ func (b *LocalBackend) NetworkLockInit(keys []tka.Key, disablementValues [][]byt
 	// We use an in-memory tailchonk because we don't want to commit to
 	// the filesystem until we've finished the initialization sequence,
 	// just in case something goes wrong.
-	_, genesisAUM, err := tka.Create(&tka.Mem{}, tka.State{
+	_, genesisAUM, err := tka.Create(tka.ChonkMem(), tka.State{
 		Keys: keys,
 		// TODO(tom): s/tka.State.DisablementSecrets/tka.State.DisablementValues
 		//   This will center on consistent nomenclature:

ipn/ipnlocal/network-lock_test.go

@@ -17,6 +17,7 @@ import (
 	"path/filepath"
 	"reflect"
 	"testing"
+	"time"
 
 	go4mem "go4.org/mem"
 
@@ -31,6 +32,7 @@ import (
 	"tailscale.com/tailcfg"
 	"tailscale.com/tka"
 	"tailscale.com/tsd"
+	"tailscale.com/tstest"
 	"tailscale.com/types/key"
 	"tailscale.com/types/netmap"
 	"tailscale.com/types/persist"
@@ -89,7 +91,7 @@ func TestTKAEnablementFlow(t *testing.T) {
 	// our mock server can communicate.
 	nlPriv := key.NewNLPrivate()
 	key := tka.Key{Kind: tka.Key25519, Public: nlPriv.Public().Verifier(), Votes: 2}
-	a1, genesisAUM, err := tka.Create(&tka.Mem{}, tka.State{
+	a1, genesisAUM, err := tka.Create(tka.ChonkMem(), tka.State{
 		Keys:               []tka.Key{key},
 		DisablementSecrets: [][]byte{bytes.Repeat([]byte{0xa5}, 32)},
 	}, nlPriv)
@@ -399,7 +401,7 @@ func TestTKASync(t *testing.T) {
 
 			// Setup the tka authority on the control plane.
 			key := tka.Key{Kind: tka.Key25519, Public: nlPriv.Public().Verifier(), Votes: 2}
-			controlStorage := &tka.Mem{}
+			controlStorage := tka.ChonkMem()
 			controlAuthority, bootstrap, err := tka.Create(controlStorage, tka.State{
 				Keys:               []tka.Key{key, someKey},
 				DisablementSecrets: [][]byte{tka.DisablementKDF(disablementSecret)},
@@ -548,10 +550,226 @@ func TestTKASync(t *testing.T) {
 	}
 }
 
+// Whenever we run a TKA sync and get new state from control, we compact the
+// local state.
+func TestTKASyncTriggersCompact(t *testing.T) {
+	someKeyPriv := key.NewNLPrivate()
+	someKey := tka.Key{Kind: tka.Key25519, Public: someKeyPriv.Public().Verifier(), Votes: 1}
+
+	disablementSecret := bytes.Repeat([]byte{0xa5}, 32)
+
+	nodePriv := key.NewNode()
+	nlPriv := key.NewNLPrivate()
+	pm := must.Get(newProfileManager(new(mem.Store), t.Logf, health.NewTracker(eventbustest.NewBus(t))))
+	must.Do(pm.SetPrefs((&ipn.Prefs{
+		Persist: &persist.Persist{
+			PrivateNodeKey: nodePriv,
+			NetworkLockKey: nlPriv,
+		},
+	}).View(), ipn.NetworkProfile{}))
+
+	// Create a clock, and roll it back by 30 days.
+	//
+	// Our compaction algorithm preserves AUMs received in the last 14 days, so
+	// we need to backdate the commit times to make the AUMs eligible for compaction.
+	clock := tstest.NewClock(tstest.ClockOpts{})
+	clock.Advance(-30 * 24 * time.Hour)
+
+	// Set up the TKA authority on the control plane.
+	key := tka.Key{Kind: tka.Key25519, Public: nlPriv.Public().Verifier(), Votes: 2}
+	controlStorage := tka.ChonkMem()
+	controlStorage.SetClock(clock)
+	controlAuthority, bootstrap, err := tka.Create(controlStorage, tka.State{
+		Keys:               []tka.Key{key, someKey},
+		DisablementSecrets: [][]byte{tka.DisablementKDF(disablementSecret)},
+	}, nlPriv)
+	if err != nil {
+		t.Fatalf("tka.Create() failed: %v", err)
+	}
+
+	// Fill the control plane TKA authority with a lot of AUMs, enough so that:
+	//
+	//	1. the chain of AUMs includes some checkpoints
+	//	2. the chain is long enough it would be trimmed if we ran the compaction
+	//	   algorithm with the defaults
+	for range 100 {
+		upd := controlAuthority.NewUpdater(nlPriv)
+		if err := upd.RemoveKey(someKey.MustID()); err != nil {
+			t.Fatalf("RemoveKey: %v", err)
+		}
+		if err := upd.AddKey(someKey); err != nil {
+			t.Fatalf("AddKey: %v", err)
+		}
+		aums, err := upd.Finalize(controlStorage)
+		if err != nil {
+			t.Fatalf("Finalize: %v", err)
+		}
+		if err := controlAuthority.Inform(controlStorage, aums); err != nil {
+			t.Fatalf("controlAuthority.Inform() failed: %v", err)
+		}
+	}
+
+	// Set up the TKA authority on the node.
+	nodeStorage := tka.ChonkMem()
+	nodeStorage.SetClock(clock)
+	nodeAuthority, err := tka.Bootstrap(nodeStorage, bootstrap)
+	if err != nil {
+		t.Fatalf("tka.Bootstrap() failed: %v", err)
+	}
+
+	// Make a mock control server.
+	ts, client := fakeNoiseServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer r.Body.Close()
+		switch r.URL.Path {
+		case "/machine/tka/sync/offer":
+			body := new(tailcfg.TKASyncOfferRequest)
+			if err := json.NewDecoder(r.Body).Decode(body); err != nil {
+				t.Fatal(err)
+			}
+			t.Logf("got sync offer:\n%+v", body)
+			nodeOffer, err := toSyncOffer(body.Head, body.Ancestors)
+			if err != nil {
+				t.Fatal(err)
+			}
+			controlOffer, err := controlAuthority.SyncOffer(controlStorage)
+			if err != nil {
+				t.Fatal(err)
+			}
+			sendAUMs, err := controlAuthority.MissingAUMs(controlStorage, nodeOffer)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			head, ancestors, err := fromSyncOffer(controlOffer)
+			if err != nil {
+				t.Fatal(err)
+			}
+			resp := tailcfg.TKASyncOfferResponse{
+				Head:        head,
+				Ancestors:   ancestors,
+				MissingAUMs: make([]tkatype.MarshaledAUM, len(sendAUMs)),
+			}
+			for i, a := range sendAUMs {
+				resp.MissingAUMs[i] = a.Serialize()
+			}
+
+			t.Logf("responding to sync offer with:\n%+v", resp)
+			w.WriteHeader(200)
+			if err := json.NewEncoder(w).Encode(resp); err != nil {
+				t.Fatal(err)
+			}
+
+		case "/machine/tka/sync/send":
+			body := new(tailcfg.TKASyncSendRequest)
+			if err := json.NewDecoder(r.Body).Decode(body); err != nil {
+				t.Fatal(err)
+			}
+			t.Logf("got sync send:\n%+v", body)
+
+			var remoteHead tka.AUMHash
+			if err := remoteHead.UnmarshalText([]byte(body.Head)); err != nil {
+				t.Fatalf("head unmarshal: %v", err)
+			}
+			toApply := make([]tka.AUM, len(body.MissingAUMs))
+			for i, a := range body.MissingAUMs {
+				if err := toApply[i].Unserialize(a); err != nil {
+					t.Fatalf("decoding missingAUM[%d]: %v", i, err)
+				}
+			}
+
+			if len(toApply) > 0 {
+				if err := controlAuthority.Inform(controlStorage, toApply); err != nil {
+					t.Fatalf("control.Inform(%+v) failed: %v", toApply, err)
+				}
+			}
+			head, err := controlAuthority.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)
+			w.WriteHeader(404)
+		}
+	}))
+	defer ts.Close()
+
+	// Setup the client.
+	cc, _ := fakeControlClient(t, client)
+	b := LocalBackend{
+		cc:     cc,
+		ccAuto: cc,
+		logf:   t.Logf,
+		pm:     pm,
+		store:  pm.Store(),
+		tka: &tkaState{
+			authority: nodeAuthority,
+			storage:   nodeStorage,
+		},
+	}
+
+	// Trigger a sync.
+	err = b.tkaSyncIfNeeded(&netmap.NetworkMap{
+		TKAEnabled: true,
+		TKAHead:    controlAuthority.Head(),
+	}, pm.CurrentPrefs())
+	if err != nil {
+		t.Errorf("tkaSyncIfNeeded() failed: %v", err)
+	}
+
+	// Add a new AUM in control.
+	upd := controlAuthority.NewUpdater(nlPriv)
+	if err := upd.RemoveKey(someKey.MustID()); err != nil {
+		t.Fatalf("RemoveKey: %v", err)
+	}
+	aums, err := upd.Finalize(controlStorage)
+	if err != nil {
+		t.Fatalf("Finalize: %v", err)
+	}
+	if err := controlAuthority.Inform(controlStorage, aums); err != nil {
+		t.Fatalf("controlAuthority.Inform() failed: %v", err)
+	}
+
+	// Run a second sync, which should trigger a compaction.
+	err = b.tkaSyncIfNeeded(&netmap.NetworkMap{
+		TKAEnabled: true,
+		TKAHead:    controlAuthority.Head(),
+	}, pm.CurrentPrefs())
+	if err != nil {
+		t.Errorf("tkaSyncIfNeeded() failed: %v", err)
+	}
+
+	// Check that the node and control plane are in sync.
+	if nodeHead, controlHead := b.tka.authority.Head(), controlAuthority.Head(); nodeHead != controlHead {
+		t.Errorf("node head = %v, want %v", nodeHead, controlHead)
+	}
+
+	// Check the node has compacted away some of its AUMs; that it has purged some AUMs which
+	// are still kept in the control plane.
+	nodeAUMs, err := b.tka.storage.AllAUMs()
+	if err != nil {
+		t.Errorf("AllAUMs() for node failed: %v", err)
+	}
+	controlAUMS, err := controlStorage.AllAUMs()
+	if err != nil {
+		t.Errorf("AllAUMs() for control failed: %v", err)
+	}
+	if len(nodeAUMs) == len(controlAUMS) {
+		t.Errorf("node has not compacted; it has the same number of AUMs as control (node = control = %d)", len(nodeAUMs))
+	}
+}
+
 func TestTKAFilterNetmap(t *testing.T) {
 	nlPriv := key.NewNLPrivate()
 	nlKey := tka.Key{Kind: tka.Key25519, Public: nlPriv.Public().Verifier(), Votes: 2}
-	storage := &tka.Mem{}
+	storage := tka.ChonkMem()
 	authority, _, err := tka.Create(storage, tka.State{
 		Keys:               []tka.Key{nlKey},
 		DisablementSecrets: [][]byte{bytes.Repeat([]byte{0xa5}, 32)},

tka/builder_test.go

@@ -28,7 +28,7 @@ func TestAuthorityBuilderAddKey(t *testing.T) {
 	pub, priv := testingKey25519(t, 1)
 	key := Key{Kind: Key25519, Public: pub, Votes: 2}
 
-	storage := &Mem{}
+	storage := ChonkMem()
 	a, _, err := Create(storage, State{
 		Keys:               []Key{key},
 		DisablementSecrets: [][]byte{DisablementKDF([]byte{1, 2, 3})},
@@ -62,7 +62,7 @@ func TestAuthorityBuilderMaxKey(t *testing.T) {
 	pub, priv := testingKey25519(t, 1)
 	key := Key{Kind: Key25519, Public: pub, Votes: 2}
 
-	storage := &Mem{}
+	storage := ChonkMem()
 	a, _, err := Create(storage, State{
 		Keys:               []Key{key},
 		DisablementSecrets: [][]byte{DisablementKDF([]byte{1, 2, 3})},
@@ -109,7 +109,7 @@ func TestAuthorityBuilderRemoveKey(t *testing.T) {
 	pub2, _ := testingKey25519(t, 2)
 	key2 := Key{Kind: Key25519, Public: pub2, Votes: 1}
 
-	storage := &Mem{}
+	storage := ChonkMem()
 	a, _, err := Create(storage, State{
 		Keys:               []Key{key, key2},
 		DisablementSecrets: [][]byte{DisablementKDF([]byte{1, 2, 3})},
@@ -155,7 +155,7 @@ func TestAuthorityBuilderSetKeyVote(t *testing.T) {
 	pub, priv := testingKey25519(t, 1)
 	key := Key{Kind: Key25519, Public: pub, Votes: 2}
 
-	storage := &Mem{}
+	storage := ChonkMem()
 	a, _, err := Create(storage, State{
 		Keys:               []Key{key},
 		DisablementSecrets: [][]byte{DisablementKDF([]byte{1, 2, 3})},
@@ -191,7 +191,7 @@ func TestAuthorityBuilderSetKeyMeta(t *testing.T) {
 	pub, priv := testingKey25519(t, 1)
 	key := Key{Kind: Key25519, Public: pub, Votes: 2, Meta: map[string]string{"a": "b"}}
 
-	storage := &Mem{}
+	storage := ChonkMem()
 	a, _, err := Create(storage, State{
 		Keys:               []Key{key},
 		DisablementSecrets: [][]byte{DisablementKDF([]byte{1, 2, 3})},
@@ -227,7 +227,7 @@ func TestAuthorityBuilderMultiple(t *testing.T) {
 	pub, priv := testingKey25519(t, 1)
 	key := Key{Kind: Key25519, Public: pub, Votes: 2}
 
-	storage := &Mem{}
+	storage := ChonkMem()
 	a, _, err := Create(storage, State{
 		Keys:               []Key{key},
 		DisablementSecrets: [][]byte{DisablementKDF([]byte{1, 2, 3})},
@@ -275,7 +275,7 @@ func TestAuthorityBuilderCheckpointsAfterXUpdates(t *testing.T) {
 	pub, priv := testingKey25519(t, 1)
 	key := Key{Kind: Key25519, Public: pub, Votes: 2}
 
-	storage := &Mem{}
+	storage := ChonkMem()
 	a, _, err := Create(storage, State{
 		Keys:               []Key{key},
 		DisablementSecrets: [][]byte{DisablementKDF([]byte{1, 2, 3})},

tka/chaintest_test.go

@@ -285,25 +285,25 @@ func (c *testChain) makeAUM(v *testchainNode) AUM {
 
 // Chonk returns a tailchonk containing all AUMs.
 func (c *testChain) Chonk() Chonk {
-	var out Mem
+	out := ChonkMem()
 	for _, update := range c.AUMs {
 		if err := out.CommitVerifiedAUMs([]AUM{update}); err != nil {
 			panic(err)
 		}
 	}
-	return &out
+	return out
 }
 
 // ChonkWith returns a tailchonk containing the named AUMs.
 func (c *testChain) ChonkWith(names ...string) Chonk {
-	var out Mem
+	out := ChonkMem()
 	for _, name := range names {
 		update := c.AUMs[name]
 		if err := out.CommitVerifiedAUMs([]AUM{update}); err != nil {
 			panic(err)
 		}
 	}
-	return &out
+	return out
 }
 
 type testchainOpt struct {

tka/key_test.go

@@ -72,7 +72,7 @@ func TestNLPrivate(t *testing.T) {
 	// Test that key.NLPrivate implements Signer by making a new
 	// authority.
 	k := Key{Kind: Key25519, Public: pub.Verifier(), Votes: 1}
-	_, aum, err := Create(&Mem{}, State{
+	_, aum, err := Create(ChonkMem(), State{
 		Keys:               []Key{k},
 		DisablementSecrets: [][]byte{bytes.Repeat([]byte{1}, 32)},
 	}, p)