Add deep extent tree support (depth >= 2) in EXT4 reader and formatter

Author: gbfish

Sources/ContainerizationEXT4/EXT4+Formatter.swift

@@ -1039,6 +1039,130 @@ extension EXT4 {
             }
         }
 
+        /// Recursively writes extent tree nodes to disk and returns the
+        /// ExtentIndex entries that the parent should store.
+        ///
+        /// At depth 1, this writes leaf blocks (containing ExtentLeaf entries).
+        /// At depth > 1, this recurses to build child subtrees first, then
+        /// writes index blocks pointing to those children.
+        private func writeExtentSubtree(
+            depth: UInt16,
+            numExtents: UInt32,
+            numBlocks: UInt32,
+            start: UInt32,
+            entriesPerBlock: UInt32,
+            extentBlockCount: inout UInt32
+        ) throws -> [ExtentIndex] {
+            var indices: [ExtentIndex] = []
+            let childDepth = depth - 1
+
+            // How many child blocks do we need at this level?
+            // Each child block can ultimately cover entriesPerBlock^childDepth leaf extents.
+            var leafCapacityPerChild: UInt32 = entriesPerBlock
+            for _ in 1..<depth {
+                leafCapacityPerChild *= entriesPerBlock
+            }
+            let numChildren = (numExtents + leafCapacityPerChild - 1) / leafCapacityPerChild
+            var extentsWritten: UInt32 = 0
+
+            for _ in 0..<numChildren {
+                let extentsForChild = min(numExtents - extentsWritten, leafCapacityPerChild)
+                let logicalOffset = extentsWritten * EXT4.MaxBlocksPerExtent
+
+                if childDepth == 0 {
+                    // Write a leaf block containing ExtentLeaf entries.
+                    if self.pos % self.blockSize != 0 {
+                        try self.seek(block: self.currentBlock + 1)
+                    }
+                    let leafBlockAddr = self.currentBlock
+                    extentBlockCount += 1
+
+                    let leafHeader = ExtentHeader(
+                        magic: EXT4.ExtentHeaderMagic,
+                        entries: UInt16(extentsForChild),
+                        max: UInt16(entriesPerBlock),
+                        depth: 0,
+                        generation: 0
+                    )
+                    var leafNode = ExtentLeafNode(header: leafHeader, leaves: [])
+                    fillExtents(
+                        node: &leafNode, numExtents: extentsForChild, numBlocks: numBlocks,
+                        start: start,
+                        offset: logicalOffset)
+                    try withUnsafeLittleEndianBytes(of: leafNode.header) { bytes in
+                        try self.handle.write(contentsOf: bytes)
+                    }
+                    for leaf in leafNode.leaves {
+                        try withUnsafeLittleEndianBytes(of: leaf) { bytes in
+                            try self.handle.write(contentsOf: bytes)
+                        }
+                    }
+                    let checksum = leafNode.leaves.last?.block ?? 0
+                    let extentTail = ExtentTail(checksum: checksum)
+                    try withUnsafeLittleEndianBytes(of: extentTail) { bytes in
+                        try self.handle.write(contentsOf: bytes)
+                    }
+
+                    indices.append(
+                        ExtentIndex(
+                            block: logicalOffset,
+                            leafLow: leafBlockAddr,
+                            leafHigh: 0,
+                            unused: 0
+                        ))
+                } else {
+                    // Recurse to build the child subtree, then write an index
+                    // block pointing to the returned child indices.
+                    let childIndices = try writeExtentSubtree(
+                        depth: childDepth,
+                        numExtents: extentsForChild,
+                        numBlocks: numBlocks,
+                        start: start,
+                        entriesPerBlock: entriesPerBlock,
+                        extentBlockCount: &extentBlockCount
+                    )
+
+                    // Write the index block for this subtree.
+                    if self.pos % self.blockSize != 0 {
+                        try self.seek(block: self.currentBlock + 1)
+                    }
+                    let indexBlockAddr = self.currentBlock
+                    extentBlockCount += 1
+
+                    let indexHeader = ExtentHeader(
+                        magic: EXT4.ExtentHeaderMagic,
+                        entries: UInt16(childIndices.count),
+                        max: UInt16(entriesPerBlock),
+                        depth: childDepth,
+                        generation: 0
+                    )
+                    try withUnsafeLittleEndianBytes(of: indexHeader) { bytes in
+                        try self.handle.write(contentsOf: bytes)
+                    }
+                    for childIdx in childIndices {
+                        try withUnsafeLittleEndianBytes(of: childIdx) { bytes in
+                            try self.handle.write(contentsOf: bytes)
+                        }
+                    }
+                    let checksum = childIndices.last?.block ?? 0
+                    let extentTail = ExtentTail(checksum: checksum)
+                    try withUnsafeLittleEndianBytes(of: extentTail) { bytes in
+                        try self.handle.write(contentsOf: bytes)
+                    }
+
+                    indices.append(
+                        ExtentIndex(
+                            block: logicalOffset,
+                            leafLow: indexBlockAddr,
+                            leafHigh: 0,
+                            unused: 0
+                        ))
+                }
+                extentsWritten += extentsForChild
+            }
+            return indices
+        }
+
         private func writeExtents(_ inode: Inode, _ blocks: (start: UInt32, end: UInt32)) throws -> Inode {
             var inode = inode
             // rest of code assumes that extents MUST go into a new block
@@ -1048,14 +1172,31 @@ extension EXT4 {
             let dataBlocks = blocks.end - blocks.start
             let numExtents = (dataBlocks + EXT4.MaxBlocksPerExtent - 1) / EXT4.MaxBlocksPerExtent
             var usedBlocks = dataBlocks
-            let extentNodeSize = 12
-            let extentsPerBlock = self.blockSize / extentNodeSize - 1
+            let extentNodeSize: UInt32 = 12
+            let entriesPerBlock = self.blockSize / extentNodeSize - 1
             var blockData: [UInt8] = .init(repeating: 0, count: 60)
             var blockIndex: Int = 0
-            switch numExtents {
-            case 0:
+
+            guard numExtents > 0 else {
                 return inode  // noop
-            case 1..<5:
+            }
+
+            // Determine the required tree depth.
+            // Depth 0: up to 4 extents fit inline in the inode.
+            // Depth N (N >= 1): each level multiplies capacity by entriesPerBlock,
+            // with up to 4 entries at the root.
+            var depth: UInt16 = 0
+            var capacity: UInt32 = 4
+            while capacity < numExtents {
+                depth += 1
+                capacity = 4
+                for _ in 0..<depth {
+                    capacity *= entriesPerBlock
+                }
+            }
+
+            if depth == 0 {
+                // All extents fit inline in the inode's 60-byte block field.
                 let extentHeader = ExtentHeader(
                     magic: EXT4.ExtentHeaderMagic,
                     entries: UInt16(numExtents),
@@ -1079,73 +1220,42 @@ extension EXT4 {
                         }
                     }
                 }
-            case 5..<4 * UInt32(extentsPerBlock) + 1:
-                let extentBlocks = (numExtents + extentsPerBlock - 1) / extentsPerBlock
-                usedBlocks += extentBlocks
+            } else {
+                // Build the extent tree bottom-up. writeExtentSubtree writes
+                // child blocks to disk and returns the index entries for the
+                // parent level.
+                var extentBlockCount: UInt32 = 0
+                let rootIndices = try writeExtentSubtree(
+                    depth: depth,
+                    numExtents: numExtents,
+                    numBlocks: dataBlocks,
+                    start: blocks.start,
+                    entriesPerBlock: entriesPerBlock,
+                    extentBlockCount: &extentBlockCount
+                )
+                usedBlocks += extentBlockCount
+
                 let extentHeader = ExtentHeader(
                     magic: EXT4.ExtentHeaderMagic,
-                    entries: UInt16(extentBlocks),
+                    entries: UInt16(rootIndices.count),
                     max: 4,
-                    depth: 1,
+                    depth: depth,
                     generation: 0
                 )
-                var root = ExtentIndexNode(header: extentHeader, indices: [])
-                for i in 0..<extentBlocks {
-                    if self.pos % self.blockSize != 0 {
-                        try self.seek(block: self.currentBlock + 1)
-                    }
-                    let extentIdx = ExtentIndex(
-                        block: i * extentsPerBlock * EXT4.MaxBlocksPerExtent,
-                        leafLow: self.currentBlock,
-                        leafHigh: 0,
-                        unused: 0)
-                    var extentsInBlock = numExtents - i * extentsPerBlock
-                    if extentsInBlock > extentsPerBlock {
-                        extentsInBlock = extentsPerBlock
-                    }
-                    let leafHeader = ExtentHeader(
-                        magic: EXT4.ExtentHeaderMagic,
-                        entries: UInt16(extentsInBlock),
-                        max: UInt16(extentsPerBlock),
-                        depth: 0,
-                        generation: 0
-                    )
-                    var leafNode = ExtentLeafNode(header: leafHeader, leaves: [])
-                    let offset = i * extentsPerBlock * EXT4.MaxBlocksPerExtent
-                    fillExtents(
-                        node: &leafNode, numExtents: extentsInBlock, numBlocks: dataBlocks,
-                        start: blocks.start,
-                        offset: offset)
-                    try withUnsafeLittleEndianBytes(of: leafNode.header) { bytes in
-                        try self.handle.write(contentsOf: bytes)
-                    }
-                    for leaf in leafNode.leaves {
-                        try withUnsafeLittleEndianBytes(of: leaf) { bytes in
-                            try self.handle.write(contentsOf: bytes)
-                        }
-                    }
-                    let extentTail = ExtentTail(checksum: leafNode.leaves.last!.block)
-                    try withUnsafeLittleEndianBytes(of: extentTail) { bytes in
-                        try self.handle.write(contentsOf: bytes)
-                    }
-                    root.indices.append(extentIdx)
-                }
-                withUnsafeLittleEndianBytes(of: root.header) { bytes in
+                withUnsafeLittleEndianBytes(of: extentHeader) { bytes in
                     for b in bytes {
                         blockData[blockIndex] = b
                         blockIndex = blockIndex + 1
                     }
                 }
-                for leaf in root.indices {
-                    withUnsafeLittleEndianBytes(of: leaf) { bytes in
+                for idx in rootIndices {
+                    withUnsafeLittleEndianBytes(of: idx) { bytes in
                         for b in bytes {
                             blockData[blockIndex] = b
                             blockIndex = blockIndex + 1
                         }
                     }
                 }
-            default:
-                throw Error.fileTooBig(UInt64(dataBlocks) * self.blockSize)
             }
             inode.block = (
                 blockData[0], blockData[1], blockData[2], blockData[3], blockData[4], blockData[5], blockData[6],

Sources/ContainerizationEXT4/EXT4+Reader.swift

@@ -196,61 +196,61 @@ extension EXT4 {
         func getExtents(inode: InodeNumber) throws -> [(start: UInt32, end: UInt32)]? {
             let inode = try self.getInode(number: inode)
             let inodeBlock = Data(tupleToArray(inode.block))
-            var offset = 0
-            var extents: [(start: UInt32, end: UInt32)] = []
 
             let extentHeaderSize = MemoryLayout<ExtentHeader>.size
-            let extentIndexSize = MemoryLayout<ExtentIndex>.size
-            let extentLeafSize = MemoryLayout<ExtentLeaf>.size
-            // read extent header
-            let header = inodeBlock.subdata(in: offset..<offset + extentHeaderSize).withUnsafeBytes {
+            let header = inodeBlock.subdata(in: 0..<extentHeaderSize).withUnsafeBytes {
                 $0.loadLittleEndian(as: ExtentHeader.self)
             }
             guard header.magic == EXT4.ExtentHeaderMagic else {
                 return []
             }
-            offset += extentHeaderSize  // Jump to entries
-            switch header.depth {
-            case 0:
-                // When depth is 0 the extent header is followed by extent leaves
+            var extents: [(start: UInt32, end: UInt32)] = []
+            try readExtentNode(data: inodeBlock, header: header, into: &extents)
+            return extents
+        }
+
+        /// Recursively reads extent tree nodes. For leaf nodes (depth 0), appends
+        /// extent mappings directly. For index nodes (depth > 0), follows each
+        /// index entry to a child block and recurses.
+        private func readExtentNode(
+            data: Data,
+            header: ExtentHeader,
+            into extents: inout [(start: UInt32, end: UInt32)]
+        ) throws {
+            let extentHeaderSize = MemoryLayout<ExtentHeader>.size
+            let extentIndexSize = MemoryLayout<ExtentIndex>.size
+            let extentLeafSize = MemoryLayout<ExtentLeaf>.size
+            var offset = extentHeaderSize
+
+            if header.depth == 0 {
+                // Leaf node: entries are ExtentLeaf mappings
                 for _ in 0..<header.entries {
-                    let leaf = inodeBlock.subdata(in: offset..<offset + extentLeafSize).withUnsafeBytes {
-                        $0.load(as: ExtentLeaf.self)
+                    let leaf = data.subdata(in: offset..<offset + extentLeafSize).withUnsafeBytes {
+                        $0.loadLittleEndian(as: ExtentLeaf.self)
                     }
                     extents.append((leaf.startLow, leaf.startLow + UInt32(leaf.length)))
                     offset += extentLeafSize
                 }
-            case 1:
-                // When depth is 1 the extent header is followed by extent indices which point to leaves
+            } else {
+                // Index node: entries are ExtentIndex pointers to child blocks
                 for _ in 0..<header.entries {
-                    let indexNode = inodeBlock.subdata(in: offset..<offset + extentIndexSize).withUnsafeBytes {
-                        $0.load(as: ExtentIndex.self)
+                    let indexNode = data.subdata(in: offset..<offset + extentIndexSize).withUnsafeBytes {
+                        $0.loadLittleEndian(as: ExtentIndex.self)
                     }
                     try self.seek(block: indexNode.leafLow)
                     guard let block = try self.handle.read(upToCount: Int(self.blockSize)) else {
                         throw EXT4.Error.couldNotReadBlock(indexNode.leafLow)
                     }
-                    var blockOffset = 0
-                    let leafHeader = block.subdata(in: blockOffset..<extentHeaderSize).withUnsafeBytes {
+                    let childHeader = block.subdata(in: 0..<extentHeaderSize).withUnsafeBytes {
                         $0.loadLittleEndian(as: ExtentHeader.self)
                     }
-                    guard leafHeader.magic == EXT4.ExtentHeaderMagic else {
+                    guard childHeader.magic == EXT4.ExtentHeaderMagic else {
                         throw Error.invalidExtents
                     }
-                    blockOffset += extentHeaderSize
-                    for _ in 0..<leafHeader.entries {
-                        let leaf = block.subdata(in: blockOffset..<blockOffset + extentLeafSize).withUnsafeBytes {
-                            $0.loadLittleEndian(as: ExtentLeaf.self)
-                        }
-                        extents.append((leaf.startLow, leaf.startLow + UInt32(leaf.length)))
-                        blockOffset += extentLeafSize
-                    }
+                    try readExtentNode(data: block, header: childHeader, into: &extents)
                     offset += extentIndexSize
                 }
-            default:
-                throw Error.deepExtentsUnimplemented
             }
-            return extents
         }
 
         // MARK: Internal functions

Tests/ContainerizationEXT4Tests/TestEXT4Reader+IO.swift

@@ -548,6 +548,38 @@ struct EXT4PathIOTests {
         #expect(String(decoding: partialData, as: UTF8.self) == expectedContent)
     }
 
+    @Test
+    func multiExtentFileReadback() throws {
+        // Create an image large enough that the formatter produces multiple extents
+        // by interleaving file creation (which can cause non-contiguous allocation).
+        // This exercises the recursive readExtentNode() path for depth 0 and depth 1.
+        let url = try buildFS(minDiskSize: 32 * 1024 * 1024) { fmt in
+            // Create many files to consume blocks, then a large file that may
+            // span multiple extents due to intervening allocations.
+            for i in 0..<50 {
+                let content = String(repeating: Character(UnicodeScalar(65 + (i % 26))!), count: 8192)
+                try self.createFile(fmt, "/filler_\(i).txt", content)
+            }
+            // Create a large file — the formatter writes this contiguously, but the
+            // reader still exercises getExtents() → readExtentNode() with depth 0.
+            let bigContent = String(repeating: "X", count: 512 * 1024)
+            try self.createFile(fmt, "/big.bin", bigContent)
+        }
+        defer { try? FileManager.default.removeItem(at: url) }
+
+        let r = try openReader(url)
+
+        // Verify the large file reads back correctly
+        let data = try r.readFile(at: FilePath("/big.bin"))
+        #expect(data.count == 512 * 1024)
+        #expect(data.allSatisfy { $0 == UInt8(ascii: "X") })
+
+        // Verify filler files also read correctly
+        let filler0 = try r.readFile(at: FilePath("/filler_0.txt"))
+        #expect(filler0.count == 8192)
+        #expect(filler0.allSatisfy { $0 == UInt8(ascii: "A") })
+    }
+
     @Test
     func largeFileReadAcrossBlocks() throws {
         // Keep this modest to avoid slow CI while still crossing multiple blocks.