Parse an ELF file in Zig

elf.zig

const std = @import("std");

pub const ELF_MAGIC = [_]u8{ 0x7F, 0x45, 0x4C, 0x46 };

pub const ElfClass = enum(u8) {
    Elf32 = 1,
    Elf64 = 2,
};

pub const ElfData = enum(u8) {
    ElfLittle = 1,
    ElfBig = 2,

    pub fn to_zig_endian(self: ElfData) std.builtin.Endian {
        return switch (self) {
            ElfData.ElfLittle => .little,
            ElfData.ElfBig => .big,
        };
    }
};

pub const ElfOsAbi = enum(u8) {
    SystemV = 0,
    HP_UX = 1,
    NetBSD = 2,
    Linux = 3,
    GNU_Hurd = 4,
    Solaris = 6,
    AIX = 7,
    IRIX = 8,
    FreeBSD = 9,
    Tru64 = 10,
    NovellModesto = 11,
    OpenBSD = 12,
    OpenVMS = 13,
    NonStopKernel = 14,
    AROS = 15,
    FenixOS = 16,
    CloudABI = 17,
    StratusTechnologiesOpenVOS = 18,
};

pub const ElfType = enum(u16) {
    None = 0,
    Relocatable = 1,
    Executable = 2,
    SharedObject = 3,
    Core = 4,
    LoOS = 0xFE00,
    HiOS = 0xFEFF,
    LoProcessor = 0xFF00,
    HiProcessor = 0xFFFF,
};

// NOTE: This is not a complete list of machine types
pub const ElfMachine = enum(u16) {
    NoMachine = 0,
    SPARC = 2,
    x86 = 3,
    MIPS = 8,
    PowerPC = 20,
    ARM = 40,
    SuperH = 42,
    IA_64 = 50,
    x86_64 = 62,
    AArch64 = 183,
};

pub const ElfHeader = extern struct {
    magic: [4]u8,
    class: ElfClass,
    data: ElfData,
    version: u8,
    os_abi: ElfOsAbi,
    abi_version: u8,
    padding: [7]u8,
};

pub const ElfHeaderTable = union(ElfClass) {
    Elf32: Elf32HeaderTable,
    Elf64: Elf64HeaderTable,
};

pub const Elf32HeaderTable = extern struct {
    elf_type: ElfType,
    machine: ElfMachine,
    version2: u32,
    entry: u32,
    program_header_offset: u32,
    section_header_offset: u32,
    flags: u32,
    header_size: u16,
    program_header_entry_size: u16,
    program_header_count: u16,
    section_header_entry_size: u16,
    section_header_count: u16,
    section_header_string_index: u16,
};

pub const Elf64HeaderTable = extern struct {
    elf_type: ElfType,
    machine: ElfMachine,
    version2: u32,
    entry: u64,
    program_header_offset: u64,
    section_header_offset: u64,
    flags: u32,
    header_size: u16,
    program_header_entry_size: u16,
    program_header_count: u16,
    section_header_entry_size: u16,
    section_header_count: u16,
    section_header_string_index: u16,
};

pub const ElfPHEntry = union(ElfClass) {
    Elf32: Elf32ProgramHeaderEntry,
    Elf64: Elf64ProgramHeaderEntry,
};

pub const ProgramFlags = packed struct(u32) {
    Executable: bool,
    Writable: bool,
    Readable: bool,
    reserved: u29,
};

pub const Elf32ProgramHeaderEntry = extern struct {
    type: u32,
    offset: u32,
    virtual_address: u32,
    physical_address: u32,
    file_size: u32,
    memory_size: u32,
    flags: u32,
    alignment: u32,

    pub fn program_flags(self: Elf32ProgramHeaderEntry) ProgramFlags {
        return @bitCast(self.flags);
    }
};

pub const Elf64ProgramHeaderEntry = extern struct {
    type: u32,
    flags: u32,
    offset: u64,
    virtual_address: u64,
    physical_address: u64,
    file_size: u64,
    memory_size: u64,
    alignment: u64,

    pub fn program_flags(self: Elf32ProgramHeaderEntry) ProgramFlags {
        return @bitCast(self.flags);
    }
};

pub const ElfSectionEntry = union(ElfClass) {
    Elf32: Elf32SectionEntry,
    Elf64: Elf64SectionEntry,
};

pub const Elf32SectionEntry = extern struct {
    name: u32,
    type: u32,
    flags: u32,
    address: u32,
    offset: u32,
    size: u32,
    link: u32,
    info: u32,
    address_alignment: u32,
    entry_size: u32,
};

pub const Elf64SectionEntry = extern struct {
    name: u32,
    type: u32,
    flags: u64,
    address: u64,
    offset: u64,
    size: u64,
    link: u32,
    info: u32,
    address_alignment: u64,
    entry_size: u64,
};

pub const ELF = struct {
    header: ElfHeader,
    header_table: ElfHeaderTable,
    program_header: []ElfPHEntry,
    section_header: []ElfSectionEntry,

    pub fn parse(allocator: std.mem.Allocator, data: []const u8) !ELF {
        var stream = std.io.fixedBufferStream(data);
        var reader = stream.reader();

        const header = try reader.readStruct(ElfHeader);

        if (!std.mem.eql(u8, &header.magic, &ELF_MAGIC)) {
            return error.ELFInvalidMagic;
        }

        const endian = header.data.to_zig_endian();
        const header_table = switch (header.class) {
            .Elf32 => @unionInit(ElfHeaderTable, @tagName(.Elf32), try reader.readStructEndian(Elf32HeaderTable, endian)),
            .Elf64 => @unionInit(ElfHeaderTable, @tagName(.Elf64), try reader.readStructEndian(Elf64HeaderTable, endian)),
        };

        const active = std.meta.activeTag(header_table);

        const ph_ent_size = switch (header_table) {
            .Elf32 => |h| h.program_header_entry_size,
            .Elf64 => |h| h.program_header_entry_size,
        };

        const ph_ent_count = switch (header_table) {
            .Elf32 => |h| h.program_header_count,
            .Elf64 => |h| h.program_header_count,
        };

        const ph_offset = switch (header_table) {
            .Elf32 => |h| h.program_header_offset,
            .Elf64 => |h| h.program_header_offset,
        };

        const sh_ent_size = switch (header_table) {
            .Elf32 => |h| h.section_header_entry_size,
            .Elf64 => |h| h.section_header_entry_size,
        };

        const sh_ent_count = switch (header_table) {
            .Elf32 => |h| h.section_header_count,
            .Elf64 => |h| h.section_header_count,
        };

        const sh_offset = switch (header_table) {
            .Elf32 => |h| h.section_header_offset,
            .Elf64 => |h| h.section_header_offset,
        };

        if (active == .Elf32 and ph_ent_size != @sizeOf(Elf32ProgramHeaderEntry) or active == .Elf64 and ph_ent_size != @sizeOf(Elf64ProgramHeaderEntry)) {
            @panic("Invalid program header entry size! (Maybe nonstandard ELF file?)");
        }

        if (active == .Elf32 and sh_ent_size != @sizeOf(Elf32SectionEntry) or active == .Elf64 and sh_ent_size != @sizeOf(Elf64SectionEntry)) {
            @panic("Invalid section header entry size! (Maybe nonstandard ELF file?)");
        }

        stream = std.io.fixedBufferStream(data[ph_offset..]);
        reader = stream.reader();

        var entries = std.ArrayList(ElfPHEntry).init(allocator);
        for (0..ph_ent_count) |_| {
            try entries.append(switch (header.class) {
                .Elf32 => @unionInit(ElfPHEntry, @tagName(.Elf32), try reader.readStructEndian(Elf32ProgramHeaderEntry, endian)),
                .Elf64 => @unionInit(ElfPHEntry, @tagName(.Elf64), try reader.readStructEndian(Elf64ProgramHeaderEntry, endian)),
            });
        }

        stream = std.io.fixedBufferStream(data[sh_offset..]);
        reader = stream.reader();

        var sections = std.ArrayList(ElfSectionEntry).init(allocator);
        for (0..sh_ent_count) |_| {
            try sections.append(switch (header.class) {
                .Elf32 => @unionInit(ElfSectionEntry, @tagName(.Elf32), try reader.readStructEndian(Elf32SectionEntry, endian)),
                .Elf64 => @unionInit(ElfSectionEntry, @tagName(.Elf64), try reader.readStructEndian(Elf64SectionEntry, endian)),
            });
        }

        return ELF{
            .header = header,
            .header_table = header_table,
            .program_header = try entries.toOwnedSlice(),
            .section_header = try sections.toOwnedSlice(),
        };
    }
};