btrfstree: Rethink 'Path' yet again

5 files changed, 217 insertions, 168 deletions

diff --git a/lib/btrfs/btrfstree/btree.go b/lib/btrfs/btrfstree/btree.go
index 4f5d21b..19c7c68 100644
--- a/lib/btrfs/btrfstree/btree.go
+++ b/lib/btrfs/btrfstree/btree.go
@@ -14,6 +14,17 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 )
 
+type Tree interface {
+	// CheckOwner returns whether it is permissible for a node
+	// with .Head.Owner=owner and .Head.Generation=gen to be in
+	// this tree.
+	//
+	// If there is an error determining this, then `failOpen`
+	// specifies whether it should return an error (false) or nil
+	// (true).
+	TreeCheckOwner(ctx context.Context, failOpen bool, owner btrfsprim.ObjID, gen btrfsprim.Generation) error
+}
+
 type TreeSearcher interface {
 	// How the search should be described in the event of an
 	// error.
diff --git a/lib/btrfs/btrfstree/btree_forrest.go b/lib/btrfs/btrfstree/btree_forrest.go
index b04bfc0..1875c2f 100644
--- a/lib/btrfs/btrfstree/btree_forrest.go
+++ b/lib/btrfs/btrfstree/btree_forrest.go
@@ -22,7 +22,9 @@ type TreeRoot struct {
 	Level      uint8
 	Generation btrfsprim.Generation
 
-	RootInode btrfsprim.ObjID // only for subvolume trees
+	RootInode  btrfsprim.ObjID // only for subvolume trees
+	ParentUUID btrfsprim.UUID
+	ParentGen  btrfsprim.Generation // offset of this tree's root item
 }
 
 // LookupTreeRoot is a utility function to help with implementing the
@@ -72,7 +74,10 @@ func LookupTreeRoot(_ context.Context, fs TreeOperator, sb Superblock, treeID bt
 				RootNode:   rootItemBody.ByteNr,
 				Level:      rootItemBody.Level,
 				Generation: rootItemBody.Generation,
+
 				RootInode:  rootItemBody.RootDirID,
+				ParentUUID: rootItemBody.ParentUUID,
+				ParentGen:  btrfsprim.Generation(rootItem.Key.Offset),
 			}, nil
 		case *btrfsitem.Error:
 			return nil, fmt.Errorf("malformed ROOT_ITEM for tree %v: %w", treeID, rootItemBody.Err)
@@ -83,10 +88,7 @@ func LookupTreeRoot(_ context.Context, fs TreeOperator, sb Superblock, treeID bt
 }
 
 type TreeOperatorImpl struct {
-	NodeSource interface {
-		NodeSource
-		NodeFile
-	}
+	NodeSource NodeSource
 }
 
 func (fs TreeOperatorImpl) RawTree(ctx context.Context, treeID btrfsprim.ObjID) (*RawTree, error) {
@@ -108,7 +110,7 @@ func (fs TreeOperatorImpl) RawTree(ctx context.Context, treeID btrfsprim.ObjID)
 func (fs TreeOperatorImpl) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, errHandle func(*TreeError), cbs TreeWalkHandler) {
 	tree, err := fs.RawTree(ctx, treeID)
 	if err != nil {
-		errHandle(&TreeError{Path: Path{{FromTree: treeID, ToMaxKey: btrfsprim.MaxKey}}, Err: err})
+		errHandle(&TreeError{Path: Path{PathRoot{TreeID: treeID}}, Err: err})
 		return
 	}
 	tree.TreeWalk(ctx, cbs)
diff --git a/lib/btrfs/btrfstree/btree_tree.go b/lib/btrfs/btrfstree/btree_tree.go
index 86eab11..27d6548 100644
--- a/lib/btrfs/btrfstree/btree_tree.go
+++ b/lib/btrfs/btrfstree/btree_tree.go
@@ -22,27 +22,35 @@ type RawTree struct {
 }
 
 func (tree *RawTree) TreeWalk(ctx context.Context, cbs TreeWalkHandler) {
+	sb, err := tree.Forrest.NodeSource.Superblock()
+	if err != nil {
+		if cbs.BadSuperblock != nil {
+			cbs.BadSuperblock(err)
+		}
+		return
+	}
 	if tree.RootNode == 0 {
 		return
 	}
-	path := Path{{
-		FromTree:         tree.ID,
-		FromItemSlot:     -1,
-		ToNodeAddr:       tree.RootNode,
-		ToNodeGeneration: tree.Generation,
-		ToNodeLevel:      tree.Level,
-		ToMaxKey:         btrfsprim.MaxKey,
-	}}
-	tree.walk(ctx, path, cbs)
+	path := Path{
+		PathRoot{
+			Tree:         tree,
+			TreeID:       tree.ID,
+			ToAddr:       tree.RootNode,
+			ToGeneration: tree.Generation,
+			ToLevel:      tree.Level,
+		},
+	}
+	tree.walk(ctx, *sb, path, cbs)
 }
 
-func (tree *RawTree) walk(ctx context.Context, path Path, cbs TreeWalkHandler) {
+func (tree *RawTree) walk(ctx context.Context, sb Superblock, path Path, cbs TreeWalkHandler) {
 	if ctx.Err() != nil {
 		return
 	}
 
 	// 001
-	nodeAddr, nodeExp, ok := path.NodeExpectations(tree.Forrest.NodeSource)
+	nodeAddr, nodeExp, ok := path.NodeExpectations(ctx, true) // TODO(lukeshu): Consider whether failing open is the right thing here
 	if !ok {
 		return
 	}
@@ -74,18 +82,20 @@ func (tree *RawTree) walk(ctx context.Context, path Path, cbs TreeWalkHandler) {
 	}
 	// branch a (interior)
 	for i, item := range node.BodyInterior {
-		toMaxKey := path.Node(-1).ToMaxKey
+		toMaxKey := nodeExp.MaxItem.Val
 		if i+1 < len(node.BodyInterior) {
 			toMaxKey = node.BodyInterior[i+1].Key.Mm()
 		}
-		itemPath := append(path, PathElem{
-			FromTree:         node.Head.Owner,
-			FromItemSlot:     i,
-			ToNodeAddr:       item.BlockPtr,
-			ToNodeGeneration: item.Generation,
-			ToNodeLevel:      node.Head.Level - 1,
-			ToKey:            item.Key,
-			ToMaxKey:         toMaxKey,
+		itemPath := append(path, PathKP{
+			FromTree: node.Head.Owner,
+			FromSlot: i,
+
+			ToAddr:       item.BlockPtr,
+			ToGeneration: item.Generation,
+			ToMinKey:     item.Key,
+
+			ToMaxKey: toMaxKey,
+			ToLevel:  node.Head.Level - 1,
 		})
 		// 003a
 		recurse := cbs.KeyPointer == nil || cbs.KeyPointer(itemPath, item)
@@ -94,7 +104,7 @@ func (tree *RawTree) walk(ctx context.Context, path Path, cbs TreeWalkHandler) {
 		}
 		// 004a
 		if recurse {
-			tree.walk(ctx, itemPath, cbs)
+			tree.walk(ctx, sb, itemPath, cbs)
 			if ctx.Err() != nil {
 				return
 			}
@@ -105,11 +115,11 @@ func (tree *RawTree) walk(ctx context.Context, path Path, cbs TreeWalkHandler) {
 		return
 	}
 	for i, item := range node.BodyLeaf {
-		itemPath := append(path, PathElem{
-			FromTree:     node.Head.Owner,
-			FromItemSlot: i,
-			ToKey:        item.Key,
-			ToMaxKey:     item.Key,
+		itemPath := append(path, PathItem{
+			FromTree: node.Head.Owner,
+			FromSlot: i,
+
+			ToKey: item.Key,
 		})
 		// 003b
 		switch item.Body.(type) {
@@ -271,3 +281,64 @@ func (tree *RawTree) TreeSubrange(ctx context.Context, min int, searcher TreeSea
 	}
 	return nil
 }
+
+// TreeCheckOwner implements the 'Tree' interface.
+func (tree *RawTree) TreeCheckOwner(ctx context.Context, failOpen bool, owner btrfsprim.ObjID, gen btrfsprim.Generation) error {
+	var uuidTree *RawTree
+	for {
+		// Main.
+		if owner == tree.ID {
+			return nil
+		}
+		if tree.ParentUUID == (btrfsprim.UUID{}) {
+			return fmt.Errorf("owner=%v is not acceptable in this tree",
+				owner)
+		}
+		if gen > tree.ParentGen {
+			return fmt.Errorf("claimed owner=%v might be acceptable in this tree (if generation<=%v) but not with claimed generation=%v",
+				owner, tree.ParentGen, gen)
+		}
+
+		// Loop update.
+		if uuidTree == nil {
+			var err error
+			uuidTree, err = tree.Forrest.RawTree(ctx, btrfsprim.UUID_TREE_OBJECTID)
+			if err != nil {
+				if failOpen {
+					return nil
+				}
+				return fmt.Errorf("unable to determine whether owner=%v generation=%v is acceptable: %w",
+					owner, gen, err)
+			}
+		}
+		parentIDItem, err := uuidTree.TreeLookup(ctx, btrfsitem.UUIDToKey(tree.ParentUUID))
+		if err != nil {
+			if failOpen {
+				return nil
+			}
+			return fmt.Errorf("unable to determine whether owner=%v generation=%v is acceptable: %w",
+				owner, gen, err)
+		}
+		switch parentIDBody := parentIDItem.Body.(type) {
+		case *btrfsitem.UUIDMap:
+			tree, err = tree.Forrest.RawTree(ctx, parentIDBody.ObjID)
+			if err != nil {
+				if failOpen {
+					return nil
+				}
+				return fmt.Errorf("unable to determine whether owner=%v generation=%v is acceptable: %w",
+					owner, gen, err)
+			}
+		case *btrfsitem.Error:
+			if failOpen {
+				return nil
+			}
+			return fmt.Errorf("unable to determine whether owner=%v generation=%v is acceptable: %w",
+				owner, gen, parentIDBody.Err)
+		default:
+			// This is a panic because the item decoder should not emit UUID_SUBVOL items as anything but
+			// btrfsitem.UUIDMap or btrfsitem.Error without this code also being updated.
+			panic(fmt.Errorf("should not happen: UUID_SUBVOL item has unexpected type: %T", parentIDBody))
+		}
+	}
+}
diff --git a/lib/btrfs/btrfstree/path.go b/lib/btrfs/btrfstree/path.go
index 1335b76..537df9a 100644
--- a/lib/btrfs/btrfstree/path.go
+++ b/lib/btrfs/btrfstree/path.go
@@ -5,8 +5,8 @@
 package btrfstree
 
 import (
+	"context"
 	"fmt"
-	"io"
 	"strings"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
@@ -14,11 +14,10 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 )
 
-// Path is a path from the superblock (i.e. the root of the btrfs
-// system) to the a node or item within one of the btrees in the
-// system.
+// Path is a path from the superblock or a ROOT_ITEM to a node or
+// item within one of the btrees in the system.
 //
-//   - The first element will always have an ItemSlot of -1.
+//   - The first element will always have a FromSlot of -1.
 //
 //   - For .Item() callbacks, the last element will always have a
 //     NodeAddr of 0.
@@ -27,164 +26,156 @@ import (
 //
 //	[superblock: tree=B, lvl=3, gen=6]
 //	     |
-//	     | <------------------------------------------ pathElem={from_tree:B, from_slot=-1,
-//	     |                                                       to_addr:0x01, to_gen=6, to_lvl=3}
+//	     | <------------------------------------------ PathRoot={Tree:B,
+//	     |                                                       ToAddr:0x01, ToGen=6, ToLvl=3}
 //	  +[0x01]-------------+
 //	  | lvl=3 gen=6 own=B |
 //	  +-+-+-+-+-+-+-+-+-+-+
 //	  |0|1|2|3|4|5|6|7|8|9|
 //	  +-+-+-+-+-+-+-+-+-+-+
 //	                 |
-//	                 | <------------------------------ pathElem:{from_tree:B, from_slot:7,
-//	                 |                                           to_addr:0x02, to_gen:5, to_lvl:2}
+//	                 | <------------------------------ PathKP={FromTree:B, FromSlot:7,
+//	                 |                                         ToAddr:0x02, ToGen:5, ToLvl:2}
 //	              +[0x02]--------------+
 //	              | lvl=2 gen=5 own=B  |
 //	              +-+-+-+-+-+-+-+-+-+-+
 //	              |0|1|2|3|4|5|6|7|8|9|
 //	              +-+-+-+-+-+-+-+-+-+-+
 //	                           |
-//	                           | <-------------------- pathElem={from_tree:B, from_slot:6,
-//	                           |                                 to_addr:0x03, to_gen:5, to_lvl:1}
+//	                           | <-------------------- PathKP={FromTree:B, FromSlot:6,
+//	                           |                               ToAddr:0x03, ToGen:5, ToLvl:1}
 //	                        +[0x03]-------------+
 //	                        | lvl=1 gen=5 own=A |
 //	                        +-+-+-+-+-+-+-+-+-+-+
 //	                        |0|1|2|3|4|5|6|7|8|9|
 //	                        +-+-+-+-+-+-+-+-+-+-+
 //	                               |
-//	                               | <---------------- pathElem={from_tree:A, from_slot:3,
-//	                               |                             to_addr:0x04, to_gen:2, lvl:0}
+//	                               | <---------------- PathKP={FromTree:A, FromSlot:3,
+//	                               |                           ToAddr:0x04, ToGen:2, ToLvl:0}
 //	                             +[0x04]-------------+
 //	                             | lvl=0 gen=2 own=A |
 //	                             +-+-+-+-+-+-+-+-+-+-+
 //	                             |0|1|2|3|4|5|6|7|8|9|
 //	                             +-+-+-+-+-+-+-+-+-+-+
 //	                                |
-//	                                | <--------------- pathElem={from_tree:A, from_slot:1,
-//	                                |                            to_addr:0, to_gen: 0, to_lvl:0}
+//	                                | <--------------- PathItem={FromTree:A, FromSlot:1}
+//	                                |
 //	                              [item]
 type Path []PathElem
 
-// A PathElem essentially represents a KeyPointer.
-type PathElem struct {
-	// FromTree is the owning tree ID of the parent node; or the
-	// well-known tree ID if this is the root.
-	FromTree btrfsprim.ObjID
-	// FromItemSlot is the index of this KeyPointer in the parent
-	// Node; or -1 if this is the root and there is no KeyPointer.
-	FromItemSlot int
+// A PathElem is either a PathRoot, a PathKP, or a PathItem.
+type PathElem interface {
+	isPathElem()
+}
 
-	// ToNodeAddr is the address of the node that the KeyPointer
-	// points at, or 0 if this is a leaf item and nothing is being
-	// pointed at.
-	ToNodeAddr btrfsvol.LogicalAddr
-	// ToNodeGeneration is the expected generation of the node at
-	// ToNodeAddr, or 0 if this is a leaf item and nothing is
-	// being pointed at.
-	ToNodeGeneration btrfsprim.Generation
-	// ToNodeLevel is the expected level of the node at
-	// ToNodeAddr, or 0 if this is a leaf item and nothing is
-	// being pointed at.
-	ToNodeLevel uint8
-	// ToKey is either
-	//  - btrfprim.Key{} if this is the root node being pointed
-	//    to,
-	//  - the KeyPointer.Key if this is a non-root node being
-	//    pointed to, or
-	//  - the key of the leaf item being pointed to.
-	ToKey    btrfsprim.Key
+type PathRoot struct {
+	Tree Tree
+	// It should be no surprise that these 4 members mimic the 4
+	// members of a 'RawTree'.
+	TreeID       btrfsprim.ObjID
+	ToAddr       btrfsvol.LogicalAddr
+	ToGeneration btrfsprim.Generation
+	ToLevel      uint8
+}
+
+func (PathRoot) isPathElem() {}
+
+type PathKP struct {
+	// From the containing Node.
+	FromTree btrfsprim.ObjID
+	FromSlot int
+	// From the KP itself.
+	ToAddr       btrfsvol.LogicalAddr
+	ToGeneration btrfsprim.Generation
+	ToMinKey     btrfsprim.Key
+	// From the structure of the tree.
 	ToMaxKey btrfsprim.Key
+	ToLevel  uint8
 }
 
-func (elem PathElem) writeNodeTo(w io.Writer) {
-	fmt.Fprintf(w, "node:%d@%v", elem.ToNodeLevel, elem.ToNodeAddr)
+func (PathKP) isPathElem() {}
+
+type PathItem struct {
+	// From the containing Node.
+	FromTree btrfsprim.ObjID
+	FromSlot int
+	// From the Item itself.
+	ToKey btrfsprim.Key
 }
 
+func (PathItem) isPathElem() {}
+
 func (path Path) String() string {
 	if len(path) == 0 {
 		return "(empty-path)"
 	}
 	var ret strings.Builder
-	fmt.Fprintf(&ret, "%s->", path[0].FromTree.Format(btrfsprim.ROOT_TREE_OBJECTID))
-	if len(path) == 1 && path[0] == (PathElem{FromTree: path[0].FromTree, FromItemSlot: -1}) {
-		ret.WriteString("(empty-path)")
-	} else {
-		path[0].writeNodeTo(&ret)
-	}
-	for _, elem := range path[1:] {
-		fmt.Fprintf(&ret, "[%v]", elem.FromItemSlot)
-		if elem.ToNodeAddr != 0 {
-			ret.WriteString("->")
-			elem.writeNodeTo(&ret)
+	for _, elem := range path {
+		switch elem := elem.(type) {
+		case PathRoot:
+			fmt.Fprintf(&ret, "%s->node:%d@%v",
+				elem.TreeID.Format(btrfsprim.ROOT_TREE_OBJECTID),
+				elem.ToLevel, elem.ToAddr)
+		case PathKP:
+			fmt.Fprintf(&ret, "[%d]->node:%d@%v",
+				elem.FromSlot,
+				elem.ToLevel, elem.ToAddr)
+		case PathItem:
+			// fmt.Fprintf(&ret, "[%d]->item:%v",
+			// 	elem.FromSlot,
+			// 	elem.ToKey)
+			fmt.Fprintf(&ret, "[%d]",
+				elem.FromSlot)
+		default:
+			panic(fmt.Errorf("should not happen: unexpected PathElem type: %T", elem))
 		}
 	}
 	return ret.String()
 }
 
-func (path Path) DeepCopy() Path {
-	return append(Path(nil), path...)
-}
-
 // NodeExpectations returns the address to read and the expectations
 // to have when reading the node pointed to by this Path.
 //
 // `ok` is false if the path is empty or if this Path points to an
 // item rather than a node.
-func (path Path) NodeExpectations(fs NodeFile) (_ btrfsvol.LogicalAddr, _ NodeExpectations, ok bool) {
-	if path.Node(-1).ToNodeAddr == 0 && path.Node(-1).ToNodeGeneration == 0 && path.Node(-1).ToNodeLevel == 0 {
+func (path Path) NodeExpectations(ctx context.Context, failOpen bool) (_ btrfsvol.LogicalAddr, _ NodeExpectations, ok bool) {
+	if len(path) == 0 {
 		return 0, NodeExpectations{}, false
 	}
-
-	checkOwner := func(owner btrfsprim.ObjID, gen btrfsprim.Generation) error {
-		var treeParents []btrfsprim.ObjID
-
-		tree := path.Node(-1).FromTree
-		for {
-			if owner == tree {
-				// OK!
-				return nil
-			}
-
-			treeParents = append(treeParents, tree)
-			parent, parentGen, parentOK := fs.ParentTree(tree)
-			if !parentOK {
-				// Failed look up parent info; fail open.
-				return nil
-			}
-
-			if parent == 0 {
-				// End of the line.
-				return fmt.Errorf("expected owner in %v but claims to have owner=%v",
-					treeParents, owner)
-			}
-			if gen > parentGen {
-				return fmt.Errorf("claimed owner=%v might be acceptable in this tree (if generation<=%v) but not with claimed generation=%v",
-					owner, parentGen, gen)
-			}
-			tree = parent
-		}
+	firstElem, ok := path[0].(PathRoot)
+	if !ok {
+		panic(fmt.Errorf("should not happen: first PathElem is not PathRoot: %T", path[0]))
+	}
+	switch lastElem := path[len(path)-1].(type) {
+	case PathRoot:
+		return lastElem.ToAddr, NodeExpectations{
+			LAddr:      containers.OptionalValue(lastElem.ToAddr),
+			Level:      containers.OptionalValue(lastElem.ToLevel),
+			Generation: containers.OptionalValue(lastElem.ToGeneration),
+			Owner: func(owner btrfsprim.ObjID, gen btrfsprim.Generation) error {
+				return firstElem.Tree.TreeCheckOwner(ctx, failOpen, owner, gen)
+			},
+			MinItem: containers.OptionalValue(btrfsprim.Key{}),
+			MaxItem: containers.OptionalValue(btrfsprim.MaxKey),
+		}, true
+	case PathKP:
+		return lastElem.ToAddr, NodeExpectations{
+			LAddr:      containers.OptionalValue(lastElem.ToAddr),
+			Level:      containers.OptionalValue(lastElem.ToLevel),
+			Generation: containers.OptionalValue(lastElem.ToGeneration),
+			Owner: func(owner btrfsprim.ObjID, gen btrfsprim.Generation) error {
+				return firstElem.Tree.TreeCheckOwner(ctx, failOpen, owner, gen)
+			},
+			MinItem: containers.OptionalValue(lastElem.ToMinKey),
+			MaxItem: containers.OptionalValue(lastElem.ToMaxKey),
+		}, true
+	case PathItem:
+		return 0, NodeExpectations{}, false
+	default:
+		panic(fmt.Errorf("should not happen: unexpected PathElem type: %T", lastElem))
 	}
-
-	return path.Node(-1).ToNodeAddr, NodeExpectations{
-		LAddr:      containers.OptionalValue(path.Node(-1).ToNodeAddr),
-		Level:      containers.OptionalValue(path.Node(-1).ToNodeLevel),
-		Generation: containers.OptionalValue(path.Node(-1).ToNodeGeneration),
-		Owner:      checkOwner,
-		MinItem:    containers.OptionalValue(path.Node(-1).ToKey),
-		MaxItem:    containers.OptionalValue(path.Node(-1).ToMaxKey),
-	}, true
 }
 
 func (path Path) Parent() Path {
 	return path[:len(path)-1]
 }
-
-// Node is returns an element from the path; `path.Node(x)` is like
-// `&path[x]`, but negative values of x move down from the end of path
-// (similar to how lists work in many other languages, such as
-// Python).
-func (path Path) Node(x int) *PathElem {
-	if x < 0 {
-		x += len(path)
-	}
-	return &path[x]
-}
diff --git a/lib/btrfs/btrfstree/readnode.go b/lib/btrfs/btrfstree/readnode.go
deleted file mode 100644
index c9642da..0000000
--- a/lib/btrfs/btrfstree/readnode.go
+++ /dev/null
@@ -1,26 +0,0 @@
-// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package btrfstree
-
-import (
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
-	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
-)
-
-type NodeFile interface {
-	diskio.ReaderAt[btrfsvol.LogicalAddr]
-
-	// ParentTree, given a tree ID, returns that tree's parent
-	// tree, if it has one.
-	//
-	//  - non-zero, ?, true : the parent tree ID
-	//
-	//  - 0, 0, true : the tree does not have a parent
-	//
-	//  - ?, ?, false : the tree's parent information could not be
-	//    looked up
-	ParentTree(btrfsprim.ObjID) (btrfsprim.ObjID, btrfsprim.Generation, bool)
-}