tka: don't look up parent/child information from purged AUMs

We soft-delete AUMs when they're purged, but when we call `ChildAUMs()`,
we look up soft-deleted AUMs to find the `Children` field.

This patch changes the behaviour of `ChildAUMs()` so it only looks at
not-deleted AUMs. This means we don't need to record child information
on AUMs any more, which is a minor space saving for any newly-recorded
AUMs.

Updates https://github.com/tailscale/tailscale/issues/17566
Updates https://github.com/tailscale/corp/issues/27166

Signed-off-by: Alex Chan <alexc@tailscale.com>

tka/tailchonk.go

@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"slices"
 	"sync"
 	"time"
 
@@ -206,10 +207,14 @@ func ChonkDir(dir string) (*FS, error) {
 // CBOR was chosen because we are already using it and it serializes
 // much smaller than JSON for AUMs. The 'keyasint' thing isn't essential
 // but again it saves a bunch of bytes.
+//
+// We have removed the following fields from fsHashInfo, but they may be
+// present in data stored in existing deployments. Do not reuse these values,
+// to avoid getting unexpected values from legacy data:
+//   - cbor:1, Children
 type fsHashInfo struct {
-	Children    []AUMHash `cbor:"1,keyasint"`
+	AUM         *AUM  `cbor:"2,keyasint"`
-	AUM         *AUM      `cbor:"2,keyasint"`
+	CreatedUnix int64 `cbor:"3,keyasint,omitempty"`
-	CreatedUnix int64     `cbor:"3,keyasint,omitempty"`
 
 	// PurgedUnix is set when the AUM is deleted. The value is
 	// the unix epoch at the time it was deleted.
@@ -285,32 +290,15 @@ func (c *FS) ChildAUMs(prevAUMHash AUMHash) ([]AUM, error) {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
 
-	info, err := c.get(prevAUMHash)
+	var out []AUM
-	if err != nil {
-		if os.IsNotExist(err) {
-			// not knowing about this hash is not an error
-			return nil, nil
-		}
-		return nil, err
-	}
-	// NOTE(tom): We don't check PurgedUnix here because 'purged'
-	// only applies to that specific AUM (i.e. info.AUM) and not to
-	// any information about children stored against that hash.
 
-	out := make([]AUM, len(info.Children))
+	err := c.scanHashes(func(info *fsHashInfo) {
-	for i, h := range info.Children {
+		if info.AUM != nil && bytes.Equal(info.AUM.PrevAUMHash, prevAUMHash[:]) {
-		c, err := c.get(h)
+			out = append(out, *info.AUM)
-		if err != nil {
-			// We expect any AUM recorded as a child on its parent to exist.
-			return nil, fmt.Errorf("reading child %d of %x: %v", i, h, err)
-		}
-		if c.AUM == nil || c.PurgedUnix > 0 {
-			return nil, fmt.Errorf("child %d of %x: AUM not stored", i, h)
 		}
-		out[i] = *c.AUM
+	})
-	}
 
-	return out, nil
+	return out, err
 }
 
 func (c *FS) get(h AUMHash) (*fsHashInfo, error) {
@@ -346,13 +334,45 @@ func (c *FS) Heads() ([]AUM, error) {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
 
+	// Scan the complete list of AUMs, and build a list of all parent hashes.
+	// This tells us which AUMs have children.
+	var parentHashes []AUMHash
+
+	allAUMs, err := c.AllAUMs()
+	if err != nil {
+		return nil, err
+	}
+
+	for _, h := range allAUMs {
+		aum, err := c.AUM(h)
+		if err != nil {
+			return nil, err
+		}
+		parent, hasParent := aum.Parent()
+		if !hasParent {
+			continue
+		}
+		if !slices.Contains(parentHashes, parent) {
+			parentHashes = append(parentHashes, parent)
+		}
+	}
+
+	// Now scan a second time, and only include AUMs which weren't marked as
+	// the parent of any other AUM.
 	out := make([]AUM, 0, 6) // 6 is arbitrary.
-	err := c.scanHashes(func(info *fsHashInfo) {
+
-		if len(info.Children) == 0 && info.AUM != nil && info.PurgedUnix == 0 {
+	for _, h := range allAUMs {
-			out = append(out, *info.AUM)
+		if slices.Contains(parentHashes, h) {
+			continue
 		}
-	})
+		aum, err := c.AUM(h)
-	return out, err
+		if err != nil {
+			return nil, err
+		}
+		out = append(out, aum)
+	}
+
+	return out, nil
 }
 
 // AllAUMs returns all AUMs stored in the chonk.
@@ -362,7 +382,7 @@ func (c *FS) AllAUMs() ([]AUMHash, error) {
 
 	out := make([]AUMHash, 0, 6) // 6 is arbitrary.
 	err := c.scanHashes(func(info *fsHashInfo) {
-		if info.AUM != nil && info.PurgedUnix == 0 {
+		if info.AUM != nil {
 			out = append(out, info.AUM.Hash())
 		}
 	})
@@ -391,6 +411,9 @@ func (c *FS) scanHashes(eachHashInfo func(*fsHashInfo)) error {
 			if err != nil {
 				return fmt.Errorf("reading %x: %v", h, err)
 			}
+			if info.PurgedUnix > 0 {
+				continue
+			}
 
 			eachHashInfo(info)
 		}
@@ -445,24 +468,6 @@ func (c *FS) CommitVerifiedAUMs(updates []AUM) error {
 
 	for i, aum := range updates {
 		h := aum.Hash()
-		// We keep track of children against their parent so that
-		// ChildAUMs() do not need to scan all AUMs.
-		parent, hasParent := aum.Parent()
-		if hasParent {
-			err := c.commit(parent, func(info *fsHashInfo) {
-				// Only add it if its not already there.
-				for i := range info.Children {
-					if info.Children[i] == h {
-						return
-					}
-				}
-				info.Children = append(info.Children, h)
-			})
-			if err != nil {
-				return fmt.Errorf("committing update[%d] to parent %x: %v", i, parent, err)
-			}
-		}
-
 		err := c.commit(h, func(info *fsHashInfo) {
 			info.PurgedUnix = 0 // just in-case it was set for some reason
 			info.AUM = &aum

tka/tailchonk_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"golang.org/x/crypto/blake2s"
+	"tailscale.com/util/must"
 )
 
 // randHash derives a fake blake2s hash from the test name
@@ -144,9 +145,6 @@ func TestTailchonkFS_Commit(t *testing.T) {
 	if _, err := os.Stat(filepath.Join(dir, base)); err != nil {
 		t.Errorf("stat of AUM file failed: %v", err)
 	}
-	if _, err := os.Stat(filepath.Join(chonk.base, "M7", "M7LL2NDB4NKCZIUPVS6RDM2GUOIMW6EEAFVBWMVCPUANQJPHT3SQ")); err != nil {
-		t.Errorf("stat of AUM parent failed: %v", err)
-	}
 
 	info, err := chonk.get(aum.Hash())
 	if err != nil {
@@ -199,6 +197,14 @@ func TestTailchonkFS_PurgeAUMs(t *testing.T) {
 	}
 }
 
+func hashesLess(x, y AUMHash) bool {
+	return bytes.Compare(x[:], y[:]) < 0
+}
+
+func aumHashesLess(x, y AUM) bool {
+	return hashesLess(x.Hash(), y.Hash())
+}
+
 func TestTailchonkFS_AllAUMs(t *testing.T) {
 	chonk := &FS{base: t.TempDir()}
 	genesis := AUM{MessageKind: AUMRemoveKey, KeyID: []byte{1, 2}}
@@ -220,14 +226,54 @@ func TestTailchonkFS_AllAUMs(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	hashesLess := func(a, b AUMHash) bool {
-		return bytes.Compare(a[:], b[:]) < 0
-	}
 	if diff := cmp.Diff([]AUMHash{genesis.Hash(), intermediate.Hash(), leaf.Hash()}, hashes, cmpopts.SortSlices(hashesLess)); diff != "" {
 		t.Fatalf("AllAUMs() output differs (-want, +got):\n%s", diff)
 	}
 }
 
+func TestTailchonkFS_ChildAUMsOfPurgedAUM(t *testing.T) {
+	chonk := &FS{base: t.TempDir()}
+	parent := AUM{MessageKind: AUMRemoveKey, KeyID: []byte{0, 0}}
+
+	parentHash := parent.Hash()
+
+	child1 := AUM{MessageKind: AUMAddKey, KeyID: []byte{1, 1}, PrevAUMHash: parentHash[:]}
+	child2 := AUM{MessageKind: AUMAddKey, KeyID: []byte{2, 2}, PrevAUMHash: parentHash[:]}
+	child3 := AUM{MessageKind: AUMAddKey, KeyID: []byte{3, 3}, PrevAUMHash: parentHash[:]}
+
+	child2Hash := child2.Hash()
+	grandchild2A := AUM{MessageKind: AUMAddKey, KeyID: []byte{2, 2, 2, 2}, PrevAUMHash: child2Hash[:]}
+	grandchild2B := AUM{MessageKind: AUMAddKey, KeyID: []byte{2, 2, 2, 2, 2}, PrevAUMHash: child2Hash[:]}
+
+	commitSet := []AUM{parent, child1, child2, child3, grandchild2A, grandchild2B}
+
+	if err := chonk.CommitVerifiedAUMs(commitSet); err != nil {
+		t.Fatalf("CommitVerifiedAUMs failed: %v", err)
+	}
+
+	// Check the set of hashes is correct
+	childHashes := must.Get(chonk.ChildAUMs(parentHash))
+	if diff := cmp.Diff([]AUM{child1, child2, child3}, childHashes, cmpopts.SortSlices(aumHashesLess)); diff != "" {
+		t.Fatalf("ChildAUMs() output differs (-want, +got):\n%s", diff)
+	}
+
+	// Purge the parent AUM, and check the set of child AUMs is unchanged
+	chonk.PurgeAUMs([]AUMHash{parent.Hash()})
+
+	childHashes = must.Get(chonk.ChildAUMs(parentHash))
+	if diff := cmp.Diff([]AUM{child1, child2, child3}, childHashes, cmpopts.SortSlices(aumHashesLess)); diff != "" {
+		t.Fatalf("ChildAUMs() output differs (-want, +got):\n%s", diff)
+	}
+
+	// Now purge one of the child AUMs, and check it no longer appears as a child of the parent
+	chonk.PurgeAUMs([]AUMHash{child3.Hash()})
+
+	childHashes = must.Get(chonk.ChildAUMs(parentHash))
+	if diff := cmp.Diff([]AUM{child1, child2}, childHashes, cmpopts.SortSlices(aumHashesLess)); diff != "" {
+		t.Fatalf("ChildAUMs() output differs (-want, +got):\n%s", diff)
+	}
+}
+
 func TestMarkActiveChain(t *testing.T) {
 	type aumTemplate struct {
 		AUM AUM
@@ -585,10 +631,7 @@ func (c *compactingChonkFake) CommitTime(hash AUMHash) (time.Time, error) {
 }
 
 func (c *compactingChonkFake) PurgeAUMs(hashes []AUMHash) error {
-	lessHashes := func(a, b AUMHash) bool {
+	if diff := cmp.Diff(c.wantDelete, hashes, cmpopts.SortSlices(hashesLess)); diff != "" {
-		return bytes.Compare(a[:], b[:]) < 0
-	}
-	if diff := cmp.Diff(c.wantDelete, hashes, cmpopts.SortSlices(lessHashes)); diff != "" {
 		c.t.Errorf("deletion set differs (-want, +got):\n%s", diff)
 	}
 	return nil