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>
This commit is contained in:
Tom DNetto 2022-07-07 11:54:04 -07:00 committed by Tom
parent e37167b3ef
commit cca25f6107
2 changed files with 286 additions and 0 deletions

158
tka/tailchonk.go Normal file
View File

@ -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
}

128
tka/tailchonk_test.go Normal file
View File

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