Have a go at rearranging things in to a lib/btrfsprogs

3 files changed, 202 insertions, 0 deletions

diff --git a/lib/btrfsprogs/btrfsutil/open.go b/lib/btrfsprogs/btrfsutil/open.go
new file mode 100644
index 0000000..cc081a6
--- /dev/null
+++ b/lib/btrfsprogs/btrfsutil/open.go
@@ -0,0 +1,28 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrfsutil
+
+import (
+	"fmt"
+	"os"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+)
+
+func Open(flag int, filenames ...string) (*btrfs.FS, error) {
+	fs := new(btrfs.FS)
+	for _, filename := range filenames {
+		fh, err := os.OpenFile(filename, flag, 0)
+		if err != nil {
+			_ = fs.Close()
+			return nil, fmt.Errorf("file %q: %w", filename, err)
+		}
+		if err := fs.AddDevice(&btrfs.Device{File: fh}); err != nil {
+			_ = fs.Close()
+			return nil, fmt.Errorf("file %q: %w", filename, err)
+		}
+	}
+	return fs, nil
+}
diff --git a/lib/btrfsprogs/btrfsutil/scan.go b/lib/btrfsprogs/btrfsutil/scan.go
new file mode 100644
index 0000000..d83525c
--- /dev/null
+++ b/lib/btrfsprogs/btrfsutil/scan.go
@@ -0,0 +1,55 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrfsutil
+
+import (
+	"errors"
+	"fmt"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/util"
+)
+
+// ScanForNodes mimics btrfs-progs
+// cmds/rescue-chunk-recover.c:scan_one_device(), except rather than
+// doing something itself when it finds a node, it simply calls a
+// callback function.
+func ScanForNodes(dev *btrfs.Device, sb btrfs.Superblock, fn func(*util.Ref[btrfsvol.PhysicalAddr, btrfs.Node], error), prog func(btrfsvol.PhysicalAddr)) error {
+	devSize, err := dev.Size()
+	if err != nil {
+		return err
+	}
+
+	if sb.NodeSize < sb.SectorSize {
+		return fmt.Errorf("node_size(%v) < sector_size(%v)",
+			sb.NodeSize, sb.SectorSize)
+	}
+
+	for pos := btrfsvol.PhysicalAddr(0); pos+btrfsvol.PhysicalAddr(sb.NodeSize) < devSize; pos += btrfsvol.PhysicalAddr(sb.SectorSize) {
+		if util.InSlice(pos, btrfs.SuperblockAddrs) {
+			//fmt.Printf("sector@%v is a superblock\n", pos)
+			continue
+		}
+
+		if prog != nil {
+			prog(pos)
+		}
+
+		nodeRef, err := btrfs.ReadNode[btrfsvol.PhysicalAddr](dev, sb, pos, nil)
+		if err != nil && errors.Is(err, btrfs.ErrNotANode) {
+			continue
+		}
+		fn(nodeRef, err)
+
+		pos += btrfsvol.PhysicalAddr(sb.NodeSize) - btrfsvol.PhysicalAddr(sb.SectorSize)
+	}
+
+	if prog != nil {
+		prog(devSize)
+	}
+
+	return nil
+}
diff --git a/lib/btrfsprogs/btrfsutil/walk.go b/lib/btrfsprogs/btrfsutil/walk.go
new file mode 100644
index 0000000..0c54384
--- /dev/null
+++ b/lib/btrfsprogs/btrfsutil/walk.go
@@ -0,0 +1,119 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrfsutil
+
+import (
+	"fmt"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/util"
+)
+
+type WalkErr struct {
+	TreeName string
+	Path     btrfs.TreePath
+	Err      error
+}
+
+func (e WalkErr) Unwrap() error { return e.Err }
+
+func (e WalkErr) Error() string {
+	if len(e.Path) == 0 {
+		return fmt.Sprintf("%v: %v", e.TreeName, e.Err)
+	}
+	return fmt.Sprintf("%v: %v: %v", e.TreeName, e.Path, e.Err)
+}
+
+type WalkAllTreesHandler struct {
+	Err func(error)
+	// Callbacks for entire trees
+	PreTree  func(name string, id btrfs.ObjID)
+	PostTree func(name string, id btrfs.ObjID)
+	// Callbacks for nodes or smaller
+	UnsafeNodes bool
+	btrfs.TreeWalkHandler
+}
+
+// WalkAllTrees walks all trees in a *btrfs.FS.  Rather than returning
+// an error, it calls errCb each time an error is encountered.  The
+// error will always be of type WalkErr.
+func WalkAllTrees(fs *btrfs.FS, cbs WalkAllTreesHandler) {
+	var treeName string
+	handleErr := func(path btrfs.TreePath, err error) {
+		cbs.Err(WalkErr{
+			TreeName: treeName,
+			Path:     path,
+			Err:      err,
+		})
+	}
+
+	trees := []struct {
+		Name string
+		ID   btrfs.ObjID
+	}{
+		{
+			Name: "root tree",
+			ID:   btrfs.ROOT_TREE_OBJECTID,
+		},
+		{
+			Name: "chunk tree",
+			ID:   btrfs.CHUNK_TREE_OBJECTID,
+		},
+		{
+			Name: "log tree",
+			ID:   btrfs.TREE_LOG_OBJECTID,
+		},
+		{
+			Name: "block group tree",
+			ID:   btrfs.BLOCK_GROUP_TREE_OBJECTID,
+		},
+	}
+	origItem := cbs.Item
+	cbs.Item = func(path btrfs.TreePath, item btrfs.Item) error {
+		if item.Head.Key.ItemType == btrfsitem.ROOT_ITEM_KEY {
+			trees = append(trees, struct {
+				Name string
+				ID   btrfs.ObjID
+			}{
+				Name: fmt.Sprintf("tree %v (via %v %v)",
+					item.Head.Key.ObjectID.Format(0), treeName, path),
+				ID: item.Head.Key.ObjectID,
+			})
+		}
+		if origItem != nil {
+			return origItem(path, item)
+		}
+		return nil
+	}
+
+	if !cbs.UnsafeNodes {
+		origNode := cbs.Node
+		cbs.Node = func(path btrfs.TreePath, node *util.Ref[btrfsvol.LogicalAddr, btrfs.Node], err error) error {
+			if err != nil {
+				handleErr(path, err)
+			}
+			if node != nil && origNode != nil {
+				return origNode(path, node, nil)
+			}
+			return nil
+		}
+	}
+
+	for i := 0; i < len(trees); i++ {
+		tree := trees[i]
+		treeName = tree.Name
+		if cbs.PreTree != nil {
+			cbs.PreTree(treeName, tree.ID)
+		}
+		if err := fs.TreeWalk(tree.ID, cbs.TreeWalkHandler); err != nil {
+			handleErr(nil, err)
+		}
+		if cbs.PostTree != nil {
+			cbs.PostTree(treeName, tree.ID)
+		}
+	}
+}