tka: Checkpoint every 50 updates

Signed-off-by: Tom DNetto <tom@tailscale.com>

tka/aum.go

@@ -50,7 +50,6 @@ func (h AUMHash) IsZero() bool {
 	return h == (AUMHash{})
 }
 
-
 // AUMKind describes valid AUM types.
 type AUMKind uint8
 

tka/builder.go

@@ -6,6 +6,7 @@ package tka
 
 import (
 	"fmt"
+	"os"
 
 	"tailscale.com/types/tkatype"
 )
@@ -92,8 +93,62 @@ func (b *UpdateBuilder) SetKeyMeta(keyID tkatype.KeyID, meta map[string]string)
 	return b.mkUpdate(AUM{MessageKind: AUMUpdateKey, Meta: meta, KeyID: keyID})
 }
 
+func (b *UpdateBuilder) generateCheckpoint() error {
+	// Compute the checkpoint state.
+	state := b.a.state
+	for i, update := range b.out {
+		var err error
+		if state, err = state.applyVerifiedAUM(update); err != nil {
+			return fmt.Errorf("applying update %d: %v", i, err)
+		}
+	}
+
+	// Checkpoints cant specify a parent AUM.
+	state.LastAUMHash = nil
+	return b.mkUpdate(AUM{MessageKind: AUMCheckpoint, State: &state})
+}
+
+// checkpointEvery sets how often a checkpoint AUM should be generated.
+const checkpointEvery = 50
+
 // Finalize returns the set of update message to actuate the update.
-func (b *UpdateBuilder) Finalize() ([]AUM, error) {
+func (b *UpdateBuilder) Finalize(storage Chonk) ([]AUM, error) {
+	var (
+		needCheckpoint bool    = true
+		cursor         AUMHash = b.a.Head()
+	)
+	for i := len(b.out); i < checkpointEvery; i++ {
+		aum, err := storage.AUM(cursor)
+		if err != nil {
+			if err == os.ErrNotExist {
+				// The available chain is shorter than the interval to checkpoint at.
+				needCheckpoint = false
+				break
+			}
+			return nil, fmt.Errorf("reading AUM: %v", err)
+		}
+
+		if aum.MessageKind == AUMCheckpoint {
+			needCheckpoint = false
+			break
+		}
+
+		parent, hasParent := aum.Parent()
+		if !hasParent {
+			// We've hit the genesis update, so the chain is shorter than the interval to checkpoint at.
+			needCheckpoint = false
+			break
+		}
+		cursor = parent
+	}
+
+	if needCheckpoint {
+		if err := b.generateCheckpoint(); err != nil {
+			return nil, fmt.Errorf("generating checkpoint: %v", err)
+		}
+	}
+
+	// Check no AUMs were applied in the meantime
 	if len(b.out) > 0 {
 		if parent, _ := b.out[0].Parent(); parent != b.a.Head() {
 			return nil, fmt.Errorf("updates no longer apply to head: based on %x but head is %x", parent, b.a.Head())

tka/builder_test.go

@@ -44,7 +44,7 @@ func TestAuthorityBuilderAddKey(t *testing.T) {
 	if err := b.AddKey(key2); err != nil {
 		t.Fatalf("AddKey(%v) failed: %v", key2, err)
 	}
-	updates, err := b.Finalize()
+	updates, err := b.Finalize(storage)
 	if err != nil {
 		t.Fatalf("Finalize() failed: %v", err)
 	}
@@ -78,7 +78,7 @@ func TestAuthorityBuilderRemoveKey(t *testing.T) {
 	if err := b.RemoveKey(key2.ID()); err != nil {
 		t.Fatalf("RemoveKey(%v) failed: %v", key2, err)
 	}
-	updates, err := b.Finalize()
+	updates, err := b.Finalize(storage)
 	if err != nil {
 		t.Fatalf("Finalize() failed: %v", err)
 	}
@@ -110,7 +110,7 @@ func TestAuthorityBuilderSetKeyVote(t *testing.T) {
 	if err := b.SetKeyVote(key.ID(), 5); err != nil {
 		t.Fatalf("SetKeyVote(%v) failed: %v", key.ID(), err)
 	}
-	updates, err := b.Finalize()
+	updates, err := b.Finalize(storage)
 	if err != nil {
 		t.Fatalf("Finalize() failed: %v", err)
 	}
@@ -146,7 +146,7 @@ func TestAuthorityBuilderSetKeyMeta(t *testing.T) {
 	if err := b.SetKeyMeta(key.ID(), map[string]string{"b": "c"}); err != nil {
 		t.Fatalf("SetKeyMeta(%v) failed: %v", key, err)
 	}
-	updates, err := b.Finalize()
+	updates, err := b.Finalize(storage)
 	if err != nil {
 		t.Fatalf("Finalize() failed: %v", err)
 	}
@@ -191,7 +191,7 @@ func TestAuthorityBuilderMultiple(t *testing.T) {
 	if err := b.RemoveKey(key.ID()); err != nil {
 		t.Fatalf("RemoveKey(%v) failed: %v", key, err)
 	}
-	updates, err := b.Finalize()
+	updates, err := b.Finalize(storage)
 	if err != nil {
 		t.Fatalf("Finalize() failed: %v", err)
 	}
@@ -212,3 +212,60 @@ func TestAuthorityBuilderMultiple(t *testing.T) {
 		t.Errorf("GetKey(key).err = %v, want %v", err, ErrNoSuchKey)
 	}
 }
+
+func TestAuthorityBuilderCheckpointsAfterXUpdates(t *testing.T) {
+	pub, priv := testingKey25519(t, 1)
+	key := Key{Kind: Key25519, Public: pub, Votes: 2}
+
+	storage := &Mem{}
+	a, _, err := Create(storage, State{
+		Keys:               []Key{key},
+		DisablementSecrets: [][]byte{DisablementKDF([]byte{1, 2, 3})},
+	}, signer25519(priv))
+	if err != nil {
+		t.Fatalf("Create() failed: %v", err)
+	}
+
+	for i := 0; i <= checkpointEvery; i++ {
+		pub2, _ := testingKey25519(t, int64(i+2))
+		key2 := Key{Kind: Key25519, Public: pub2, Votes: 1}
+
+		b := a.NewUpdater(signer25519(priv))
+		if err := b.AddKey(key2); err != nil {
+			t.Fatalf("AddKey(%v) failed: %v", key2, err)
+		}
+		updates, err := b.Finalize(storage)
+		if err != nil {
+			t.Fatalf("Finalize() failed: %v", err)
+		}
+		// See if the update is valid by applying it to the authority
+		// + checking if the new key is there.
+		if err := a.Inform(storage, updates); err != nil {
+			t.Fatalf("could not apply generated updates: %v", err)
+		}
+		if _, err := a.state.GetKey(key2.ID()); err != nil {
+			t.Fatal(err)
+		}
+
+		wantKind := AUMAddKey
+		if i == checkpointEvery-1 { // Genesis + 49 updates == 50 (the value of checkpointEvery)
+			wantKind = AUMCheckpoint
+		}
+		lastAUM, err := storage.AUM(a.Head())
+		if err != nil {
+			t.Fatal(err)
+		}
+		if lastAUM.MessageKind != wantKind {
+			t.Errorf("[%d] HeadAUM.MessageKind = %v, want %v", i, lastAUM.MessageKind, wantKind)
+		}
+	}
+
+	// Try starting an authority just based on storage.
+	a2, err := Open(storage)
+	if err != nil {
+		t.Fatalf("Failed to open from stored AUMs: %v", err)
+	}
+	if a.Head() != a2.Head() {
+		t.Errorf("stored and computed HEAD differ: got %v, want %v", a2.Head(), a.Head())
+	}
+}

tka/tka.go

@@ -10,6 +10,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"reflect"
 	"sort"
 
 	"github.com/fxamacker/cbor/v2"
@@ -181,6 +182,18 @@ func advanceByPrimary(state State, candidates []AUM) (next *AUM, out State, err
 	}
 
 	aum := pickNextAUM(state, candidates)
+
+	// TODO(tom): Remove this before GA, this is just a correctness check during implementation.
+	// Post-GA, we want clients to not error if they dont recognize additional fields in State.
+	if aum.MessageKind == AUMCheckpoint {
+		dupe := state
+		dupe.LastAUMHash = nil
+		// aum.State is non-nil (see aum.StaticValidate).
+		if !reflect.DeepEqual(dupe, *aum.State) {
+			return nil, State{}, errors.New("checkpoint includes changes not represented in earlier AUMs")
+		}
+	}
+
 	if state, err = state.applyVerifiedAUM(aum); err != nil {
 		return nil, State{}, fmt.Errorf("advancing state: %v", err)
 	}
@@ -244,10 +257,6 @@ func fastForward(storage Chonk, maxIter int, startState State, done func(curAUM
 
 // computeStateAt returns the State at wantHash.
 func computeStateAt(storage Chonk, maxIter int, wantHash AUMHash) (State, error) {
-	// TODO(tom): This is going to get expensive for really long
-	//            chains. We should make nodes emit a checkpoint every
-	//            X updates or something.
-
 	topAUM, err := storage.AUM(wantHash)
 	if err != nil {
 		return State{}, err