tka: implement State and applying AUMs

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

5 files changed, 677 insertions, 42 deletions

diff --git a/tka/aum.go b/tka/aum.go
index 88166a4c5..7d0897cd8 100644
--- a/tka/aum.go
+++ b/tka/aum.go
@@ -7,6 +7,7 @@ package tka
 import (
 	"bytes"
 	"crypto/ed25519"
+	"encoding/binary"
 	"errors"
 	"fmt"
 
@@ -104,8 +105,7 @@ type AUM struct {
 
 	// State describes the full state of the key authority.
 	// This field is used for Checkpoint AUMs.
-	// TODO(tom): Use type *State once a future PR brings in that type.
-	State interface{} `cbor:"5,keyasint,omitempty"`
+	State *State `cbor:"5,keyasint,omitempty"`
 
 	// DisablementSecret is used to transmit a secret for disabling
 	// the TKA.
@@ -122,6 +122,13 @@ type AUM struct {
 	Signatures []Signature `cbor:"23,keyasint,omitempty"`
 }
 
+// Upper bound on checkpoint elements, chosen arbitrarily. Intended to
+// cap out insanely large AUMs.
+const (
+	maxDisablementSecrets = 32
+	maxKeys               = 512
+)
+
 // StaticValidate returns a nil error if the AUM is well-formed.
 func (a *AUM) StaticValidate() error {
 	if a.Key != nil {
@@ -138,7 +145,36 @@ func (a *AUM) StaticValidate() error {
 		}
 	}
 
-	// TODO(tom): Validate State once a future PR brings in that type.
+	if a.State != nil {
+		if len(a.State.LastAUMHash) != 0 {
+			return errors.New("checkpoint state cannot specify a parent AUM")
+		}
+		if len(a.State.DisablementSecrets) == 0 {
+			return errors.New("at least one disablement secret required")
+		}
+		if numDS := len(a.State.DisablementSecrets); numDS > maxDisablementSecrets {
+			return fmt.Errorf("too many disablement secrets (%d, max %d)", numDS, maxDisablementSecrets)
+		}
+		for i, ds := range a.State.DisablementSecrets {
+			if len(ds) != disablementLength {
+				return fmt.Errorf("disablement[%d]: invalid length (got %d, want %d)", i, len(ds), disablementLength)
+			}
+		}
+		// TODO(tom): Check for duplicate disablement secrets.
+
+		if len(a.State.Keys) == 0 {
+			return errors.New("at least one key is required")
+		}
+		if numKeys := len(a.State.Keys); numKeys > maxKeys {
+			return fmt.Errorf("too many keys (%d, max %d)", numKeys, maxKeys)
+		}
+		for i, k := range a.State.Keys {
+			if err := k.StaticValidate(); err != nil {
+				return fmt.Errorf("key[%d]: %v", i, err)
+			}
+		}
+		// TODO(tom): Check for duplicate keys.
+	}
 
 	switch a.MessageKind {
 	case AUMAddKey:
@@ -253,4 +289,50 @@ func (a *AUM) sign25519(priv ed25519.PrivateKey) {
 	})
 }
 
-// TODO(tom): Implement Weight() once a future PR brings in the State type.
+// Weight computes the 'signature weight' of the AUM
+// based on keys in the state machine. The caller must
+// ensure that all signatures are valid.
+//
+// More formally: W = Sum(key.votes)
+//
+// AUMs with a higher weight than their siblings
+// are preferred when resolving forks in the AUM chain.
+func (a *AUM) Weight(state State) uint {
+	var weight uint
+
+	// Track the keys that have already been used, so two
+	// signatures with the same key do not result in 2x
+	// the weight.
+	//
+	// We use the first 8 bytes as the key for this map,
+	// because KeyIDs are either a blake2s hash or
+	// the 25519 public key, both of which approximate
+	// random distribution.
+	seenKeys := make(map[uint64]struct{}, 6)
+	for _, sig := range a.Signatures {
+		if len(sig.KeyID) < 8 {
+			// Invalid, don't count it
+			continue
+		}
+
+		keyID := binary.LittleEndian.Uint64(sig.KeyID)
+
+		key, err := state.GetKey(sig.KeyID)
+		if err != nil {
+			if err == ErrNoSuchKey {
+				// Signatures with an unknown key do not contribute
+				// to the weight.
+				continue
+			}
+			panic(err)
+		}
+		if _, seen := seenKeys[keyID]; seen {
+			continue
+		}
+
+		weight += key.Votes
+		seenKeys[keyID] = struct{}{}
+	}
+
+	return weight
+}
diff --git a/tka/aum_test.go b/tka/aum_test.go
index ad8e5b971..6fb14cc3e 100644
--- a/tka/aum_test.go
+++ b/tka/aum_test.go
@@ -10,10 +10,12 @@ import (
 
 	"github.com/fxamacker/cbor/v2"
 	"github.com/google/go-cmp/cmp"
+	"golang.org/x/crypto/blake2s"
 )
 
 func TestSerialization(t *testing.T) {
 	uint2 := uint(2)
+	var fakeAUMHash AUMHash
 
 	tcs := []struct {
 		Name   string
@@ -94,44 +96,45 @@ func TestSerialization(t *testing.T) {
 				0x04, //    |- major type 0 (int), value 4 (byte 4)
 			},
 		},
-		// TODO(tom): Uncomment once a future PR brings in the State type.
-		// {
-		// 	"Checkpoint",
-		// 	AUM{MessageKind: AUMCheckpoint, PrevAUMHash: []byte{1, 2}, State: &State{
-		// 		LastAUMHash: []byte{3, 4},
-		// 		Keys: []Key{
-		// 			{Kind: Key25519, Public: []byte{5, 6}},
-		// 		},
-		// 	}},
-		// 	[]byte{
-		// 		0xa3, // major type 5 (map), 3 items
-		// 		0x01, // |- major type 0 (int), value 1 (first key, MessageKind)
-		// 		0x06, // |- major type 0 (int), value 6 (first value, AUMCheckpoint)
-		// 		0x02, // |- major type 0 (int), value 2 (second key, PrevAUMHash)
-		// 		0x42, // |- major type 2 (byte string), 2 items (second value)
-		// 		0x01, //    |- major type 0 (int), value 1 (byte 1)
-		// 		0x02, //    |- major type 0 (int), value 2 (byte 2)
-		// 		0x05, // |- major type 0 (int), value 5 (third key, State)
-		// 		0xa3, // |- major type 5 (map), 3 items (third value, State type)
-		// 		0x01, //    |- major type 0 (int), value 1 (first key, LastAUMHash)
-		// 		0x42, //    |- major type 2 (byte string), 2 items (first value)
-		// 		0x03, //       |- major type 0 (int), value 3 (byte 3)
-		// 		0x04, //    |- major type 0 (int), value 4 (byte 4)
-		// 		0x02, //    |- major type 0 (int), value 2 (second key, DisablementSecrets)
-		// 		0xf6, //    |- major type 7 (val), value null (second value, nil)
-		// 		0x03, //    |- major type 0 (int), value 3 (third key, Keys)
-		// 		0x81, //    |- major type 4 (array), value 1 (one item in array)
-		// 		0xa3, //       |- major type 5 (map), 3 items (Key type)
-		// 		0x01, //          |- major type 0 (int), value 1 (first key, Kind)
-		// 		0x01, //          |- major type 0 (int), value 1 (first value, Key25519)
-		// 		0x02, //          |- major type 0 (int), value 2 (second key, Votes)
-		// 		0x00, //          |- major type 0 (int), value 0 (second value, 0)
-		// 		0x03, //          |- major type 0 (int), value 3 (third key, Public)
-		// 		0x42, //          |- major type 2 (byte string), 2 items (third value)
-		// 		0x05, //             |- major type 0 (int), value 5 (byte 5)
-		// 		0x06, //             |- major type 0 (int), value 6 (byte 6)
-		// 	},
-		// },
+		{
+			"Checkpoint",
+			AUM{MessageKind: AUMCheckpoint, PrevAUMHash: []byte{1, 2}, State: &State{
+				LastAUMHash: &fakeAUMHash,
+				Keys: []Key{
+					{Kind: Key25519, Public: []byte{5, 6}},
+				},
+			}},
+			append(
+				append([]byte{
+					0xa3,       // major type 5 (map), 3 items
+					0x01,       // |- major type 0 (int), value 1 (first key, MessageKind)
+					0x06,       // |- major type 0 (int), value 6 (first value, AUMCheckpoint)
+					0x02,       // |- major type 0 (int), value 2 (second key, PrevAUMHash)
+					0x42,       // |- major type 2 (byte string), 2 items (second value)
+					0x01,       //    |- major type 0 (int), value 1 (byte 1)
+					0x02,       //    |- major type 0 (int), value 2 (byte 2)
+					0x05,       // |- major type 0 (int), value 5 (third key, State)
+					0xa3,       // |- major type 5 (map), 3 items (third value, State type)
+					0x01,       //   |- major type 0 (int), value 1 (first key, LastAUMHash)
+					0x58, 0x20, //   |- major type 2 (byte string), 32 items (first value)
+				},
+					bytes.Repeat([]byte{0}, 32)...),
+				[]byte{
+					0x02, //     |- major type 0 (int), value 2 (second key, DisablementSecrets)
+					0xf6, //     |- major type 7 (val), value null (second value, nil)
+					0x03, //     |- major type 0 (int), value 3 (third key, Keys)
+					0x81, //     |- major type 4 (array), value 1 (one item in array)
+					0xa3, //       |- major type 5 (map), 3 items (Key type)
+					0x01, //          |- major type 0 (int), value 1 (first key, Kind)
+					0x01, //          |- major type 0 (int), value 1 (first value, Key25519)
+					0x02, //          |- major type 0 (int), value 2 (second key, Votes)
+					0x00, //          |- major type 0 (int), value 0 (second value, 0)
+					0x03, //          |- major type 0 (int), value 3 (third key, Public)
+					0x42, //          |- major type 2 (byte string), 2 items (third value)
+					0x05, //             |- major type 0 (int), value 5 (byte 5)
+					0x06, //             |- major type 0 (int), value 6 (byte 6)
+				}...),
+		},
 		{
 			"Signature",
 			AUM{MessageKind: AUMAddKey, Signatures: []Signature{{KeyID: []byte{1}}}},
@@ -171,6 +174,77 @@ func TestSerialization(t *testing.T) {
 	}
 }
 
+func TestAUMWeight(t *testing.T) {
+	var fakeKeyID [blake2s.Size]byte
+	testingRand(t, 1).Read(fakeKeyID[:])
+
+	pub, _ := testingKey25519(t, 1)
+	key := Key{Kind: Key25519, Public: pub, Votes: 2}
+	pub, _ = testingKey25519(t, 2)
+	key2 := Key{Kind: Key25519, Public: pub, Votes: 2}
+
+	tcs := []struct {
+		Name  string
+		AUM   AUM
+		State State
+		Want  uint
+	}{
+		{
+			"Empty",
+			AUM{},
+			State{},
+			0,
+		},
+		{
+			"Key unknown",
+			AUM{
+				Signatures: []Signature{{KeyID: fakeKeyID[:]}},
+			},
+			State{},
+			0,
+		},
+		{
+			"Unary key",
+			AUM{
+				Signatures: []Signature{{KeyID: key.ID()}},
+			},
+			State{
+				Keys: []Key{key},
+			},
+			2,
+		},
+		{
+			"Multiple keys",
+			AUM{
+				Signatures: []Signature{{KeyID: key.ID()}, {KeyID: key2.ID()}},
+			},
+			State{
+				Keys: []Key{key, key2},
+			},
+			4,
+		},
+		{
+			"Double use",
+			AUM{
+				Signatures: []Signature{{KeyID: key.ID()}, {KeyID: key.ID()}},
+			},
+			State{
+				Keys: []Key{key},
+			},
+			2,
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.Name, func(t *testing.T) {
+			got := tc.AUM.Weight(tc.State)
+			if got != tc.Want {
+				t.Errorf("Weight() = %d, want %d", got, tc.Want)
+			}
+		})
+	}
+}
+
 func TestAUMHashes(t *testing.T) {
 	// .Hash(): a hash over everything.
 	// .SigHash(): a hash over everything except the signatures.
diff --git a/tka/key.go b/tka/key.go
index eabdccdfb..be275ecd9 100644
--- a/tka/key.go
+++ b/tka/key.go
@@ -51,6 +51,29 @@ type Key struct {
 	Meta map[string]string `cbor:"12,keyasint,omitempty"`
 }
 
+// Clone makes an independent copy of Key.
+//
+// NOTE: There is a difference between a nil slice and an empty
+//       slice for encoding purposes, so an implementation of Clone()
+//       must take care to preserve this.
+func (k Key) Clone() Key {
+	out := k
+
+	if k.Public != nil {
+		out.Public = make([]byte, len(k.Public))
+		copy(out.Public, k.Public)
+	}
+
+	if k.Meta != nil {
+		out.Meta = make(map[string]string, len(k.Meta))
+		for k, v := range k.Meta {
+			out.Meta[k] = v
+		}
+	}
+
+	return out
+}
+
 func (k Key) ID() KeyID {
 	switch k.Kind {
 	// Because 25519 public keys are so short, we just use the 32-byte
diff --git a/tka/state.go b/tka/state.go
new file mode 100644
index 000000000..52bb7eb8f
--- /dev/null
+++ b/tka/state.go
@@ -0,0 +1,204 @@
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package tka
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+
+	"golang.org/x/crypto/argon2"
+)
+
+// ErrNoSuchKey is returned if the key referenced by a KeyID does not exist.
+var ErrNoSuchKey = errors.New("key not found")
+
+// State describes Tailnet Key Authority state at an instant in time.
+//
+// State is mutated by applying Authority Update Messages (AUMs), resulting
+// in a new State.
+type State struct {
+	// LastAUMHash is the blake2s digest of the last-applied AUM.
+	// Because AUMs are strictly ordered and form a hash chain, we
+	// check the previous AUM hash in an update we are applying
+	// is the same as the LastAUMHash.
+	LastAUMHash *AUMHash `cbor:"1,keyasint"`
+
+	// DisablementSecrets are KDF-derived values which can be used
+	// to turn off the TKA in the event of a consensus-breaking bug.
+	// An AUM of type DisableNL should contain a secret when results
+	// in one of these values when run through the disablement KDF.
+	//
+	// TODO(tom): This is an alpha feature, remove this mechanism once
+	//            we have confidence in our implementation.
+	DisablementSecrets [][]byte `cbor:"2,keyasint"`
+
+	// Keys are the public keys currently trusted by the TKA.
+	Keys []Key `cbor:"3,keyasint"`
+}
+
+// GetKey returns the trusted key with the specified KeyID.
+func (s State) GetKey(key KeyID) (Key, error) {
+	for _, k := range s.Keys {
+		if bytes.Equal(k.ID(), key) {
+			return k, nil
+		}
+	}
+
+	return Key{}, ErrNoSuchKey
+}
+
+// Clone makes an independent copy of State.
+//
+// NOTE: There is a difference between a nil slice and an empty
+//       slice for encoding purposes, so an implementation of Clone()
+//       must take care to preserve this.
+func (s State) Clone() State {
+	out := State{}
+
+	if s.LastAUMHash != nil {
+		dupe := *s.LastAUMHash
+		out.LastAUMHash = &dupe
+	}
+
+	if s.DisablementSecrets != nil {
+		out.DisablementSecrets = make([][]byte, len(s.DisablementSecrets))
+		for i := range s.DisablementSecrets {
+			out.DisablementSecrets[i] = make([]byte, len(s.DisablementSecrets[i]))
+			copy(out.DisablementSecrets[i], s.DisablementSecrets[i])
+		}
+	}
+
+	if s.Keys != nil {
+		out.Keys = make([]Key, len(s.Keys))
+		for i := range s.Keys {
+			out.Keys[i] = s.Keys[i].Clone()
+		}
+	}
+
+	return out
+}
+
+// cloneForUpdate is like Clone, except LastAUMHash is set based
+// on the hash of the given update.
+func (s State) cloneForUpdate(update *AUM) State {
+	out := s.Clone()
+	aumHash := update.Hash()
+	out.LastAUMHash = &aumHash
+	return out
+}
+
+const disablementLength = 32
+
+var disablementSalt = []byte("tailscale network-lock disablement salt")
+
+func disablementKDF(secret []byte) []byte {
+	// time = 4 (3 recommended, booped to 4 to compensate for less memory)
+	// memory = 16 (32 recommended)
+	// threads = 4
+	// keyLen = 32 (256 bits)
+	return argon2.Key(secret, disablementSalt, 4, 16*1024, 4, disablementLength)
+}
+
+// checkDisablement returns true for a valid disablement secret.
+func (s State) checkDisablement(secret []byte) bool {
+	derived := disablementKDF(secret)
+	for _, candidate := range s.DisablementSecrets {
+		if bytes.Equal(derived, candidate) {
+			return true
+		}
+	}
+	return false
+}
+
+// parentMatches returns true if an AUM can chain to (be applied)
+// to the current state.
+//
+// Specifically, the rules are:
+//  - The last AUM hash must match (transitively, this implies that this
+//    update follows the last update message applied to the state machine)
+//  - Or, the state machine knows no parent (its brand new).
+func (s State) parentMatches(update AUM) bool {
+	if s.LastAUMHash == nil {
+		return true
+	}
+	return bytes.Equal(s.LastAUMHash[:], update.PrevAUMHash)
+}
+
+// applyVerifiedAUM computes a new state based on the update provided.
+//
+// The provided update MUST be verified: That is, the AUM must be well-formed
+// (as defined by StaticValidate()), and signatures over the AUM must have
+// been verified.
+func (s State) applyVerifiedAUM(update AUM) (State, error) {
+	// Validate that the update message has the right parent.
+	if !s.parentMatches(update) {
+		return State{}, errors.New("parent AUMHash mismatch")
+	}
+
+	switch update.MessageKind {
+	case AUMNoOp:
+		out := s.cloneForUpdate(&update)
+		return out, nil
+
+	case AUMCheckpoint:
+		return update.State.cloneForUpdate(&update), nil
+
+	case AUMAddKey:
+		if _, err := s.GetKey(update.Key.ID()); err == nil {
+			return State{}, errors.New("key already exists")
+		}
+		out := s.cloneForUpdate(&update)
+		out.Keys = append(out.Keys, *update.Key)
+		return out, nil
+
+	case AUMUpdateKey:
+		k, err := s.GetKey(update.KeyID)
+		if err != nil {
+			return State{}, err
+		}
+		if update.Votes != nil {
+			k.Votes = *update.Votes
+		}
+		if update.Meta != nil {
+			k.Meta = update.Meta
+		}
+		out := s.cloneForUpdate(&update)
+		for i := range out.Keys {
+			if bytes.Equal(out.Keys[i].ID(), update.KeyID) {
+				out.Keys[i] = k
+			}
+		}
+		return out, nil
+
+	case AUMRemoveKey:
+		idx := -1
+		for i := range s.Keys {
+			if bytes.Equal(update.KeyID, s.Keys[i].ID()) {
+				idx = i
+				break
+			}
+		}
+		if idx < 0 {
+			return State{}, ErrNoSuchKey
+		}
+		out := s.cloneForUpdate(&update)
+		out.Keys = append(out.Keys[:idx], out.Keys[idx+1:]...)
+		return out, nil
+
+	case AUMDisableNL:
+		// TODO(tom): We should handle this at a higher level than State.
+		if !s.checkDisablement(update.DisablementSecret) {
+			return State{}, errors.New("incorrect disablement secret")
+		}
+		// Valid disablement secret, lets reset
+		return State{}, nil
+
+	default:
+		// TODO(tom): Instead of erroring, update lastHash and
+		// continue (to preserve future compatibility).
+		return State{}, fmt.Errorf("unhandled message: %v", update.MessageKind)
+	}
+}
diff --git a/tka/state_test.go b/tka/state_test.go
new file mode 100644
index 000000000..d2aa16b25
--- /dev/null
+++ b/tka/state_test.go
@@ -0,0 +1,252 @@
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package tka
+
+import (
+	"bytes"
+	"encoding/hex"
+	"errors"
+	"testing"
+
+	"github.com/fxamacker/cbor/v2"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+)
+
+func fromHex(in string) []byte {
+	out, err := hex.DecodeString(in)
+	if err != nil {
+		panic(err)
+	}
+	return out
+}
+
+func hashFromHex(in string) *AUMHash {
+	var out AUMHash
+	copy(out[:], fromHex(in))
+	return &out
+}
+
+func TestCloneState(t *testing.T) {
+	tcs := []struct {
+		Name  string
+		State State
+	}{
+		{
+			"Empty",
+			State{},
+		},
+		{
+			"Key",
+			State{
+				Keys: []Key{{Kind: Key25519, Votes: 2, Public: []byte{5, 6, 7, 8}, Meta: map[string]string{"a": "b"}}},
+			},
+		},
+		{
+			"DisablementSecrets",
+			State{
+				DisablementSecrets: [][]byte{
+					{1, 2, 3, 4},
+					{5, 6, 7, 8},
+				},
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.Name, func(t *testing.T) {
+			if diff := cmp.Diff(tc.State, tc.State.Clone()); diff != "" {
+				t.Errorf("output state differs (-want, +got):\n%s", diff)
+			}
+
+			// Make sure the cloned State is the same even after
+			// an encode + decode into + from CBOR.
+			t.Run("cbor", func(t *testing.T) {
+				out := bytes.NewBuffer(nil)
+				encoder, err := cbor.CTAP2EncOptions().EncMode()
+				if err != nil {
+					t.Fatal(err)
+				}
+				if err := encoder.NewEncoder(out).Encode(tc.State.Clone()); err != nil {
+					t.Fatal(err)
+				}
+
+				var decodedState State
+				if err := cbor.Unmarshal(out.Bytes(), &decodedState); err != nil {
+					t.Fatalf("Unmarshal failed: %v", err)
+				}
+				if diff := cmp.Diff(tc.State, decodedState); diff != "" {
+					t.Errorf("decoded state differs (-want, +got):\n%s", diff)
+				}
+			})
+		})
+	}
+}
+
+func TestApplyUpdatesChain(t *testing.T) {
+	intOne := uint(1)
+	tcs := []struct {
+		Name    string
+		Updates []AUM
+		Start   State
+		End     State
+	}{
+		{
+			"AddKey",
+			[]AUM{{MessageKind: AUMAddKey, Key: &Key{Kind: Key25519, Public: []byte{1, 2, 3, 4}}}},
+			State{},
+			State{
+				Keys:        []Key{{Kind: Key25519, Public: []byte{1, 2, 3, 4}}},
+				LastAUMHash: hashFromHex("53898e4311d0b6087fcbb871563868a16c629d9267df851fcfa7b52b31d2bd03"),
+			},
+		},
+		{
+			"RemoveKey",
+			[]AUM{{MessageKind: AUMRemoveKey, KeyID: []byte{1, 2, 3, 4}, PrevAUMHash: fromHex("53898e4311d0b6087fcbb871563868a16c629d9267df851fcfa7b52b31d2bd03")}},
+			State{
+				Keys:        []Key{{Kind: Key25519, Public: []byte{1, 2, 3, 4}}},
+				LastAUMHash: hashFromHex("53898e4311d0b6087fcbb871563868a16c629d9267df851fcfa7b52b31d2bd03"),
+			},
+			State{
+				LastAUMHash: hashFromHex("15d65756abfafbb592279503f40759898590c9c59056be1e2e9f02684c15ba4b"),
+			},
+		},
+		{
+			"UpdateKey",
+			[]AUM{{MessageKind: AUMUpdateKey, KeyID: []byte{1, 2, 3, 4}, Votes: &intOne, Meta: map[string]string{"a": "b"}, PrevAUMHash: fromHex("53898e4311d0b6087fcbb871563868a16c629d9267df851fcfa7b52b31d2bd03")}},
+			State{
+				Keys:        []Key{{Kind: Key25519, Public: []byte{1, 2, 3, 4}}},
+				LastAUMHash: hashFromHex("53898e4311d0b6087fcbb871563868a16c629d9267df851fcfa7b52b31d2bd03"),
+			},
+			State{
+				LastAUMHash: hashFromHex("828fe04c16032cf3e0b021abca0b4d79924b0a18b2e627b308347aa87ce7c21c"),
+				Keys:        []Key{{Kind: Key25519, Votes: 1, Meta: map[string]string{"a": "b"}, Public: []byte{1, 2, 3, 4}}},
+			},
+		},
+		{
+			"ChainedKeyUpdates",
+			[]AUM{
+				{MessageKind: AUMAddKey, Key: &Key{Kind: Key25519, Public: []byte{5, 6, 7, 8}}},
+				{MessageKind: AUMRemoveKey, KeyID: []byte{1, 2, 3, 4}, PrevAUMHash: fromHex("f09bda3bb7cf6756ea9adc25770aede4b3ca8142949d6ef5ca0add29af912fd4")},
+			},
+			State{
+				Keys: []Key{{Kind: Key25519, Public: []byte{1, 2, 3, 4}}},
+			},
+			State{
+				Keys:        []Key{{Kind: Key25519, Public: []byte{5, 6, 7, 8}}},
+				LastAUMHash: hashFromHex("218165fe5f757304b9deaff4ac742890364f5f509e533c74e80e0ce35e44ee1d"),
+			},
+		},
+		{
+			"Disablement",
+			[]AUM{{MessageKind: AUMDisableNL, DisablementSecret: []byte{1, 2, 3, 4}, PrevAUMHash: fromHex("53898e4311d0b6087fcbb871563868a16c629d9267df851fcfa7b52b31d2bd03")}},
+			State{
+				DisablementSecrets: [][]byte{disablementKDF([]byte{1, 2, 3, 4})},
+				LastAUMHash:        hashFromHex("53898e4311d0b6087fcbb871563868a16c629d9267df851fcfa7b52b31d2bd03"),
+			},
+			State{},
+		},
+		{
+			"Checkpoint",
+			[]AUM{
+				{MessageKind: AUMAddKey, Key: &Key{Kind: Key25519, Public: []byte{5, 6, 7, 8}}},
+				{MessageKind: AUMCheckpoint, State: &State{
+					Keys: []Key{{Kind: Key25519, Public: []byte{1, 2, 3, 4}}},
+				}, PrevAUMHash: fromHex("f09bda3bb7cf6756ea9adc25770aede4b3ca8142949d6ef5ca0add29af912fd4")},
+			},
+			State{DisablementSecrets: [][]byte{[]byte{1, 2, 3, 4}}},
+			State{
+				Keys:        []Key{{Kind: Key25519, Public: []byte{1, 2, 3, 4}}},
+				LastAUMHash: hashFromHex("2e34f7e21883c35c8e34ec06e735f7ed8a14c3ceeb11ccb18fcbc11d51c8dabb"),
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.Name, func(t *testing.T) {
+			state := tc.Start
+			for i := range tc.Updates {
+				var err error
+				// t.Logf("update[%d] start-state = %+v", i, state)
+				state, err = state.applyVerifiedAUM(tc.Updates[i])
+				if err != nil {
+					t.Fatalf("Apply message[%d] failed: %v", i, err)
+				}
+				// t.Logf("update[%d] end-state = %+v", i, state)
+
+				updateHash := tc.Updates[i].Hash()
+				if tc.Updates[i].MessageKind != AUMDisableNL {
+					if got, want := *state.LastAUMHash, updateHash[:]; !bytes.Equal(got[:], want) {
+						t.Errorf("expected state.LastAUMHash = %x (update %d), got %x", want, i, got)
+					}
+				}
+			}
+
+			if diff := cmp.Diff(tc.End, state, cmpopts.EquateEmpty()); diff != "" {
+				t.Errorf("output state differs (+got, -want):\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestApplyUpdateErrors(t *testing.T) {
+	tcs := []struct {
+		Name    string
+		Updates []AUM
+		Start   State
+		Error   error
+	}{
+		{
+			"AddKey exists",
+			[]AUM{{MessageKind: AUMAddKey, Key: &Key{Kind: Key25519, Public: []byte{1, 2, 3, 4}}}},
+			State{Keys: []Key{{Kind: Key25519, Public: []byte{1, 2, 3, 4}}}},
+			errors.New("key already exists"),
+		},
+		{
+			"RemoveKey notfound",
+			[]AUM{{MessageKind: AUMRemoveKey, Key: &Key{Kind: Key25519, Public: []byte{1, 2, 3, 4}}}},
+			State{},
+			ErrNoSuchKey,
+		},
+		{
+			"UpdateKey notfound",
+			[]AUM{{MessageKind: AUMUpdateKey, KeyID: []byte{1}}},
+			State{},
+			ErrNoSuchKey,
+		},
+		{
+			"Bad lastAUMHash",
+			[]AUM{
+				{MessageKind: AUMAddKey, Key: &Key{Kind: Key25519, Public: []byte{5, 6, 7, 8}}},
+				{MessageKind: AUMRemoveKey, KeyID: []byte{1, 2, 3, 4}, PrevAUMHash: fromHex("1234")},
+			},
+			State{
+				Keys: []Key{{Kind: Key25519, Public: []byte{1, 2, 3, 4}}},
+			},
+			errors.New("parent AUMHash mismatch"),
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.Name, func(t *testing.T) {
+			state := tc.Start
+			for i := range tc.Updates {
+				var err error
+				// t.Logf("update[%d] start-state = %+v", i, state)
+				state, err = state.applyVerifiedAUM(tc.Updates[i])
+				if err != nil {
+					if err.Error() != tc.Error.Error() {
+						t.Errorf("state[%d].Err = %v, want %v", i, err, tc.Error)
+					} else {
+						return
+					}
+				}
+				// t.Logf("update[%d] end-state = %+v", i, state)
+			}
+
+			t.Errorf("did not error, expected %v", tc.Error)
+		})
+	}
+}