Coff

pub const Coff = struct {
    data: []const u8,
    // Set if `data` is backed by the image as loaded by the loader
    is_loaded: bool,
    is_image: bool,
    coff_header_offset: usize,

    guid: [16]u8 = undefined,
    age: u32 = undefined,

    // The lifetime of `data` must be longer than the lifetime of the returned Coff
    pub fn init(data: []const u8, is_loaded: bool) error{ EndOfStream, MissingPEHeader }!Coff {
        if (data.len < pe_pointer_offset + 4) return error.EndOfStream;
        const header_offset = mem.readInt(u32, data[pe_pointer_offset..][0..4], .little);
        if (data.len < header_offset + 4) return error.EndOfStream;
        const is_image = mem.eql(u8, data[header_offset..][0..4], pe_signature);

        const coff: Coff = .{
            .data = data,
            .is_image = is_image,
            .is_loaded = is_loaded,
            .coff_header_offset = o: {
                if (is_image) break :o header_offset + 4;
                break :o header_offset;
            },
        };

        // Do some basic validation upfront
        if (is_image) {
            const coff_header = coff.getHeader();
            if (coff_header.size_of_optional_header == 0) return error.MissingPEHeader;
        }

        // JK: we used to check for architecture here and throw an error if not x86 or derivative.
        // However I am willing to take a leap of faith and let aarch64 have a shot also.

        return coff;
    }

    pub fn getPdbPath(self: *Coff) !?[]const u8 {
        assert(self.is_image);

        const data_dirs = self.getDataDirectories();
        if (@backingInt(IMAGE.DIRECTORY_ENTRY.DEBUG) >= data_dirs.len) return null;

        const debug_dir = data_dirs[@backingInt(IMAGE.DIRECTORY_ENTRY.DEBUG)];
        var reader: std.Io.Reader = .fixed(self.data);

        if (self.is_loaded) {
            reader.seek = debug_dir.virtual_address;
        } else {
            // Find what section the debug_dir is in, in order to convert the RVA to a file offset
            for (self.getSectionHeaders()) |*sect| {
                if (debug_dir.virtual_address >= sect.virtual_address and debug_dir.virtual_address < sect.virtual_address + sect.virtual_size) {
                    reader.seek = sect.pointer_to_raw_data + (debug_dir.virtual_address - sect.virtual_address);
                    break;
                }
            } else return error.InvalidDebugDirectory;
        }

        // Find the correct DebugDirectoryEntry, and where its data is stored.
        // It can be in any section.
        const debug_dir_entry_count = debug_dir.size / @sizeOf(DebugDirectoryEntry);
        var i: u32 = 0;
        while (i < debug_dir_entry_count) : (i += 1) {
            const debug_dir_entry = try reader.takeStruct(DebugDirectoryEntry, .little);
            if (debug_dir_entry.type == .CODEVIEW) {
                const dir_offset = if (self.is_loaded) debug_dir_entry.address_of_raw_data else debug_dir_entry.pointer_to_raw_data;
                reader.seek = dir_offset;
                break;
            }
        } else return null;

        const code_view_signature = try reader.takeArray(4);
        // 'RSDS' indicates PDB70 format, used by lld.
        if (!mem.eql(u8, code_view_signature, "RSDS"))
            return error.InvalidPEMagic;
        try reader.readSliceAll(self.guid[0..]);
        self.age = try reader.takeInt(u32, .little);

        // Finally read the null-terminated string.
        const start = reader.seek;
        const len = std.mem.findScalar(u8, self.data[start..], 0) orelse return null;
        return self.data[start .. start + len];
    }

    pub fn getHeader(self: Coff) Header {
        return @as(*align(1) const Header, @ptrCast(self.data[self.coff_header_offset..][0..@sizeOf(Header)])).*;
    }

    pub fn getOptionalHeader(self: Coff) OptionalHeader {
        assert(self.is_image);
        const offset = self.coff_header_offset + @sizeOf(Header);
        return @as(*align(1) const OptionalHeader, @ptrCast(self.data[offset..][0..@sizeOf(OptionalHeader)])).*;
    }

    pub fn getOptionalHeader32(self: Coff) OptionalHeader.PE32 {
        assert(self.is_image);
        const offset = self.coff_header_offset + @sizeOf(Header);
        return @as(*align(1) const OptionalHeader.PE32, @ptrCast(self.data[offset..][0..@sizeOf(OptionalHeader.PE32)])).*;
    }

    pub fn getOptionalHeader64(self: Coff) OptionalHeader.@"PE32+" {
        assert(self.is_image);
        const offset = self.coff_header_offset + @sizeOf(Header);
        return @as(*align(1) const OptionalHeader.@"PE32+", @ptrCast(self.data[offset..][0..@sizeOf(OptionalHeader.@"PE32+")])).*;
    }

    pub fn getImageBase(self: Coff) u64 {
        const hdr = self.getOptionalHeader();
        return switch (@backingInt(hdr.magic)) {
            IMAGE_NT_OPTIONAL_HDR32_MAGIC => self.getOptionalHeader32().image_base,
            IMAGE_NT_OPTIONAL_HDR64_MAGIC => self.getOptionalHeader64().image_base,
            else => unreachable, // We assume we have validated the header already
        };
    }

    pub fn getNumberOfDataDirectories(self: Coff) u32 {
        const hdr = self.getOptionalHeader();
        return switch (@backingInt(hdr.magic)) {
            IMAGE_NT_OPTIONAL_HDR32_MAGIC => self.getOptionalHeader32().number_of_rva_and_sizes,
            IMAGE_NT_OPTIONAL_HDR64_MAGIC => self.getOptionalHeader64().number_of_rva_and_sizes,
            else => unreachable, // We assume we have validated the header already
        };
    }

    pub fn getDataDirectories(self: *const Coff) []align(1) const ImageDataDirectory {
        const hdr = self.getOptionalHeader();
        const size: usize = switch (@backingInt(hdr.magic)) {
            IMAGE_NT_OPTIONAL_HDR32_MAGIC => @sizeOf(OptionalHeader.PE32),
            IMAGE_NT_OPTIONAL_HDR64_MAGIC => @sizeOf(OptionalHeader.@"PE32+"),
            else => unreachable, // We assume we have validated the header already
        };
        const offset = self.coff_header_offset + @sizeOf(Header) + size;
        return @as([*]align(1) const ImageDataDirectory, @ptrCast(self.data[offset..]))[0..self.getNumberOfDataDirectories()];
    }

    pub fn getSymtab(self: *const Coff) ?Symtab {
        const coff_header = self.getHeader();
        if (coff_header.pointer_to_symbol_table == 0) return null;

        const offset = coff_header.pointer_to_symbol_table;
        const size = coff_header.number_of_symbols * Symbol.sizeOf();
        return .{ .buffer = self.data[offset..][0..size] };
    }

    pub fn getStrtab(self: *const Coff) error{InvalidStrtabSize}!?Strtab {
        const coff_header = self.getHeader();
        if (coff_header.pointer_to_symbol_table == 0) return null;

        const offset = coff_header.pointer_to_symbol_table + Symbol.sizeOf() * coff_header.number_of_symbols;
        const size = mem.readInt(u32, self.data[offset..][0..4], .little);
        if ((offset + size) > self.data.len) return error.InvalidStrtabSize;

        return Strtab{ .buffer = self.data[offset..][0..size] };
    }

    pub fn strtabRequired(self: *const Coff) bool {
        for (self.getSectionHeaders()) |*sect_hdr| if (sect_hdr.getName() == null) return true;
        return false;
    }

    pub fn getSectionHeaders(self: *const Coff) []align(1) const SectionHeader {
        const coff_header = self.getHeader();
        const offset = self.coff_header_offset + @sizeOf(Header) + coff_header.size_of_optional_header;
        return @as([*]align(1) const SectionHeader, @ptrCast(self.data.ptr + offset))[0..coff_header.number_of_sections];
    }

    pub fn getSectionHeadersAlloc(self: *const Coff, allocator: mem.Allocator) ![]SectionHeader {
        const section_headers = self.getSectionHeaders();
        const out_buff = try allocator.alloc(SectionHeader, section_headers.len);
        for (out_buff, 0..) |*section_header, i| {
            section_header.* = section_headers[i];
        }

        return out_buff;
    }

    pub fn getSectionName(self: *const Coff, sect_hdr: *align(1) const SectionHeader) error{InvalidStrtabSize}![]const u8 {
        const name = sect_hdr.getName() orelse blk: {
            const strtab = (try self.getStrtab()).?;
            const name_offset = sect_hdr.getNameOffset().?;
            break :blk strtab.get(name_offset);
        };
        return name;
    }

    pub fn getSectionByName(self: *const Coff, comptime name: []const u8) ?*align(1) const SectionHeader {
        for (self.getSectionHeaders()) |*sect| {
            const section_name = self.getSectionName(sect) catch |e| switch (e) {
                error.InvalidStrtabSize => continue, //ignore invalid(?) strtab entries - see also GitHub issue #15238
            };
            if (mem.eql(u8, section_name, name)) {
                return sect;
            }
        }
        return null;
    }

    pub fn getSectionData(self: *const Coff, sec: *align(1) const SectionHeader) []const u8 {
        const offset = if (self.is_loaded) sec.virtual_address else sec.pointer_to_raw_data;
        return self.data[offset..][0..sec.virtual_size];
    }

    pub fn getSectionDataAlloc(self: *const Coff, sec: *align(1) const SectionHeader, allocator: mem.Allocator) ![]u8 {
        const section_data = self.getSectionData(sec);
        return allocator.dupe(u8, section_data);
    }
}