/*
 * elf.c
 *
 * Copyright (C) 2016 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <exec/elf.h>
#include <process.h>
#include <module.h>
#include <log.h>

extern void *get_kernel_symbol(const char *name);

static dword_t elf_hash(const byte_t *name)
{
    dword_t hash = 0;

    while (*name)
    {
        hash = (hash << 4) + *name++;
        hash = (hash ^ ((hash >> 24) & 0xF0)) & ~(hash & 0xF0000000);
    }

    return hash;
}

dword_t load_elf(handle_t file, thread_state_t *initial_state)
{
    elf_ident_t ident;
    char filename[MAX_PATH];
    qword_t filesize;
    size_t bytes_read;

    sysret_t ret = syscall(SYSCALL_QUERY_FILE, file, FILE_INFO_NAME, filename, sizeof(filename));
    if (ret != ERR_SUCCESS) return ret;

    ret = syscall(SYSCALL_QUERY_FILE, file, FILE_INFO_SIZE, &filesize, sizeof(filesize));
    if (ret != ERR_SUCCESS) return ret;

    ret = syscall(SYSCALL_READ_FILE, file, &ident, 0ULL, (size_t)sizeof(ident), &bytes_read);
    if (ret != ERR_SUCCESS) return ret;

    if (ident.magic != ELF_MAGIC) return ERR_INVALID;
    if (ident.class != ELF_CLASS_32BIT) return ERR_INVALID;

    elf32_header_t header;
    ret = syscall(SYSCALL_READ_FILE, file, &header, 16ULL, (size_t)sizeof(header), &bytes_read);
    if (ret != ERR_SUCCESS) return ret;

    handle_t section;
    ret = syscall(SYSCALL_CREATE_MEMORY_SECTION, filename, file, (size_t)filesize, MEMORY_SECTION_WRITABLE, &section);
    if (ret != ERR_SUCCESS) return ret;

    dword_t i;
    elf32_segment_t *segments = __builtin_alloca(header.segment_count * sizeof(elf32_segment_t));

    ret = syscall(SYSCALL_READ_FILE,
                  file,
                  segments,
                  (qword_t)header.segment_offset,
                  (size_t)(header.segment_count * sizeof(elf32_segment_t)),
                  &bytes_read);
    if (ret != ERR_SUCCESS) return ret;

    for (i = 0; i < header.segment_count; i++)
    {
        if (segments[i].type ==  ELF_SEGMENT_NULL)
        {
            continue;
        }
        else if (ELF_SEGMENT_LOAD)
        {
            dword_t flags = MEMORY_BLOCK_USERMODE | MEMORY_BLOCK_EVICTABLE;
            if (segments[i].flags & ELF_SEGMENT_READ) flags |= MEMORY_BLOCK_ACCESSIBLE;
            if (segments[i].flags & ELF_SEGMENT_WRITE) flags |= MEMORY_BLOCK_ACCESSIBLE | MEMORY_BLOCK_WRITABLE;
            if (segments[i].flags & ELF_SEGMENT_EXEC) flags |= MEMORY_BLOCK_EXECUTABLE;

            qword_t offset = (qword_t)segments[i].offset - PAGE_OFFSET(segments[i].virtual_address);

            ret = syscall(SYSCALL_MAP_MEMORY_SECTION,
                          INVALID_HANDLE,
                          section,
                          &segments[i].virtual_address,
                          offset,
                          PAGE_ALIGN_UP(segments[i].mem_size),
                          flags);
            if (ret != ERR_SUCCESS) break;
        }
        else
        {
            ret = ERR_INVALID;
            break;
        }
    }

    initial_state->regs.eip = header.entry;

    void *stack_address = NULL;
    ret = syscall(SYSCALL_ALLOC_MEMORY,
                  INVALID_HANDLE,
                  &stack_address,
                  0x200000,
                  MEMORY_BLOCK_ACCESSIBLE | MEMORY_BLOCK_WRITABLE | MEMORY_BLOCK_USERMODE | MEMORY_BLOCK_EVICTABLE);
    if (ret != ERR_SUCCESS) return ret;

    initial_state->regs.esp = (uintptr_t)stack_address + 0x200000;
    syscall(SYSCALL_CLOSE_OBJECT, section);
    return ret;
}

dword_t load_module_elf(void *address, module_t *module)
{
    dword_t ret = ERR_SUCCESS;
    elf_ident_t *ident = (elf_ident_t*)address;
    if (ident->magic != ELF_MAGIC) return ERR_INVALID;
    if (ident->class != ELF_CLASS_32BIT) return ERR_INVALID;

    elf32_header_t *header = (elf32_header_t*)((uintptr_t)address + 16);
    elf32_segment_t *segments = (elf32_segment_t*)((uintptr_t)address + header->segment_offset);
    uintptr_t lowest_address = (uintptr_t)-1;
    uintptr_t highest_address = 0;
    size_t total_size = 0;
    const dword_t *hash = NULL;
    const elf32_symbol_t *symtab = NULL;
    const char *strtab = NULL;
    const elf32_rel_t *rel = NULL;
    size_t rel_size = 0, rel_entsize = 0;
    const elf32_rela_t *rela = NULL;
    size_t rela_size = 0, rela_entsize = 0;
    const void *plt_rel = NULL;
    dword_t plt_rel_type = 0;
    size_t plt_rel_size = 0;

    inline void *lookup_symbol(const char *name)
    {
        const dword_t *bucket = &hash[2];
        const dword_t *chain = &hash[2 + hash[0]];
        dword_t index = bucket[elf_hash((const byte_t*)name) % hash[0]];

        while (index != ELF_SYMBOL_UNDEF)
        {
            if (strcmp(name, &strtab[symtab[index].st_name]) == 0)
            {
                return (void*)((uintptr_t)module->base_address - lowest_address + symtab[index].st_value);
            }

            index = chain[index];
        }

        return NULL;
    }

    inline void relocate_with_addend(const elf32_rela_t *rela)
    {
        void *location = (void*)(module->base_address - lowest_address + rela->r_offset);
        dword_t index = ELF_R_SYM(rela->r_info);
        uintptr_t symbol = 0;

        if (index != ELF_SYMBOL_UNDEF)
        {
            symbol = symtab[index].st_value
                ? (uintptr_t)module->base_address - lowest_address + symtab[index].st_value
                : (uintptr_t)get_kernel_symbol(&strtab[symtab[index].st_name]);
        }

        switch (ELF_R_TYPE(rela->r_info))
        {
        case ELF_R_386_32: *(dword_t*)location += symbol + rela->r_addend; break;
        case ELF_R_386_PC32: *(dword_t*)location += symbol + rela->r_addend - (uintptr_t)location; break;

        case ELF_R_386_GOT32:
            log_write(LOG_WARNING, "missing support: GOT32\n");
            /* *(dword_t*)location = G + rela->r_addend */;
            break;

        case ELF_R_386_PLT32:
            log_write(LOG_WARNING, "missing support: PLT32\n");
            /* *(dword_t*)location = L + rela->r_addend - (uintptr_t)location; */ break;

        case ELF_R_386_RELATIVE: *(dword_t*)location += (uintptr_t)module->base_address + rela->r_addend; break;

        case ELF_R_386_GOTOFF:
            log_write(LOG_WARNING, "missing support: GOTOFF\n");
            /* *(dword_t*)location = symbol + rela->r_addend - GOT; */ break;

        case ELF_R_386_GOTPC:
            log_write(LOG_WARNING, "missing support: GOTPC\n");
            /* *(dword_t*)location = GOT + rela->r_addend - (uintptr_t)location; */
            break;

        case ELF_R_386_32PLT:
            log_write(LOG_WARNING, "missing support: 32PLT\n");
            /* *(dword_t*)location = L + rela->r_addend */;
            break;

        case ELF_R_386_16: *(word_t*)location += symbol + rela->r_addend; break;
        case ELF_R_386_PC16: *(word_t*)location += symbol + rela->r_addend - (uintptr_t)location; break;
        case ELF_R_386_8: *(byte_t*)location += symbol + rela->r_addend; break;
        case ELF_R_386_PC8: *(byte_t*)location += symbol + rela->r_addend - (uintptr_t)location; break;

        case ELF_R_386_GLOB_DAT:
        case ELF_R_386_JMP_SLOT:
            *(dword_t*)location = symbol; break;
        }
    }

    inline void relocate(const elf32_rel_t *rel)
    {
        elf32_rela_t rela;
        rela.r_offset = rel->r_offset;
        rela.r_info = rel->r_info;
        rela.r_addend = 0;
        relocate_with_addend(&rela);
    }

    int i;
    for (i = 0; i < header->segment_count; i++)
    {
        if (segments[i].type == ELF_SEGMENT_LOAD)
        {
            uintptr_t segment_highest = PAGE_ALIGN_UP(segments[i].virtual_address + segments[i].mem_size);
            if (segments[i].virtual_address < lowest_address) lowest_address = segments[i].virtual_address;
            if (segment_highest > highest_address) highest_address = segment_highest;
        }
    }

    if (highest_address <= lowest_address) return ERR_INVALID;
    total_size = highest_address - lowest_address;

    while (TRUE)
    {
        /* Temporary HACK until memory region reservation is implemented */
        module->base_address = alloc_pool(NULL, PAGE_ALIGN_UP(total_size), 0);
        if (ret != ERR_SUCCESS) return ret;

        bool_t success = TRUE;
        free_pool(module->base_address);

        for (i = 0; i < header->segment_count; i++) if (segments[i].type == ELF_SEGMENT_LOAD)
        {
            dword_t flags = 0;
            if (segments[i].flags & ELF_SEGMENT_READ) flags |= MEMORY_BLOCK_ACCESSIBLE;
            if (segments[i].flags & ELF_SEGMENT_WRITE) flags |= MEMORY_BLOCK_ACCESSIBLE | MEMORY_BLOCK_WRITABLE;
            if (segments[i].flags & ELF_SEGMENT_EXEC) flags |= MEMORY_BLOCK_ACCESSIBLE | MEMORY_BLOCK_EXECUTABLE;

            module_segment_t *mod_seg = (module_segment_t*)malloc(sizeof(module_segment_t));
            mod_seg->address = (void*)((uintptr_t)module->base_address + segments[i].virtual_address - lowest_address);
            mod_seg->size = segments[i].mem_size;

            if (!alloc_pool(mod_seg->address, mod_seg->size, flags | MEMORY_BLOCK_WRITABLE))
            {
                success = FALSE;
                free(mod_seg);
                break;
            }

            memcpy(mod_seg->address, (void*)((uintptr_t)address + segments[i].offset), segments[i].file_size);
            list_append(&module->segments, &mod_seg->link);
        }

        if (!success)
        {
            module->base_address = NULL;

            while (module->segments.next != &module->segments)
            {
                module_segment_t *mod_seg = CONTAINER_OF(module->segments.next, module_segment_t, link);
                free_pool(mod_seg->address);
                list_remove(&mod_seg->link);
                free(mod_seg);
            }
        }
        else
        {
            break;
        }
    }

    elf32_dynamic_t *dynamic;
    for (i = 0; i < header->segment_count; i++) if (segments[i].type == ELF_SEGMENT_DYNAMIC)
    {
        for (dynamic = (elf32_dynamic_t*)((uintptr_t)address + segments[i].offset);
             dynamic->d_tag != ELF_DYNAMIC_TAG_NULL;
             dynamic++)
        {
            switch (dynamic->d_tag)
            {
            case ELF_DYNAMIC_TAG_HASH:
                hash = (const dword_t*)((uintptr_t)module->base_address - lowest_address + dynamic->d_ptr);
                break;
            case ELF_DYNAMIC_TAG_SYMTAB:
                symtab = (const elf32_symbol_t*)((uintptr_t)module->base_address - lowest_address + dynamic->d_ptr);
                break;
            case ELF_DYNAMIC_TAG_STRTAB:
                strtab = (const char*)((uintptr_t)module->base_address - lowest_address + dynamic->d_ptr);
                break;
            case ELF_DYNAMIC_TAG_REL:
                rel = (const elf32_rel_t*)((uintptr_t)module->base_address - lowest_address + dynamic->d_ptr);
                break;
            case ELF_DYNAMIC_TAG_RELSZ:
                rel_size = dynamic->d_val;
                break;
            case ELF_DYNAMIC_TAG_RELENT:
                rel_entsize = dynamic->d_val;
                break;
            case ELF_DYNAMIC_TAG_RELA:
                rela = (const elf32_rela_t*)((uintptr_t)module->base_address - lowest_address + dynamic->d_ptr);
                break;
            case ELF_DYNAMIC_TAG_RELASZ:
                rela_size = dynamic->d_val;
                break;
            case ELF_DYNAMIC_TAG_RELAENT:
                rela_entsize = dynamic->d_val;
                break;
            case ELF_DYNAMIC_TAG_JMPREL:
                plt_rel = (const void*)((uintptr_t)module->base_address - lowest_address + dynamic->d_ptr);
                break;
            case ELF_DYNAMIC_TAG_PLTRELSZ:
                plt_rel_size = dynamic->d_val;
                break;
            case ELF_DYNAMIC_TAG_PLTREL:
                plt_rel_type = dynamic->d_val;
                break;
            }
        }
    }

    if (!hash || !symtab || !strtab)
    {
        while (module->segments.next != &module->segments)
        {
            module_segment_t *mod_seg = CONTAINER_OF(module->segments.next, module_segment_t, link);
            free_pool(mod_seg->address);
            list_remove(&mod_seg->link);
            free(mod_seg);
        }

        return ERR_INVALID;
    }

    if (rel && rel_entsize == sizeof(elf32_rel_t))
    {
        for (i = 0; i < rel_size / rel_entsize; i++) relocate(&rel[i]);
    }

    if (rela && rela_entsize == sizeof(elf32_rela_t))
    {
        for (i = 0; i < rela_size / rela_entsize; i++) relocate_with_addend(&rela[i + 1]);
    }

    if (plt_rel)
    {
        if (plt_rel_type == ELF_DYNAMIC_TAG_REL)
        {
            for (i = 0; i < plt_rel_size / sizeof(elf32_rel_t); i++) relocate(&((elf32_rel_t*)plt_rel)[i]);
        }
        else if (plt_rel_type == ELF_DYNAMIC_TAG_RELA)
        {
            for (i = 0; i < plt_rel_size / sizeof(elf32_rela_t); i++) relocate_with_addend(&((elf32_rela_t*)plt_rel)[i]);
        }
    }

    if (ret == ERR_SUCCESS)
    {
        const char *driver_name = lookup_symbol("driver_name");
        if (driver_name != NULL) module->name = strdup(driver_name);
        module->load_proc = lookup_symbol("driver_load");
        module->unload_proc = lookup_symbol("driver_unload");
    }

    return ret;
}
