diff --git a/tka/tailchonk.go b/tka/tailchonk.go new file mode 100644 index 000000000..4a5abdadb --- /dev/null +++ b/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 +} diff --git a/tka/tailchonk_test.go b/tka/tailchonk_test.go new file mode 100644 index 000000000..334323275 --- /dev/null +++ b/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() 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() failed: %v", err) + } + if diff := cmp.Diff(genesis, gotGenesis); diff != "" { + t.Errorf("genesis AUM differs (-want, +got):\n%s", diff) + } +}