tka: implement Chonk type & an in-memory implementation

Chonks are responsible for efficient storage of AUMs and other TKA state.

For testing/prototyping I've implemented an in-memory version, but once we
start to use this from tailscaled we'll need a file-based version.

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

tka/tailchonk.go

@@ -0,0 +1,158 @@
+// 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 (
+	"os"
+	"sync"
+)
+
+// Chonk implementations provide durable storage for AUMs and other
+// TKA state.
+//
+// All methods must be thread-safe.
+//
+// The name 'tailchonk' was coined by @catzkorn.
+type Chonk interface {
+	// AUM returns the AUM with the specified digest.
+	//
+	// If the AUM does not exist, then os.ErrNotExist is returned.
+	AUM(hash AUMHash) (AUM, error)
+
+	// ChildAUMs returns all AUMs with a specified previous
+	// AUM hash.
+	ChildAUMs(prevAUMHash AUMHash) ([]AUM, error)
+
+	// CommitVerifiedAUMs durably stores the provided AUMs.
+	// Callers MUST ONLY provide AUMs which are verified (specifically,
+	// a call to aumVerify() must return a nil error).
+	// as the implementation assumes that only verified AUMs are stored.
+	CommitVerifiedAUMs(updates []AUM) error
+
+	// Heads returns AUMs for which there are no children. In other
+	// words, the latest AUM in all possible chains (the 'leaves').
+	Heads() ([]AUM, error)
+
+	// SetLastActiveAncestor is called to record the oldest-known AUM
+	// that contributed to the current state. This value is used as
+	// a hint on next startup to determine which chain to pick when computing
+	// the current state, if there are multiple distinct chains.
+	SetLastActiveAncestor(hash AUMHash) error
+
+	// LastActiveAncestor returns the oldest-known AUM that was (in a
+	// previous run) an ancestor of the current state. This is used
+	// as a hint to pick the correct chain in the event that the Chonk stores
+	// multiple distinct chains.
+	LastActiveAncestor() (*AUMHash, error)
+}
+
+// Mem implements in-memory storage of TKA state, suitable for
+// tests.
+//
+// Mem implements the Chonk interface.
+type Mem struct {
+	l           sync.RWMutex
+	aums        map[AUMHash]AUM
+	parentIndex map[AUMHash][]AUMHash
+
+	lastActiveAncestor *AUMHash
+}
+
+func (c *Mem) SetLastActiveAncestor(hash AUMHash) error {
+	c.l.Lock()
+	defer c.l.Unlock()
+	c.lastActiveAncestor = &hash
+	return nil
+}
+
+func (c *Mem) LastActiveAncestor() (*AUMHash, error) {
+	c.l.RLock()
+	defer c.l.RUnlock()
+	return c.lastActiveAncestor, nil
+}
+
+// Heads returns AUMs for which there are no children. In other
+// words, the latest AUM in all chains (the 'leaf').
+func (c *Mem) Heads() ([]AUM, error) {
+	c.l.RLock()
+	defer c.l.RUnlock()
+	out := make([]AUM, 0, 6)
+
+	// An AUM is a 'head' if there are no nodes for which it is the parent.
+	for _, a := range c.aums {
+		if len(c.parentIndex[a.Hash()]) == 0 {
+			out = append(out, a)
+		}
+	}
+	return out, nil
+}
+
+// AUM returns the AUM with the specified digest.
+func (c *Mem) AUM(hash AUMHash) (AUM, error) {
+	c.l.RLock()
+	defer c.l.RUnlock()
+	aum, ok := c.aums[hash]
+	if !ok {
+		return AUM{}, os.ErrNotExist
+	}
+	return aum, nil
+}
+
+// Orphans returns all AUMs which do not have a parent.
+func (c *Mem) Orphans() ([]AUM, error) {
+	c.l.RLock()
+	defer c.l.RUnlock()
+	out := make([]AUM, 0, 6)
+	for _, a := range c.aums {
+		if _, ok := a.Parent(); !ok {
+			out = append(out, a)
+		}
+	}
+	return out, nil
+}
+
+// ChildAUMs returns all AUMs with a specified previous
+// AUM hash.
+func (c *Mem) ChildAUMs(prevAUMHash AUMHash) ([]AUM, error) {
+	c.l.RLock()
+	defer c.l.RUnlock()
+	out := make([]AUM, 0, 6)
+	for _, entry := range c.parentIndex[prevAUMHash] {
+		out = append(out, c.aums[entry])
+	}
+
+	return out, nil
+}
+
+// CommitVerifiedAUMs durably stores the provided AUMs.
+// Callers MUST ONLY provide well-formed and verified AUMs,
+// as the rest of the TKA implementation assumes that only
+// verified AUMs are stored.
+func (c *Mem) CommitVerifiedAUMs(updates []AUM) error {
+	c.l.Lock()
+	defer c.l.Unlock()
+	if c.aums == nil {
+		c.parentIndex = make(map[AUMHash][]AUMHash, 64)
+		c.aums = make(map[AUMHash]AUM, 64)
+	}
+
+updateLoop:
+	for _, aum := range updates {
+		aumHash := aum.Hash()
+		c.aums[aumHash] = aum
+
+		parent, ok := aum.Parent()
+		if ok {
+			for _, exists := range c.parentIndex[parent] {
+				if exists == aumHash {
+					continue updateLoop
+				}
+			}
+			c.parentIndex[parent] = append(c.parentIndex[parent], aumHash)
+		}
+	}
+
+	return nil
+}

tka/tailchonk_test.go

@@ -0,0 +1,128 @@
+// 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 (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"golang.org/x/crypto/blake2s"
+)
+
+// randHash derives a fake blake2s hash from the test name
+// and the given seed.
+func randHash(t *testing.T, seed int64) [blake2s.Size]byte {
+	var out [blake2s.Size]byte
+	testingRand(t, seed).Read(out[:])
+	return out
+}
+
+func TestImplementsChonk(t *testing.T) {
+	impls := []Chonk{&Mem{}}
+	t.Logf("chonks: %v", impls)
+}
+
+func TestTailchonkMem_ChildAUMs(t *testing.T) {
+	chonk := Mem{}
+	parentHash := randHash(t, 1)
+	data := []AUM{
+		{
+			MessageKind: AUMRemoveKey,
+			KeyID:       []byte{1, 2},
+			PrevAUMHash: parentHash[:],
+		},
+		{
+			MessageKind: AUMRemoveKey,
+			KeyID:       []byte{3, 4},
+			PrevAUMHash: parentHash[:],
+		},
+	}
+
+	if err := chonk.CommitVerifiedAUMs(data); err != nil {
+		t.Fatalf("CommitVerifiedAUMs failed: %v", err)
+	}
+	stored, err := chonk.ChildAUMs(parentHash)
+	if err != nil {
+		t.Fatalf("ChildAUMs failed: %v", err)
+	}
+	if diff := cmp.Diff(data, stored); diff != "" {
+		t.Errorf("stored AUM differs (-want, +got):\n%s", diff)
+	}
+}
+
+func TestTailchonkMem_Orphans(t *testing.T) {
+	chonk := Mem{}
+
+	parentHash := randHash(t, 1)
+	orphan := AUM{MessageKind: AUMNoOp}
+	aums := []AUM{
+		orphan,
+		// A parent is specified, so we shouldnt see it in GetOrphans()
+		{
+			MessageKind: AUMRemoveKey,
+			KeyID:       []byte{3, 4},
+			PrevAUMHash: parentHash[:],
+		},
+	}
+	if err := chonk.CommitVerifiedAUMs(aums); err != nil {
+		t.Fatalf("CommitVerifiedAUMs failed: %v", err)
+	}
+
+	stored, err := chonk.Orphans()
+	if err != nil {
+		t.Fatalf("Orphans failed: %v", err)
+	}
+	if diff := cmp.Diff([]AUM{orphan}, stored); diff != "" {
+		t.Errorf("stored AUM differs (-want, +got):\n%s", diff)
+	}
+}
+
+func TestTailchonkMem_ReadChainFromHead(t *testing.T) {
+	chonk := Mem{}
+	genesis := AUM{MessageKind: AUMRemoveKey, KeyID: []byte{1, 2}}
+	gHash := genesis.Hash()
+	intermediate := AUM{PrevAUMHash: gHash[:]}
+	iHash := intermediate.Hash()
+	leaf := AUM{PrevAUMHash: iHash[:]}
+
+	commitSet := []AUM{
+		genesis,
+		intermediate,
+		leaf,
+	}
+	if err := chonk.CommitVerifiedAUMs(commitSet); err != nil {
+		t.Fatalf("CommitVerifiedAUMs failed: %v", err)
+	}
+	// t.Logf("genesis hash = %X", genesis.Hash())
+	// t.Logf("intermediate hash = %X", intermediate.Hash())
+	// t.Logf("leaf hash = %X", leaf.Hash())
+
+	// Read the chain from the leaf backwards.
+	gotLeafs, err := chonk.Heads()
+	if err != nil {
+		t.Fatalf("Heads failed: %v", err)
+	}
+	if diff := cmp.Diff([]AUM{leaf}, gotLeafs); diff != "" {
+		t.Fatalf("leaf AUM differs (-want, +got):\n%s", diff)
+	}
+
+	parent, _ := gotLeafs[0].Parent()
+	gotIntermediate, err := chonk.AUM(parent)
+	if err != nil {
+		t.Fatalf("AUM(<intermediate>) failed: %v", err)
+	}
+	if diff := cmp.Diff(intermediate, gotIntermediate); diff != "" {
+		t.Errorf("intermediate AUM differs (-want, +got):\n%s", diff)
+	}
+
+	parent, _ = gotIntermediate.Parent()
+	gotGenesis, err := chonk.AUM(parent)
+	if err != nil {
+		t.Fatalf("AUM(<genesis>) failed: %v", err)
+	}
+	if diff := cmp.Diff(genesis, gotGenesis); diff != "" {
+		t.Errorf("genesis AUM differs (-want, +got):\n%s", diff)
+	}
+}