/*
 * aout.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/aout.h>
#include <memory.h>
#include <heap.h>
#include <syscalls.h>

dword_t load_aout(handle_t file, thread_state_t *initial_state)
{
    aout_header_t header;
    size_t bytes_read;
    char filename[MAX_PATH];
    qword_t filesize;
    uintptr_t i;

    dword_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, &header, 0ULL, sizeof(aout_header_t), &bytes_read);
    if (ret != ERR_SUCCESS) return ret;

    void *address = (void*)0x7FE00000;
    void *data_address;
    void *bss_address;

    ret = syscall(SYSCALL_ALLOC_MEMORY,
                  INVALID_HANDLE,
                  &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)address + 0x200000;

    address = (void*)0x1000;

    switch (header.midmag & 0xFFFF)
    {
    case OMAGIC:
    {
        data_address = (void*)((uintptr_t)address + header.text_size + sizeof(aout_header_t));
        bss_address = (void*)((uintptr_t)address + header.text_size + header.data_size + sizeof(aout_header_t));

        ret = syscall(SYSCALL_ALLOC_MEMORY,
                      INVALID_HANDLE,
                      &address,
                      sizeof(aout_header_t) + header.text_size + header.data_size + header.bss_size,
                      MEMORY_BLOCK_ACCESSIBLE | MEMORY_BLOCK_WRITABLE | MEMORY_BLOCK_EXECUTABLE | MEMORY_BLOCK_USERMODE | MEMORY_BLOCK_EVICTABLE);
        if (ret != ERR_SUCCESS) return ret;

        ret = syscall(SYSCALL_READ_FILE,
                      file,
                      address,
                      0ULL,
                      sizeof(aout_header_t) + header.text_size + header.data_size,
                      &bytes_read);
        if (ret != ERR_SUCCESS) syscall(SYSCALL_FREE_MEMORY, INVALID_HANDLE, address);

        initial_state->regs.eip = (uintptr_t)address + header.entry_point + sizeof(aout_header_t);
        break;
    }

    case NMAGIC:
        return ERR_NOSYSCALL; // TODO

    case ZMAGIC:
    {
        handle_t section;
        data_address = (void*)((uintptr_t)address + PAGE_ALIGN_UP(header.text_size));
        bss_address = (void*)((uintptr_t)address + PAGE_ALIGN_UP(header.text_size) + PAGE_ALIGN_UP(header.data_size));

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

        ret = syscall(SYSCALL_MAP_MEMORY_SECTION,
                      INVALID_HANDLE,
                      section,
                      &address,
                      (qword_t)sizeof(aout_header_t),
                      PAGE_ALIGN_UP(header.text_size),
                      MEMORY_BLOCK_ACCESSIBLE
                      | MEMORY_BLOCK_WRITABLE
                      | MEMORY_BLOCK_USERMODE
                      | MEMORY_BLOCK_EXECUTABLE
                      | MEMORY_BLOCK_EVICTABLE);
        if (ret != ERR_SUCCESS)
        {
            syscall(SYSCALL_CLOSE_OBJECT, section);
            return ret;
        }

        if (header.data_size > 0)
        {
            ret = syscall(SYSCALL_MAP_MEMORY_SECTION,
                          INVALID_HANDLE,
                          section,
                          &data_address,
                          (qword_t)(sizeof(aout_header_t) + header.text_size),
                          PAGE_ALIGN_UP(header.data_size),
                          MEMORY_BLOCK_ACCESSIBLE | MEMORY_BLOCK_WRITABLE | MEMORY_BLOCK_USERMODE | MEMORY_BLOCK_EVICTABLE);
            if (ret != ERR_SUCCESS)
            {
                syscall(SYSCALL_FREE_MEMORY, address);
                syscall(SYSCALL_CLOSE_OBJECT, section);
                return ret;
            }
        }

        if (header.bss_size > 0)
        {
            ret = syscall(SYSCALL_ALLOC_MEMORY,
                          INVALID_HANDLE,
                          &bss_address,
                          header.bss_size,
                          MEMORY_BLOCK_ACCESSIBLE | MEMORY_BLOCK_WRITABLE | MEMORY_BLOCK_USERMODE | MEMORY_BLOCK_EVICTABLE);
            if (ret != ERR_SUCCESS)
            {
                syscall(SYSCALL_FREE_MEMORY, address);
                if (header.data_size > 0) syscall(SYSCALL_FREE_MEMORY, data_address);
                syscall(SYSCALL_CLOSE_OBJECT, section);
                return ret;
            }
        }

        syscall(SYSCALL_CLOSE_OBJECT, section);
        initial_state->regs.eip = (uintptr_t)address + header.entry_point;
        break;
    }

    default:
        return ERR_INVALID;
    }

    uintptr_t num_text_relocs = header.text_reloc_size / sizeof(aout_reloc_t);
    uintptr_t num_data_relocs = header.data_reloc_size / sizeof(aout_reloc_t);
    uintptr_t num_symbols = header.sym_size / sizeof(aout_symbol_t);
    aout_reloc_t *text_relocs = NULL;
    aout_reloc_t *data_relocs = NULL;
    aout_symbol_t *symbols = NULL;
    uintptr_t string_table_size = 0;
    char *string_table = NULL;

    if (header.text_reloc_size)
    {
        text_relocs = (aout_reloc_t*)malloc(header.text_reloc_size);
        if (text_relocs == NULL) return ERR_NOMEMORY;

        ret = syscall(SYSCALL_READ_FILE,
                      file,
                      text_relocs,
                      (qword_t)(sizeof(aout_header_t) + header.text_size + header.data_size),
                      header.text_reloc_size,
                      &bytes_read);
        if (ret != ERR_SUCCESS) goto cleanup;
    }

    if (header.data_reloc_size)
    {
        data_relocs = (aout_reloc_t*)malloc(header.data_reloc_size);
        if (data_relocs == NULL) return ERR_NOMEMORY;

        ret = syscall(SYSCALL_READ_FILE,
                      file,
                      data_relocs,
                      (qword_t)(sizeof(aout_header_t)
                                + header.text_size
                                + header.data_size
                                + header.text_reloc_size),
                      header.data_reloc_size,
                      &bytes_read);
        if (ret != ERR_SUCCESS) goto cleanup;
    }

    if (header.sym_size)
    {
        symbols = (aout_symbol_t*)malloc(header.sym_size);
        if (symbols == NULL) return ERR_NOMEMORY;

        ret = syscall(SYSCALL_READ_FILE,
                      file,
                      symbols,
                      (qword_t)(sizeof(aout_header_t)
                                + header.text_size
                                + header.data_size
                                + header.text_reloc_size
                                + header.data_reloc_size),
                      header.sym_size,
                      &bytes_read);
        if (ret != ERR_SUCCESS) goto cleanup;
    }

    ret = syscall(SYSCALL_READ_FILE,
                  file,
                  &string_table_size,
                  (qword_t)(sizeof(aout_header_t)
                            + header.text_size
                            + header.data_size
                            + header.sym_size
                            + header.text_reloc_size
                            + header.data_reloc_size),
                  sizeof(string_table_size),
                  &bytes_read);
    if (ret != ERR_SUCCESS) goto cleanup;

    string_table = malloc(string_table_size);
    if (string_table == NULL)
    {
        ret = ERR_NOMEMORY;
        goto cleanup;
    }

    ret = syscall(SYSCALL_READ_FILE,
                  file,
                  string_table,
                  (qword_t)(sizeof(aout_header_t)
                            + header.text_size
                            + header.data_size
                            + header.sym_size
                            + header.text_reloc_size
                            + header.data_reloc_size),
                  string_table_size,
                  &bytes_read);
    if (ret != ERR_SUCCESS) goto cleanup;

    for (i = 0; i < num_symbols; i++)
    {
        symbols[i].name = (char*)((uintptr_t)string_table + symbols[i].name_offset);

        switch (AOUT_TYPE_FIELD(symbols[i].type))
        {
        case AOUT_SYM_UNDEFINED:
            // TODO: Resolve symbol
            break;

        case AOUT_SYM_TEXT:
            symbols[i].value += (uintptr_t)address + sizeof(aout_header_t);
            break;

        case AOUT_SYM_DATA:
            symbols[i].value += (uintptr_t)data_address;
            break;

        case AOUT_SYM_BSS:
            symbols[i].value += (uintptr_t)bss_address;
            break;
        }
    }

    for (i = 0; i < num_text_relocs + num_data_relocs; i++)
    {
        aout_reloc_t *reloc = (i < num_text_relocs) ? &text_relocs[i] : &data_relocs[i];
        uintptr_t base_address = (i < num_text_relocs) ? ((uintptr_t)address + sizeof(aout_header_t)) : (uintptr_t)data_address;
        dword_t offset;

        if (!reloc->external)
        {
            offset = (uintptr_t)address + sizeof(aout_header_t);
        }
        else
        {
            if (reloc->symbol_num >= num_symbols)
            {
                ret = ERR_INVALID;
                goto cleanup;
            }

            offset = symbols[reloc->symbol_num].value;
        }

        switch (reloc->length)
        {
        case 0:
            *(byte_t*)(base_address + reloc->address) += offset;
            break;
        case 1:
            *(word_t*)(base_address + reloc->address) += offset;
            break;
        case 2:
            *(dword_t*)(base_address + reloc->address) += offset;
            break;
        }
    }

    if ((header.midmag & 0xFFFF) == ZMAGIC)
    {
        syscall(SYSCALL_SET_MEMORY_FLAGS,
                INVALID_HANDLE,
                address,
                MEMORY_BLOCK_ACCESSIBLE | MEMORY_BLOCK_USERMODE | MEMORY_BLOCK_EXECUTABLE | MEMORY_BLOCK_EVICTABLE);
    }

cleanup:
    if (symbols) free(symbols);
    if (string_table) free(string_table);
    return ret;
}
