elf: Iterate over headers w/o need for allocator

1 files changed, 135 insertions, 128 deletions

diff --git a/lib/std/elf.zig b/lib/std/elf.zig
index 469ab693c7..9765bedc29 100644
--- a/lib/std/elf.zig
+++ b/lib/std/elf.zig
@@ -341,6 +341,20 @@ const Header = struct {
     shentsize: u16,
     shnum: u16,
     shstrndx: u16,
+
+    pub fn program_header_iterator(self: Header, file: File) ProgramHeaderIterator {
+        return .{
+            .elf_header = self,
+            .file = file,
+        };
+    }
+
+    pub fn section_header_iterator(self: Header, file: File) SectionHeaderIterator {
+        return .{
+            .elf_header = self,
+            .file = file,
+        };
+    }
 };
 
 pub fn readHeader(file: File) !Header {
@@ -378,144 +392,137 @@ pub fn readHeader(file: File) !Header {
     });
 }
 
-/// All integers are native endian.
-pub const AllHeaders = struct {
-    header: Header,
-    section_headers: []Elf64_Shdr,
-    program_headers: []Elf64_Phdr,
-    allocator: *mem.Allocator,
-};
-
-pub fn readAllHeaders(allocator: *mem.Allocator, file: File) !AllHeaders {
-    var hdrs: AllHeaders = .{
-        .allocator = allocator,
-        .header = try readHeader(file),
-        .section_headers = undefined,
-        .program_headers = undefined,
-    };
-    const is_64 = hdrs.header.is_64;
-    const need_bswap = hdrs.header.endian != std.builtin.endian;
-
-    hdrs.section_headers = try allocator.alloc(Elf64_Shdr, hdrs.header.shnum);
-    errdefer allocator.free(hdrs.section_headers);
-
-    hdrs.program_headers = try allocator.alloc(Elf64_Phdr, hdrs.header.phnum);
-    errdefer allocator.free(hdrs.program_headers);
-
-    // If the ELF file is 64-bit and same-endianness, then all we have to do is
-    // yeet the bytes into memory.
-    // If only the endianness is different, they can be simply byte swapped.
-    if (is_64) {
-        const shdr_buf = std.mem.sliceAsBytes(hdrs.section_headers);
-        const phdr_buf = std.mem.sliceAsBytes(hdrs.program_headers);
-        try preadNoEof(file, shdr_buf, hdrs.header.shoff);
-        try preadNoEof(file, phdr_buf, hdrs.header.phoff);
+pub const ProgramHeaderIterator = struct {
+    elf_header: Header,
+    file: File,
+    index: usize = 0,
+
+    pub fn next(self: *ProgramHeaderIterator) !?Elf64_Phdr {
+        if (self.index >= self.elf_header.phnum) return null;
+        defer self.index += 1;
+
+        if (self.elf_header.is_64) {
+            var phdr: Elf64_Phdr = undefined;
+            const offset = self.elf_header.phoff + @sizeOf(@TypeOf(phdr)) * self.index;
+            try preadNoEof(self.file, mem.asBytes(&phdr), offset);
+
+            // ELF endianness matches native endianness.
+            if (self.elf_header.endian == std.builtin.endian) return phdr;
+
+            // Convert fields to native endianness.
+            return Elf64_Phdr{
+                .p_type = @byteSwap(@TypeOf(phdr.p_type), phdr.p_type),
+                .p_offset = @byteSwap(@TypeOf(phdr.p_offset), phdr.p_offset),
+                .p_vaddr = @byteSwap(@TypeOf(phdr.p_vaddr), phdr.p_vaddr),
+                .p_paddr = @byteSwap(@TypeOf(phdr.p_paddr), phdr.p_paddr),
+                .p_filesz = @byteSwap(@TypeOf(phdr.p_filesz), phdr.p_filesz),
+                .p_memsz = @byteSwap(@TypeOf(phdr.p_memsz), phdr.p_memsz),
+                .p_flags = @byteSwap(@TypeOf(phdr.p_flags), phdr.p_flags),
+                .p_align = @byteSwap(@TypeOf(phdr.p_align), phdr.p_align),
+            };
+        }
 
-        if (need_bswap) {
-            for (hdrs.section_headers) |*shdr| {
-                shdr.* = .{
-                    .sh_name = @byteSwap(@TypeOf(shdr.sh_name), shdr.sh_name),
-                    .sh_type = @byteSwap(@TypeOf(shdr.sh_type), shdr.sh_type),
-                    .sh_flags = @byteSwap(@TypeOf(shdr.sh_flags), shdr.sh_flags),
-                    .sh_addr = @byteSwap(@TypeOf(shdr.sh_addr), shdr.sh_addr),
-                    .sh_offset = @byteSwap(@TypeOf(shdr.sh_offset), shdr.sh_offset),
-                    .sh_size = @byteSwap(@TypeOf(shdr.sh_size), shdr.sh_size),
-                    .sh_link = @byteSwap(@TypeOf(shdr.sh_link), shdr.sh_link),
-                    .sh_info = @byteSwap(@TypeOf(shdr.sh_info), shdr.sh_info),
-                    .sh_addralign = @byteSwap(@TypeOf(shdr.sh_addralign), shdr.sh_addralign),
-                    .sh_entsize = @byteSwap(@TypeOf(shdr.sh_entsize), shdr.sh_entsize),
-                };
-            }
-            for (hdrs.program_headers) |*phdr| {
-                phdr.* = .{
-                    .p_type = @byteSwap(@TypeOf(phdr.p_type), phdr.p_type),
-                    .p_offset = @byteSwap(@TypeOf(phdr.p_offset), phdr.p_offset),
-                    .p_vaddr = @byteSwap(@TypeOf(phdr.p_vaddr), phdr.p_vaddr),
-                    .p_paddr = @byteSwap(@TypeOf(phdr.p_paddr), phdr.p_paddr),
-                    .p_filesz = @byteSwap(@TypeOf(phdr.p_filesz), phdr.p_filesz),
-                    .p_memsz = @byteSwap(@TypeOf(phdr.p_memsz), phdr.p_memsz),
-                    .p_flags = @byteSwap(@TypeOf(phdr.p_flags), phdr.p_flags),
-                    .p_align = @byteSwap(@TypeOf(phdr.p_align), phdr.p_align),
-                };
-            }
+        var phdr: Elf32_Phdr = undefined;
+        const offset = self.elf_header.phoff + @sizeOf(@TypeOf(phdr)) * self.index;
+        try preadNoEof(self.file, mem.asBytes(&phdr), offset);
+
+        // ELF endianness does NOT match native endianness.
+        if (self.elf_header.endian != std.builtin.endian) {
+            // Convert fields to native endianness.
+            phdr = .{
+                .p_type = @byteSwap(@TypeOf(phdr.p_type), phdr.p_type),
+                .p_offset = @byteSwap(@TypeOf(phdr.p_offset), phdr.p_offset),
+                .p_vaddr = @byteSwap(@TypeOf(phdr.p_vaddr), phdr.p_vaddr),
+                .p_paddr = @byteSwap(@TypeOf(phdr.p_paddr), phdr.p_paddr),
+                .p_filesz = @byteSwap(@TypeOf(phdr.p_filesz), phdr.p_filesz),
+                .p_memsz = @byteSwap(@TypeOf(phdr.p_memsz), phdr.p_memsz),
+                .p_flags = @byteSwap(@TypeOf(phdr.p_flags), phdr.p_flags),
+                .p_align = @byteSwap(@TypeOf(phdr.p_align), phdr.p_align),
+            };
         }
 
-        return hdrs;
+        // Convert 32-bit header to 64-bit.
+        return Elf64_Phdr{
+            .p_type = phdr.p_type,
+            .p_offset = phdr.p_offset,
+            .p_vaddr = phdr.p_vaddr,
+            .p_paddr = phdr.p_paddr,
+            .p_filesz = phdr.p_filesz,
+            .p_memsz = phdr.p_memsz,
+            .p_flags = phdr.p_flags,
+            .p_align = phdr.p_align,
+        };
     }
+};
 
-    const shdrs_32 = try allocator.alloc(Elf32_Shdr, hdrs.header.shnum);
-    defer allocator.free(shdrs_32);
-
-    const phdrs_32 = try allocator.alloc(Elf32_Phdr, hdrs.header.phnum);
-    defer allocator.free(phdrs_32);
-
-    const shdr_buf = std.mem.sliceAsBytes(shdrs_32);
-    const phdr_buf = std.mem.sliceAsBytes(phdrs_32);
-    try preadNoEof(file, shdr_buf, hdrs.header.shoff);
-    try preadNoEof(file, phdr_buf, hdrs.header.phoff);
-
-    if (need_bswap) {
-        for (hdrs.section_headers) |*shdr, i| {
-            const o = shdrs_32[i];
-            shdr.* = .{
-                .sh_name = @byteSwap(@TypeOf(o.sh_name), o.sh_name),
-                .sh_type = @byteSwap(@TypeOf(o.sh_type), o.sh_type),
-                .sh_flags = @byteSwap(@TypeOf(o.sh_flags), o.sh_flags),
-                .sh_addr = @byteSwap(@TypeOf(o.sh_addr), o.sh_addr),
-                .sh_offset = @byteSwap(@TypeOf(o.sh_offset), o.sh_offset),
-                .sh_size = @byteSwap(@TypeOf(o.sh_size), o.sh_size),
-                .sh_link = @byteSwap(@TypeOf(o.sh_link), o.sh_link),
-                .sh_info = @byteSwap(@TypeOf(o.sh_info), o.sh_info),
-                .sh_addralign = @byteSwap(@TypeOf(o.sh_addralign), o.sh_addralign),
-                .sh_entsize = @byteSwap(@TypeOf(o.sh_entsize), o.sh_entsize),
-            };
-        }
-        for (hdrs.program_headers) |*phdr, i| {
-            const o = phdrs_32[i];
-            phdr.* = .{
-                .p_type = @byteSwap(@TypeOf(o.p_type), o.p_type),
-                .p_offset = @byteSwap(@TypeOf(o.p_offset), o.p_offset),
-                .p_vaddr = @byteSwap(@TypeOf(o.p_vaddr), o.p_vaddr),
-                .p_paddr = @byteSwap(@TypeOf(o.p_paddr), o.p_paddr),
-                .p_filesz = @byteSwap(@TypeOf(o.p_filesz), o.p_filesz),
-                .p_memsz = @byteSwap(@TypeOf(o.p_memsz), o.p_memsz),
-                .p_flags = @byteSwap(@TypeOf(o.p_flags), o.p_flags),
-                .p_align = @byteSwap(@TypeOf(o.p_align), o.p_align),
+pub const SectionHeaderIterator = struct {
+    elf_header: Header,
+    file: File,
+    index: usize = 0,
+
+    pub fn next(self: *SectionHeaderIterator) !?Elf64_Shdr {
+        if (self.index >= self.elf_header.shnum) return null;
+        defer self.index += 1;
+
+        if (self.elf_header.is_64) {
+            var shdr: Elf64_Shdr = undefined;
+            const offset = self.elf_header.phoff + @sizeOf(@TypeOf(shdr)) * self.index;
+            try preadNoEof(self.file, mem.asBytes(&shdr), offset);
+
+            // ELF endianness matches native endianness.
+            if (self.elf_header.endian == std.builtin.endian) return shdr;
+
+            // Convert fields to native endianness.
+            return Elf64_Shdr{
+                .sh_name = @byteSwap(@TypeOf(shdr.sh_name), shdr.sh_name),
+                .sh_type = @byteSwap(@TypeOf(shdr.sh_type), shdr.sh_type),
+                .sh_flags = @byteSwap(@TypeOf(shdr.sh_flags), shdr.sh_flags),
+                .sh_addr = @byteSwap(@TypeOf(shdr.sh_addr), shdr.sh_addr),
+                .sh_offset = @byteSwap(@TypeOf(shdr.sh_offset), shdr.sh_offset),
+                .sh_size = @byteSwap(@TypeOf(shdr.sh_size), shdr.sh_size),
+                .sh_link = @byteSwap(@TypeOf(shdr.sh_link), shdr.sh_link),
+                .sh_info = @byteSwap(@TypeOf(shdr.sh_info), shdr.sh_info),
+                .sh_addralign = @byteSwap(@TypeOf(shdr.sh_addralign), shdr.sh_addralign),
+                .sh_entsize = @byteSwap(@TypeOf(shdr.sh_entsize), shdr.sh_entsize),
             };
         }
-    } else {
-        for (hdrs.section_headers) |*shdr, i| {
-            const o = shdrs_32[i];
-            shdr.* = .{
-                .sh_name = o.sh_name,
-                .sh_type = o.sh_type,
-                .sh_flags = o.sh_flags,
-                .sh_addr = o.sh_addr,
-                .sh_offset = o.sh_offset,
-                .sh_size = o.sh_size,
-                .sh_link = o.sh_link,
-                .sh_info = o.sh_info,
-                .sh_addralign = o.sh_addralign,
-                .sh_entsize = o.sh_entsize,
-            };
-        }
-        for (hdrs.program_headers) |*phdr, i| {
-            const o = phdrs_32[i];
-            phdr.* = .{
-                .p_type = o.p_type,
-                .p_offset = o.p_offset,
-                .p_vaddr = o.p_vaddr,
-                .p_paddr = o.p_paddr,
-                .p_filesz = o.p_filesz,
-                .p_memsz = o.p_memsz,
-                .p_flags = o.p_flags,
-                .p_align = o.p_align,
+
+        var shdr: Elf32_Shdr = undefined;
+        const offset = self.elf_header.shoff + @sizeOf(@TypeOf(shdr)) * self.index;
+        try preadNoEof(self.file, mem.asBytes(&shdr), offset);
+
+        // ELF endianness does NOT match native endianness.
+        if (self.elf_header.endian != std.builtin.endian) {
+            // Convert fields to native endianness.
+            shdr = .{
+                .sh_name = @byteSwap(@TypeOf(shdr.sh_name), shdr.sh_name),
+                .sh_type = @byteSwap(@TypeOf(shdr.sh_type), shdr.sh_type),
+                .sh_flags = @byteSwap(@TypeOf(shdr.sh_flags), shdr.sh_flags),
+                .sh_addr = @byteSwap(@TypeOf(shdr.sh_addr), shdr.sh_addr),
+                .sh_offset = @byteSwap(@TypeOf(shdr.sh_offset), shdr.sh_offset),
+                .sh_size = @byteSwap(@TypeOf(shdr.sh_size), shdr.sh_size),
+                .sh_link = @byteSwap(@TypeOf(shdr.sh_link), shdr.sh_link),
+                .sh_info = @byteSwap(@TypeOf(shdr.sh_info), shdr.sh_info),
+                .sh_addralign = @byteSwap(@TypeOf(shdr.sh_addralign), shdr.sh_addralign),
+                .sh_entsize = @byteSwap(@TypeOf(shdr.sh_entsize), shdr.sh_entsize),
             };
         }
-    }
 
-    return hdrs;
-}
+        // Convert 32-bit header to 64-bit.
+        return Elf64_Shdr{
+            .sh_name = shdr.sh_name,
+            .sh_type = shdr.sh_type,
+            .sh_flags = shdr.sh_flags,
+            .sh_addr = shdr.sh_addr,
+            .sh_offset = shdr.sh_offset,
+            .sh_size = shdr.sh_size,
+            .sh_link = shdr.sh_link,
+            .sh_info = shdr.sh_info,
+            .sh_addralign = shdr.sh_addralign,
+            .sh_entsize = shdr.sh_entsize,
+        };
+    }
+};
 
 pub fn int(is_64: bool, need_bswap: bool, int_32: anytype, int_64: anytype) @TypeOf(int_64) {
     if (is_64) {