scan for nodes

7 files changed, 300 insertions, 18 deletions

diff --git a/pkg/btrfs/bitfields.go b/pkg/btrfs/bitfields.go
new file mode 100644
index 0000000..ead4b0f
--- /dev/null
+++ b/pkg/btrfs/bitfields.go
@@ -0,0 +1,68 @@
+package btrfs
+
+import (
+	"fmt"
+	"strings"
+)
+
+func bitfieldString[T ~uint8 | ~uint16 | ~uint32 | ~uint64](bitfield T, bitnames []string) string {
+	if bitfield == 0 {
+		return "0"
+	}
+	var out strings.Builder
+	fmt.Fprintf(&out, "(0x%0X)", uint64(bitfield))
+	rest := bitfield
+	sep := ' '
+	for i := 0; rest != 0; i++ {
+		if rest&(1<<i) != 0 {
+			out.WriteRune(sep)
+			if i < len(bitnames) {
+				out.WriteString(bitnames[i])
+			} else {
+				fmt.Fprintf(&out, "(1<<%d)", i)
+			}
+			sep = '|'
+		}
+		rest &^= 1 << i
+	}
+	return out.String()
+}
+
+type IncompatFlags uint64
+
+const (
+	FeatureIncompatMixedBackref = IncompatFlags(1 << iota)
+	FeatureIncompatDefaultSubvol
+	FeatureIncompatMixedGroups
+	FeatureIncompatCompressLZO
+	FeatureIncompatCompressZSTD
+	FeatureIncompatBigMetadata // buggy
+	FeatureIncompatExtendedIRef
+	FeatureIncompatRAID56
+	FeatureIncompatSkinnyMetadata
+	FeatureIncompatNoHoles
+	FeatureIncompatMetadataUUID
+	FeatureIncompatRAID1C34
+	FeatureIncompatZoned
+	FeatureIncompatExtentTreeV2
+)
+
+var incompatFlagNames = []string{
+	"FeatureIncompatMixedBackref",
+	"FeatureIncompatDefaultSubvol",
+	"FeatureIncompatMixedGroups",
+	"FeatureIncompatCompressLZO",
+	"FeatureIncompatCompressZSTD",
+	"FeatureIncompatBigMetadata ",
+	"FeatureIncompatExtendedIRef",
+	"FeatureIncompatRAID56",
+	"FeatureIncompatSkinnyMetadata",
+	"FeatureIncompatNoHoles",
+	"FeatureIncompatMetadataUUID",
+	"FeatureIncompatRAID1C34",
+	"FeatureIncompatZoned",
+	"FeatureIncompatExtentTreeV2",
+}
+
+func (f IncompatFlags) Has(req IncompatFlags) bool { return f&req == req }
+func (f IncompatFlags) String() string             { return bitfieldString(f, incompatFlagNames) }
diff --git a/pkg/btrfs/crc32c.go b/pkg/btrfs/crc32c.go
index 4ea4169..5e050cc 100644
--- a/pkg/btrfs/crc32c.go
+++ b/pkg/btrfs/crc32c.go
@@ -1,12 +1,17 @@
 package btrfs
 
 import (
+	"bytes"
 	"encoding/binary"
 	"hash/crc32"
 )
 
 type CSum [0x20]byte
 
+func (a CSum) Equal(b CSum) bool {
+	return bytes.Equal(a[:], b[:])
+}
+
 func CRC32c(data []byte) CSum {
 	crc := crc32.Update(0, crc32.MakeTable(crc32.Castagnoli), data)
 
diff --git a/pkg/btrfs/image.go b/pkg/btrfs/image.go
index 6037c4d..9faded5 100644
--- a/pkg/btrfs/image.go
+++ b/pkg/btrfs/image.go
@@ -37,15 +37,15 @@ func (r *Ref[T]) Read() error {
 	return binstruct.Unmarshal(buf, &r.Data)
 }
 
+var superblockAddrs = []int64{
+	0x00_0001_0000, // 64KiB
+	0x00_0400_0000, // 64MiB
+	0x40_0000_0000, // 256GiB
+}
+
 func (img *Img) Superblocks() ([]Ref[Superblock], error) {
 	const superblockSize = 0x1000
 
-	var superblockAddrs = []int64{
-		0x00_0001_0000, // 64KiB
-		0x00_0400_0000, // 64MiB
-		0x40_0000_0000, // 256GiB
-	}
-
 	sz, err := img.Size()
 	if err != nil {
 		return nil, err
@@ -69,3 +69,48 @@ func (img *Img) Superblocks() ([]Ref[Superblock], error) {
 	}
 	return ret, nil
 }
+
+// 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 (img *Img) ScanForNodes(sb Superblock) error {
+	devSize, err := img.Size()
+	if err != nil {
+		return err
+	}
+
+	if sb.NodeSize < sb.SectorSize {
+		return fmt.Errorf("node_size(%d) < sector_size(%d)",
+			sb.NodeSize, sb.SectorSize)
+	}
+
+	nodeBuf := make([]byte, sb.NodeSize)
+	for pos := int64(0); pos < devSize; pos += int64(sb.SectorSize) {
+		if inSlice(pos, superblockAddrs) {
+			fmt.Printf("sector@%d is a superblock\n", pos)
+			continue
+		}
+		if _, err := img.ReadAt(nodeBuf, pos); err != nil {
+			return fmt.Errorf("sector@%d: %w", pos, err)
+		}
+		var nodeHeader NodeHeader
+		if err := binstruct.Unmarshal(nodeBuf, &nodeHeader); err != nil {
+			return fmt.Errorf("sector@%d: %w", pos, err)
+		}
+		if !nodeHeader.MetadataUUID.Equal(sb.EffectiveMetadataUUID()) {
+			//fmt.Printf("sector@%d does not look like a node\n", pos)
+			continue
+		}
+		if !nodeHeader.Checksum.Equal(CRC32c(nodeBuf[0x20:])) {
+			fmt.Printf("sector@%d looks like a node but is corrupt (checksum doesn't match)\n", pos)
+			continue
+		}
+
+		fmt.Printf("node@%d: physical_addr=0x%0X logical_addr=0x%0X generation=%d owner_tree=%v level=%d\n",
+			pos, pos, nodeHeader.Addr, nodeHeader.Generation, nodeHeader.OwnerTree, nodeHeader.Level)
+
+		pos += int64(sb.NodeSize) - int64(sb.SectorSize)
+	}
+
+	return nil
+}
diff --git a/pkg/btrfs/objid.go b/pkg/btrfs/objid.go
new file mode 100644
index 0000000..c0c73c5
--- /dev/null
+++ b/pkg/btrfs/objid.go
@@ -0,0 +1,95 @@
+package btrfs
+
+import (
+	"fmt"
+)
+
+type ObjID uint64
+
+const maxUint64pp = 0x1_0000_0000
+
+const (
+	// The IDs of the various trees
+	BTRFS_ROOT_TREE_OBJECTID        = ObjID(1) // holds pointers to all of the tree roots
+	BTRFS_EXTENT_TREE_OBJECTID      = ObjID(2) // stores information about which extents are in use, and reference counts
+	BTRFS_CHUNK_TREE_OBJECTID       = ObjID(3) // chunk tree stores translations from logical -> physical block numbering
+	BTRFS_DEV_TREE_OBJECTID         = ObjID(4) // stores info about which areas of a given device are in use; one per device
+	BTRFS_FS_TREE_OBJECTID          = ObjID(5) // one per subvolume, storing files and directories
+	BTRFS_ROOT_TREE_DIR_OBJECTID    = ObjID(6) // directory objectid inside the root tree
+	BTRFS_CSUM_TREE_OBJECTID        = ObjID(7) // holds checksums of all the data extents
+	BTRFS_QUOTA_TREE_OBJECTID       = ObjID(8)
+	BTRFS_UUID_TREE_OBJECTID        = ObjID(9)  // for storing items that use the BTRFS_UUID_KEY*
+	BTRFS_FREE_SPACE_TREE_OBJECTID  = ObjID(10) // tracks free space in block groups.
+	BTRFS_BLOCK_GROUP_TREE_OBJECTID = ObjID(11) // hold the block group items.
+
+	// Objects in the DEV_TREE
+	BTRFS_DEV_STATS_OBJECTID = ObjID(0) // device stats in the device tree
+
+	// ???
+	BTRFS_BALANCE_OBJECTID         = ObjID(maxUint64pp - 4) // for storing balance parameters in the root tree
+	BTRFS_ORPHAN_OBJECTID          = ObjID(maxUint64pp - 5) // orphan objectid for tracking unlinked/truncated files
+	BTRFS_TREE_LOG_OBJECTID        = ObjID(maxUint64pp - 6) // does write ahead logging to speed up fsyncs
+	BTRFS_TREE_LOG_FIXUP_OBJECTID  = ObjID(maxUint64pp - 7)
+	BTRFS_TREE_RELOC_OBJECTID      = ObjID(maxUint64pp - 8) // space balancing
+	BTRFS_DATA_RELOC_TREE_OBJECTID = ObjID(maxUint64pp - 9)
+	BTRFS_EXTENT_CSUM_OBJECTID     = ObjID(maxUint64pp - 10) // extent checksums all have this objectid
+	BTRFS_FREE_SPACE_OBJECTID      = ObjID(maxUint64pp - 11) // For storing free space cache
+	BTRFS_FREE_INO_OBJECTID        = ObjID(maxUint64pp - 12) // stores the inode number for the free-ino cache
+
+	BTRFS_MULTIPLE_OBJECTIDS = ObjID(maxUint64pp - 255) // dummy objectid represents multiple objectids
+
+	// All files have objectids in this range.
+	BTRFS_FIRST_FREE_OBJECTID       = ObjID(256)
+	BTRFS_LAST_FREE_OBJECTID        = ObjID(maxUint64pp - 256)
+	BTRFS_FIRST_CHUNK_TREE_OBJECTID = ObjID(256)
+
+	// Objects in the CHUNK_TREE
+	BTRFS_DEV_ITEMS_OBJECTID = ObjID(1)
+
+	// ???
+	BTRFS_EMPTY_SUBVOL_DIR_OBJECTID = ObjID(2)
+)
+
+func (id ObjID) String() string {
+	if id > BTRFS_LAST_FREE_OBJECTID {
+		names := map[ObjID]string{
+			BTRFS_BALANCE_OBJECTID:         "BTRFS_BALANCE_OBJECTID",
+			BTRFS_ORPHAN_OBJECTID:          "BTRFS_ORPHAN_OBJECTID",
+			BTRFS_TREE_LOG_OBJECTID:        "BTRFS_TREE_LOG_OBJECTID",
+			BTRFS_TREE_LOG_FIXUP_OBJECTID:  "BTRFS_TREE_LOG_FIXUP_OBJECTID",
+			BTRFS_TREE_RELOC_OBJECTID:      "BTRFS_TREE_RELOC_OBJECTID",
+			BTRFS_DATA_RELOC_TREE_OBJECTID: "BTRFS_DATA_RELOC_TREE_OBJECTID",
+			BTRFS_EXTENT_CSUM_OBJECTID:     "BTRFS_EXTENT_CSUM_OBJECTID",
+			BTRFS_FREE_SPACE_OBJECTID:      "BTRFS_FREE_SPACE_OBJECTID",
+			BTRFS_FREE_INO_OBJECTID:        "BTRFS_FREE_INO_OBJECTID",
+			BTRFS_MULTIPLE_OBJECTIDS:       "BTRFS_MULTIPLE_OBJECTIDS",
+		}
+		if name, ok := names[id]; ok {
+			return name
+		}
+		return fmt.Sprintf("%d", int64(id))
+	}
+	return fmt.Sprintf("%d", id)
+}
+
+type TreeObjID ObjID
+
+func (id TreeObjID) String() string {
+	names := map[ObjID]string{
+		BTRFS_ROOT_TREE_OBJECTID:        "BTRFS_ROOT_TREE_OBJECTID",
+		BTRFS_EXTENT_TREE_OBJECTID:      "BTRFS_EXTENT_TREE_OBJECTID",
+		BTRFS_CHUNK_TREE_OBJECTID:       "BTRFS_CHUNK_TREE_OBJECTID",
+		BTRFS_DEV_TREE_OBJECTID:         "BTRFS_DEV_TREE_OBJECTID",
+		BTRFS_FS_TREE_OBJECTID:          "BTRFS_FS_TREE_OBJECTID",
+		BTRFS_ROOT_TREE_DIR_OBJECTID:    "BTRFS_ROOT_TREE_DIR_OBJECTID",
+		BTRFS_CSUM_TREE_OBJECTID:        "BTRFS_CSUM_TREE_OBJECTID",
+		BTRFS_QUOTA_TREE_OBJECTID:       "BTRFS_QUOTA_TREE_OBJECTID",
+		BTRFS_UUID_TREE_OBJECTID:        "BTRFS_UUID_TREE_OBJECTID",
+		BTRFS_FREE_SPACE_TREE_OBJECTID:  "BTRFS_FREE_SPACE_TREE_OBJECTID",
+		BTRFS_BLOCK_GROUP_TREE_OBJECTID: "BTRFS_BLOCK_GROUP_TREE_OBJECTID",
+	}
+	if name, ok := names[ObjID(id)]; ok {
+		return name
+	}
+	return ObjID(id).String()
+}
diff --git a/pkg/btrfs/types.go b/pkg/btrfs/structs.go
index a91708e..92db216 100644
--- a/pkg/btrfs/types.go
+++ b/pkg/btrfs/structs.go
@@ -9,7 +9,6 @@ import (
 type (
 	PhysicalAddr int64
 	LogicalAddr  int64
-	ObjID        int64
 )
 
 type Key struct {
@@ -44,7 +43,7 @@ type Superblock struct {
 	LogRootTransID  uint64 `bin:"off=68, siz=8"` // log_root_transid
 	TotalBytes      uint64 `bin:"off=70, siz=8"` // total_bytes
 	BytesUsed       uint64 `bin:"off=78, siz=8"` // bytes_used
-	RootDirObjectID uint64 `bin:"off=80, siz=8"` // root_dir_objectid (usually 6)
+	RootDirObjectID ObjID  `bin:"off=80, siz=8"` // root_dir_objectid (usually 6)
 	NumDevices      uint64 `bin:"off=88, siz=8"` // num_devices
 
 	SectorSize        uint32 `bin:"off=90, siz=4"` // sectorsize
@@ -53,11 +52,11 @@ type Superblock struct {
 	StripeSize        uint32 `bin:"off=9c, siz=4"` // stripesize
 	SysChunkArraySize uint32 `bin:"off=a0, siz=4"` // sys_chunk_array_size
 
-	ChunkRootGeneration uint64 `bin:"off=a4, siz=8"` // chunk_root_generation
-	CompatFlags         uint64 `bin:"off=ac, siz=8"` // compat_flags
-	CompatROFlags       uint64 `bin:"off=b4, siz=8"` // compat_ro_flags - only implementations that support the flags can write to the filesystem
-	IncompatFlags       uint64 `bin:"off=bc, siz=8"` // incompat_flags - only implementations that support the flags can use the filesystem
-	ChecksumType        uint16 `bin:"off=c4, siz=2"` // csum_type - Btrfs currently uses the CRC32c little-endian hash function with seed -1.
+	ChunkRootGeneration uint64        `bin:"off=a4, siz=8"` // chunk_root_generation
+	CompatFlags         uint64        `bin:"off=ac, siz=8"` // compat_flags
+	CompatROFlags       uint64        `bin:"off=b4, siz=8"` // compat_ro_flags - only implementations that support the flags can write to the filesystem
+	IncompatFlags       IncompatFlags `bin:"off=bc, siz=8"` // incompat_flags - only implementations that support the flags can use the filesystem
+	ChecksumType        uint16        `bin:"off=c4, siz=2"` // csum_type - Btrfs currently uses the CRC32c little-endian hash function with seed -1.
 
 	RootLevel  uint8 `bin:"off=c6, siz=1"` // root_level
 	ChunkLevel uint8 `bin:"off=c7, siz=1"` // chunk_root_level
@@ -68,13 +67,24 @@ type Superblock struct {
 	CacheGeneration    uint64      `bin:"off=22b, siz=8"`   // cache_generation
 	UUIDTreeGeneration uint64      `bin:"off=233, siz=8"`   // uuid_tree_generation
 
-	Reserved [0xf0]byte `bin:"off=23b, siz=f0"` // reserved /* future expansion */
+	// FeatureIncompatMetadataUUID
+	MetadataUUID UUID `bin:"off=23b, siz=10"`
+
+	// FeatureIncompatExtentTreeV2
+	NumGlobalRoots uint64 `bin:"off=24b, siz=8"`
+
+	// FeatureIncompatExtentTreeV2
+	BlockGroupRoot           uint64 `bin:"off=253, siz=8"`
+	BlockGroupRootGeneration uint64 `bin:"off=25b, siz=8"`
+	BlockGroupRootLevel      uint8  `bin:"off=263, siz=1"`
+
+	Reserved [199]byte `bin:"off=264, siz=c7"` // future expansion
 
 	SysChunkArray  [0x800]byte `bin:"off=32b, siz=800"` // sys_chunk_array:(n bytes valid) Contains (KEY . CHUNK_ITEM) pairs for all SYSTEM chunks. This is needed to bootstrap the mapping from logical addresses to physical.
 	TODOSuperRoots [0x2a0]byte `bin:"off=b2b, siz=2a0"` // Contain super_roots (4 btrfs_root_backup)
 
-	Unused [0x235]byte `bin:"off=dcb, siz=235"` // current unused
-
+	// Padded to 4096 bytes
+	Padding       [565]byte `bin:"off=dcb, siz=235"`
 	binstruct.End `bin:"off=1000"`
 }
 
@@ -86,6 +96,13 @@ func (sb Superblock) CalculateChecksum() (CSum, error) {
 	return CRC32c(data[0x20:]), nil
 }
 
+func (sb Superblock) EffectiveMetadataUUID() UUID {
+	if !sb.IncompatFlags.Has(FeatureIncompatMetadataUUID) {
+		return sb.FSUUID
+	}
+	return sb.MetadataUUID
+}
+
 type SysChunk struct {
 	Key           `bin:"off=0, siz=11"`
 	ChunkItem     `bin:"off=11, siz=30"`
@@ -116,8 +133,45 @@ func (sb Superblock) ParseSysChunkArray() ([]SysChunk, error) {
 	return ret, nil
 }
 
+type NodeHeader struct {
+	Checksum      CSum        `bin:"off=0,  siz=20"` // Checksum of everything after this field (from 20 to the end of the node)
+	MetadataUUID  UUID        `bin:"off=20, siz=10"` // FS UUID
+	Addr          LogicalAddr `bin:"off=30, siz=8"`  // Logical address of this node
+	Flags         uint64      `bin:"off=38, siz=8"`  // Flags
+	ChunkTreeUUID UUID        `bin:"off=40, siz=10"` // Chunk tree UUID
+	Generation    uint64      `bin:"off=50, siz=8"`  // Generation
+	OwnerTree     TreeObjID   `bin:"off=58, siz=8"`  // The ID of the tree that contains this node
+	NumItems      uint32      `bin:"off=60, siz=4"`  // Number of items
+	Level         uint8       `bin:"off=64, siz=1"`  // Level (0 for leaf nodes)
+	binstruct.End `bin:"off=65"`
+}
+
+type InternalNode struct {
+	NodeHeader
+	Body []KeyPointer
+}
+
+type KeyPointer struct {
+	Key           Key    `bin:"off=0, siz=11"`
+	BlockNumber   uint64 `bin:"off=11, siz=8"`
+	Generation    uint64 `bin:"off=19, siz=8"`
+	binstruct.End `bin:"off=21"`
+}
+
+type LeafNode struct {
+	NodeHeader
+	Body []Item
+}
+
+type Item struct {
+	Key           Key    `bin:"off=0, siz=11"`
+	DataOffset    uint32 `bin:"off=11, siz=4"` // relative to the end of the header (0x65)
+	DataSize      uint32 `bin:"off=15, siz=4"`
+	binstruct.End `bin:"off=19"`
+}
+
 type DevItem struct {
-	DeviceID uint64 `bin:"off=0,    siz=8"` // device id
+	DeviceID ObjID `bin:"off=0,    siz=8"` // device ID
 
 	NumBytes     uint64 `bin:"off=8,    siz=8"` // number of bytes
 	NumBytesUsed uint64 `bin:"off=10,   siz=8"` // number of bytes used
@@ -156,7 +210,7 @@ type ChunkItem struct {
 
 type ChunkItemStripe struct {
 	// Stripes follow (for each number of stripes):
-	DeviceID      ObjID  `bin:"off=0,  siz=8"`  // device id
+	DeviceID      ObjID  `bin:"off=0,  siz=8"`  // device ID
 	Offset        uint64 `bin:"off=8,  siz=8"`  // offset
 	DeviceUUID    UUID   `bin:"off=10, siz=10"` // device UUID
 	binstruct.End `bin:"off=20"`
diff --git a/pkg/btrfs/util.go b/pkg/btrfs/util.go
new file mode 100644
index 0000000..04462f8
--- /dev/null
+++ b/pkg/btrfs/util.go
@@ -0,0 +1,10 @@
+package btrfs
+
+func inSlice[T comparable](needle T, haystack []T) bool {
+	for _, straw := range haystack {
+		if needle == straw {
+			return true
+		}
+	}
+	return false
+}
diff --git a/pkg/btrfs/uuid.go b/pkg/btrfs/uuid.go
index 5218a3f..b9e3e0c 100644
--- a/pkg/btrfs/uuid.go
+++ b/pkg/btrfs/uuid.go
@@ -1,6 +1,7 @@
 package btrfs
 
 import (
+	"bytes"
 	"encoding/hex"
 	"strings"
 )
@@ -17,3 +18,7 @@ func (uuid UUID) String() string {
 		str[20:32],
 	}, "-")
 }
+
+func (a UUID) Equal(b UUID) bool {
+	return bytes.Equal(a[:], b[:])
+}