pointers! caching! callbacks!

5 files changed, 87 insertions, 32 deletions

diff --git a/cmd/btrfs-dbg/main.go b/cmd/btrfs-dbg/main.go
index d50b44c..462b333 100644
--- a/cmd/btrfs-dbg/main.go
+++ b/cmd/btrfs-dbg/main.go
@@ -8,6 +8,7 @@ import (
 
 	"lukeshu.com/btrfs-tools/pkg/btrfs"
 	"lukeshu.com/btrfs-tools/pkg/btrfsmisc"
+	"lukeshu.com/btrfs-tools/pkg/util"
 )
 
 func main() {
@@ -60,7 +61,16 @@ func Main(imgfilename string) (err error) {
 	}
 	spew.Dump(syschunks)
 
-	if err := btrfsmisc.ScanForNodes(fs.Devices[0], superblocks[0].Data); err != nil {
+	if err := btrfsmisc.ScanForNodes(fs.Devices[0], superblocks[0].Data, func(nodeRef *util.Ref[btrfs.PhysicalAddr, btrfs.Node], err error) {
+		if err != nil {
+			fmt.Println(err)
+		} else {
+			fmt.Printf("node@%d: physical_addr=0x%0X logical_addr=0x%0X generation=%d owner=%v level=%d\n",
+				nodeRef.Addr,
+				nodeRef.Addr, nodeRef.Data.Head.Addr,
+				nodeRef.Data.Head.Generation, nodeRef.Data.Head.Owner, nodeRef.Data.Head.Level)
+		}
+	}); err != nil {
 		return err
 	}
 
diff --git a/pkg/btrfs/io1_device.go b/pkg/btrfs/io1_device.go
index 55b7525..3b84fcc 100644
--- a/pkg/btrfs/io1_device.go
+++ b/pkg/btrfs/io1_device.go
@@ -10,6 +10,9 @@ import (
 
 type Device struct {
 	*os.File
+
+	cacheSuperblocks []*util.Ref[PhysicalAddr, Superblock]
+	cacheSuperblock  *util.Ref[PhysicalAddr, Superblock]
 }
 
 func (dev Device) Size() (PhysicalAddr, error) {
@@ -30,7 +33,10 @@ func (dev *Device) ReadAt(dat []byte, paddr PhysicalAddr) (int, error) {
 	return dev.File.ReadAt(dat, int64(paddr))
 }
 
-func (dev *Device) Superblocks() ([]util.Ref[PhysicalAddr, Superblock], error) {
+func (dev *Device) Superblocks() ([]*util.Ref[PhysicalAddr, Superblock], error) {
+	if dev.cacheSuperblocks != nil {
+		return dev.cacheSuperblocks, nil
+	}
 	superblockSize := PhysicalAddr(binstruct.StaticSize(Superblock{}))
 
 	sz, err := dev.Size()
@@ -38,10 +44,10 @@ func (dev *Device) Superblocks() ([]util.Ref[PhysicalAddr, Superblock], error) {
 		return nil, err
 	}
 
-	var ret []util.Ref[PhysicalAddr, Superblock]
+	var ret []*util.Ref[PhysicalAddr, Superblock]
 	for i, addr := range SuperblockAddrs {
 		if addr+superblockSize <= sz {
-			superblock := util.Ref[PhysicalAddr, Superblock]{
+			superblock := &util.Ref[PhysicalAddr, Superblock]{
 				File: dev,
 				Addr: addr,
 			}
@@ -54,23 +60,30 @@ func (dev *Device) Superblocks() ([]util.Ref[PhysicalAddr, Superblock], error) {
 	if len(ret) == 0 {
 		return nil, fmt.Errorf("no superblocks")
 	}
+	dev.cacheSuperblocks = ret
 	return ret, nil
 }
 
-func (dev *Device) Superblock() (ret util.Ref[PhysicalAddr, Superblock], err error) {
+func (dev *Device) Superblock() (*util.Ref[PhysicalAddr, Superblock], error) {
+	if dev.cacheSuperblock != nil {
+		return dev.cacheSuperblock, nil
+	}
 	sbs, err := dev.Superblocks()
 	if err != nil {
-		return ret, err
+		return nil, err
 	}
+
 	for i, sb := range sbs {
 		if err := sb.Data.ValidateChecksum(); err != nil {
-			return ret, fmt.Errorf("superblock %d: %w", i, err)
+			return nil, fmt.Errorf("superblock %d: %w", i, err)
 		}
 		if i > 0 {
 			if !sb.Data.Equal(sbs[0].Data) {
-				return ret, fmt.Errorf("superblock %d and superblock %d disagree", 0, i)
+				return nil, fmt.Errorf("superblock %d and superblock %d disagree", 0, i)
 			}
 		}
 	}
+
+	dev.cacheSuperblock = sbs[0]
 	return sbs[0], nil
 }
diff --git a/pkg/btrfs/io2_fs.go b/pkg/btrfs/io2_fs.go
index ae0680a..ff5415a 100644
--- a/pkg/btrfs/io2_fs.go
+++ b/pkg/btrfs/io2_fs.go
@@ -16,6 +16,9 @@ type FS struct {
 	initErr  error
 	uuid2dev map[UUID]*Device
 	chunks   []SysChunk
+
+	cacheSuperblocks []*util.Ref[PhysicalAddr, Superblock]
+	cacheSuperblock  *util.Ref[PhysicalAddr, Superblock]
 }
 
 func (fs *FS) Name() string {
@@ -38,8 +41,11 @@ func (fs *FS) Size() (LogicalAddr, error) {
 	return ret, nil
 }
 
-func (fs *FS) Superblocks() ([]util.Ref[PhysicalAddr, Superblock], error) {
-	var ret []util.Ref[PhysicalAddr, Superblock]
+func (fs *FS) Superblocks() ([]*util.Ref[PhysicalAddr, Superblock], error) {
+	if fs.cacheSuperblocks != nil {
+		return fs.cacheSuperblocks, nil
+	}
+	var ret []*util.Ref[PhysicalAddr, Superblock]
 	for _, dev := range fs.Devices {
 		sbs, err := dev.Superblocks()
 		if err != nil {
@@ -47,13 +53,17 @@ func (fs *FS) Superblocks() ([]util.Ref[PhysicalAddr, Superblock], error) {
 		}
 		ret = append(ret, sbs...)
 	}
+	fs.cacheSuperblocks = ret
 	return ret, nil
 }
 
-func (fs *FS) Superblock() (ret util.Ref[PhysicalAddr, Superblock], err error) {
+func (fs *FS) Superblock() (*util.Ref[PhysicalAddr, Superblock], error) {
+	if fs.cacheSuperblock != nil {
+		return fs.cacheSuperblock, nil
+	}
 	sbs, err := fs.Superblocks()
 	if err != nil {
-		return ret, err
+		return nil, err
 	}
 
 	fname := ""
@@ -67,19 +77,20 @@ func (fs *FS) Superblock() (ret util.Ref[PhysicalAddr, Superblock], err error) {
 		}
 
 		if err := sb.Data.ValidateChecksum(); err != nil {
-			return ret, fmt.Errorf("file %q superblock %d: %w", sb.File.Name(), sbi, err)
+			return nil, fmt.Errorf("file %q superblock %d: %w", sb.File.Name(), sbi, err)
 		}
 		if i > 0 {
 			// This is probably wrong, but lots of my
 			// multi-device code is probably wrong.
 			if !sb.Data.Equal(sbs[0].Data) {
-				return ret, fmt.Errorf("file %q superblock %d and file %q superblock %d disagree",
+				return nil, fmt.Errorf("file %q superblock %d and file %q superblock %d disagree",
 					sbs[0].File.Name(), 0,
 					sb.File.Name(), sbi)
 			}
 		}
 	}
 
+	fs.cacheSuperblock = sbs[0]
 	return sbs[0], nil
 }
 
diff --git a/pkg/btrfs/types_btree.go b/pkg/btrfs/types_btree.go
index 8db00e2..5adc7ec 100644
--- a/pkg/btrfs/types_btree.go
+++ b/pkg/btrfs/types_btree.go
@@ -249,49 +249,47 @@ func (node *Node) LeafFreeSpace() uint32 {
 	return freeSpace
 }
 
-func (fs *FS) ReadNode(addr LogicalAddr) (util.Ref[LogicalAddr, Node], error) {
-	var ret util.Ref[LogicalAddr, Node]
-
+func (fs *FS) ReadNode(addr LogicalAddr) (*util.Ref[LogicalAddr, Node], error) {
 	sb, err := fs.Superblock()
 	if err != nil {
-		return ret, fmt.Errorf("btrfs.FS.ReadNode: %w", err)
+		return nil, fmt.Errorf("btrfs.FS.ReadNode: %w", err)
 	}
 
 	// read
 
 	nodeBuf := make([]byte, sb.Data.NodeSize)
 	if _, err := fs.ReadAt(nodeBuf, addr); err != nil {
-		return ret, err
+		return nil, err
 	}
 
 	var node Node
 	node.Size = sb.Data.NodeSize
 
 	if _, err := node.UnmarshalBinary(nodeBuf); err != nil {
-		return ret, fmt.Errorf("btrfs.FS.ReadNode: node@%d: %w", addr, err)
+		return nil, fmt.Errorf("btrfs.FS.ReadNode: node@%d: %w", addr, err)
 	}
 
 	// sanity checking
 
 	if node.Head.MetadataUUID != sb.Data.EffectiveMetadataUUID() {
-		return ret, fmt.Errorf("btrfs.FS.ReadNode: node@%d: does not look like a node", addr)
+		return nil, fmt.Errorf("btrfs.FS.ReadNode: node@%d: does not look like a node", addr)
 	}
 
 	if node.Head.Addr != addr {
-		return ret, fmt.Errorf("btrfs.FS.ReadNode: node@%d: read from laddr=%d but claims to be at laddr=%d",
+		return nil, fmt.Errorf("btrfs.FS.ReadNode: node@%d: read from laddr=%d but claims to be at laddr=%d",
 			addr, addr, node.Head.Addr)
 	}
 
 	stored := node.Head.Checksum
 	calced := CRC32c(nodeBuf[binstruct.StaticSize(CSum{}):])
 	if calced != stored {
-		return ret, fmt.Errorf("btrfs.FS.ReadNode: node@%d: checksum mismatch: stored=%s calculated=%s",
+		return nil, fmt.Errorf("btrfs.FS.ReadNode: node@%d: checksum mismatch: stored=%s calculated=%s",
 			addr, stored, calced)
 	}
 
 	// return
 
-	return util.Ref[LogicalAddr, Node]{
+	return &util.Ref[LogicalAddr, Node]{
 		File: fs,
 		Addr: addr,
 		Data: node,
diff --git a/pkg/btrfsmisc/fsck.go b/pkg/btrfsmisc/fsck.go
index 77305f8..a9e7d1d 100644
--- a/pkg/btrfsmisc/fsck.go
+++ b/pkg/btrfsmisc/fsck.go
@@ -11,7 +11,7 @@ import (
 // ScanForNodes mimics btrfs-progs
 // cmds/rescue-chunk-recover.c:scan_one_device(), except it doesn't do
 // anything but log when it finds a node.
-func ScanForNodes(dev *btrfs.Device, sb btrfs.Superblock) error {
+func ScanForNodes(dev *btrfs.Device, sb btrfs.Superblock, fn func(*util.Ref[btrfs.PhysicalAddr, btrfs.Node], error)) error {
 	devSize, err := dev.Size()
 	if err != nil {
 		return err
@@ -25,27 +25,50 @@ func ScanForNodes(dev *btrfs.Device, sb btrfs.Superblock) error {
 	nodeBuf := make([]byte, sb.NodeSize)
 	for pos := btrfs.PhysicalAddr(0); pos+btrfs.PhysicalAddr(sb.SectorSize) < devSize; pos += btrfs.PhysicalAddr(sb.SectorSize) {
 		if util.InSlice(pos, btrfs.SuperblockAddrs) {
-			fmt.Printf("sector@%d is a superblock\n", pos)
+			//fmt.Printf("sector@%d is a superblock\n", pos)
 			continue
 		}
+
+		// read
+
 		if _, err := dev.ReadAt(nodeBuf, pos); err != nil {
-			return fmt.Errorf("sector@%d: %w", pos, err)
+			fn(nil, fmt.Errorf("sector@%d: %w", pos, err))
+			continue
 		}
+
+		// does it look like a node?
+
 		var nodeHeader btrfs.NodeHeader
 		if _, err := binstruct.Unmarshal(nodeBuf, &nodeHeader); err != nil {
-			return fmt.Errorf("sector@%d: %w", pos, err)
+			fn(nil, fmt.Errorf("sector@%d: %w", pos, err))
 		}
 		if nodeHeader.MetadataUUID != sb.EffectiveMetadataUUID() {
 			//fmt.Printf("sector@%d does not look like a node\n", pos)
 			continue
 		}
-		if nodeHeader.Checksum != btrfs.CRC32c(nodeBuf[0x20:]) {
-			fmt.Printf("sector@%d looks like a node but is corrupt (checksum doesn't match)\n", pos)
+
+		// ok, it looks like a node; go ahead and read it as a node
+
+		nodeRef := &util.Ref[btrfs.PhysicalAddr, btrfs.Node]{
+			File: dev,
+			Addr: pos,
+			Data: btrfs.Node{
+				Size: sb.NodeSize,
+			},
+		}
+		if _, err := nodeRef.Data.UnmarshalBinary(nodeBuf); err != nil {
+			fn(nil, fmt.Errorf("sector@%d: %w", pos, err))
+			continue
+		}
+
+		// finally, process the node
+
+		if nodeRef.Data.Head.Checksum != btrfs.CRC32c(nodeBuf[0x20:]) {
+			fn(nodeRef, fmt.Errorf("sector@%d looks like a node but is corrupt (checksum doesn't match)", pos))
 			continue
 		}
 
-		fmt.Printf("node@%d: physical_addr=0x%0X logical_addr=0x%0X generation=%d owner=%v level=%d\n",
-			pos, pos, nodeHeader.Addr, nodeHeader.Generation, nodeHeader.Owner, nodeHeader.Level)
+		fn(nodeRef, nil)
 
 		pos += btrfs.PhysicalAddr(sb.NodeSize) - btrfs.PhysicalAddr(sb.SectorSize)
 	}