ElfLoader.cpp

#include "ElfLoader.hpp"

#include <fstream>
#include <vector>
#include <stdexcept>
#include <sys/mman.h>
#include <cstring>

namespace riscv {

ElfLoader::ElfLoader(Memory& memory) : memory_(memory) {}

bool ElfLoader::validate_elf_header(const Elf64_Ehdr& header) const {
    if (header.e_ident[EI_MAG0] != ELFMAG0 ||
        header.e_ident[EI_MAG1] != ELFMAG1 ||
        header.e_ident[EI_MAG2] != ELFMAG2 ||
        header.e_ident[EI_MAG3] != ELFMAG3) {
        return false; // Invalid magic number
    }

    if (header.e_ident[EI_CLASS] != ELFCLASS64) {
        return false; // Not a 64-bit ELF
    }

    if (header.e_ident[EI_DATA] != ELFDATA2LSB) {
        return false; // Not little-endian
    }

    if (header.e_type != ET_EXEC) {
        return false; // Not an executable
    }

    if (header.e_machine != EM_RISCV) {
        return false; // Not RISC-V architecture
    }

    return true;
}

std::optional<uint64_t> ElfLoader::load_elf(const std::string& filename) {
    std::ifstream file(filename, std::ios::binary);
    if (!file.is_open()) {
        return std::nullopt; // Failed to open file
    }

    Elf64_Ehdr header;
    file.read(reinterpret_cast<char*>(&header), sizeof(header));

    if (!validate_elf_header(header)) {
        return std::nullopt; // Invalid ELF header
    }

    // Read program headers
    file.seekg(header.e_phoff);
    std::vector<Elf64_Phdr> program_headers(header.e_phnum);
    file.read(reinterpret_cast<char*>(program_headers.data()), header.e_phnum * sizeof(Elf64_Phdr));

    for (const auto& phdr : program_headers) {
        if (phdr.p_type == PT_LOAD) {
            // Calculate the guest memory address
            uint8_t* guest_addr = memory_.get_data() + phdr.p_vaddr;

            // Set memory permissions to read/write
            if (mprotect(guest_addr, phdr.p_memsz, PROT_READ | PROT_WRITE) == -1) {
                throw std::runtime_error("Failed to mprotect memory for loading");
            }

            // Load segment data
            file.seekg(phdr.p_offset);
            file.read(reinterpret_cast<char*>(guest_addr), phdr.p_filesz);

            // Zero out the remaining memory if memsz > filesz
            if (phdr.p_memsz > phdr.p_filesz) {
                std::memset(guest_addr + phdr.p_filesz, 0, phdr.p_memsz - phdr.p_filesz);
            }

            // Set final memory permissions
            int prot = PROT_READ;
            if (phdr.p_flags & 0x1) { // PF_X (executable)
                prot |= PROT_EXEC;
            }
            if (phdr.p_flags & 0x2) { // PF_W (writable)
                prot |= PROT_WRITE;
            }
            if (mprotect(guest_addr, phdr.p_memsz, prot) == -1) {
                throw std::runtime_error("Failed to mprotect memory after loading");
            }
        }
    }

    return header.e_entry;
}

} // namespace riscv