#include <magenta/misc.h>
#include <magenta/errors.h>
#include <magenta/handle.h>
#include <magenta/vmo.h>
#include <magenta/vmar.h>
#include <stdio.h>
#include <string.h>
#include <stdarg.h>
#include "__elf.h"

#define HIDDEN __attribute__((visibility("hidden")))

#define VDSO_SONAME "mx.dylib"

#define NEEDS_NOTHING 0
#define NEEDS_VDSO 1
#define NEEDS_MORE 2

#define ACL (PF_R | PF_W | PF_X)
#define ACCESS(x) ((x) & ACL)

/* TODO in case we ever support ELF32 images */
#define elf_class_bits(x) (64)

#define PAGE_OFFSET(v) ((v) & (MX_PAGE_SIZE - 1))
#define PAGE_ALIGN_DOWN(v) (v) &= ~(MX_PAGE_SIZE - 1)
#define PAGE_ALIGN_UP(v)                \
    do {                                \
        if ((v) & (MX_PAGE_SIZE - 1)) { \
            v &= ~(MX_PAGE_SIZE - 1);   \
            v += MX_PAGE_SIZE;          \
        }                               \
    } while (0)

#undef DEBUG_LOG

HIDDEN void __elf_image_init(mx_handle_t local_vmar, mx_handle_t remote_vmar, struct elf_image *out)
{
    memset(out, 0x0, sizeof(*out));
    out->local_root_vmar = local_vmar;
    out->remote_root_vmar = remote_vmar;
}

static int elf_validate_ehdr(elf_ehdr_t *hdr)
{
    if (hdr->e_ident[EI_MAG0] != ELF_MAG0) {
        return -1;
    }

    if (hdr->e_ident[EI_MAG1] != ELF_MAG1) {
        return -1;
    }

    if (hdr->e_ident[EI_MAG2] != ELF_MAG2) {
        return -1;
    }

    if (hdr->e_ident[EI_MAG3] != ELF_MAG3) {
        return -1;
    }

    if (hdr->e_ident[EI_CLASS] != ELFCLASS64) {
        return -1;
    }

    if (hdr->e_machine != EM_X86_64) {
        return -1;
    }

    if (hdr->e_ident[EI_DATA] != ELFDATA2LSB) {
        return -1;
    }

    if (hdr->e_ident[EI_VERSION] != EV_CURRENT) {
        return -1;
    }

    return 0;
}

static int read_header(struct elf_image *image)
{
    mx_status_t err = mx_vmo_read(image->image_vmo, (uint8_t *)&image->hdr, sizeof(image->hdr), 0);
    if (err != MX_OK) {
        printf("vmo_read %x error %s\n", image->image_vmo, mx_status_to_string(err));
        return -1;
    }

    return elf_validate_ehdr(&image->hdr);
}

static mx_status_t allocate_vmar(struct elf_image *image, size_t size)
{
    mx_status_t err = mx_vmar_allocate(image->local_root_vmar,
        MX_VM_CAN_MAP_READ | MX_VM_CAN_MAP_WRITE |
        MX_VM_CAN_MAP_EXEC | MX_VM_CAN_MAP_SPECIFIC,
        0, size, &image->local_exec_vmar, &image->local_base_addr);

    if (err != MX_OK) {
        return err;
    }

    err = mx_vmar_allocate(image->remote_root_vmar,
        MX_VM_CAN_MAP_READ | MX_VM_CAN_MAP_WRITE |
        MX_VM_CAN_MAP_EXEC | MX_VM_CAN_MAP_SPECIFIC,
        0, size, &image->remote_exec_vmar, &image->remote_base_addr);
    image->exec_vmar_size = size;

    return err;
}

static void deduce_dynsym_count(struct elf_image *image)
{
    if (image->hash_type == REG_HASH) {
        image->dynsym_len = *(((uint32_t *)image->hash_offset) + 1);
    } else {
        /* We don't need to know the symbol count to use DT_GNU_HASH
         * (which is good because calculating it is a pain) */
        image->dynsym_len = 0;
    }
 }

static void read_dynamic_segment(struct elf_image *image)
{
    elf_dyn_t *dyn_array = (elf_dyn_t *)(image->local_base_addr + image->dynamic.p_vaddr);

    for (int i = 0;; i++) {
        elf_dyn_t *dyn = &dyn_array[i];

        if (dyn->d_tag == DT_NULL) {
            break;
        }

        switch (dyn->d_tag) {
        case DT_STRTAB:
            image->dynstr_offset = dyn->d_un.d_ptr;
            dyn->d_un.d_ptr += image->remote_base_addr;
            break;
        case DT_STRSZ:
            image->dynstr_len = dyn->d_un.d_val;
            break;
        case DT_SYMTAB:
            image->dynsym_offset = dyn->d_un.d_ptr;
            dyn->d_un.d_ptr += image->remote_base_addr;
            break;
        case DT_SYMENT:
            image->dynsym_entsz = dyn->d_un.d_val;
            break;
        case DT_HASH:
            /* prefer the GNU hash table */
            if (image->hash_type == GNU_HASH) {
                continue;
            }
            image->hash_type = REG_HASH;
            image->hash_offset = dyn->d_un.d_ptr;
            dyn->d_un.d_ptr += image->remote_base_addr;
            break;
        case DT_GNU_HASH:
            image->hash_type = GNU_HASH;
            image->hash_offset = dyn->d_un.d_ptr;
            dyn->d_un.d_ptr += image->remote_base_addr;
            break;
        default:
            break;
        }
    }
}

static void get_memory_requirements(void *phdr_buf, int entsize, int nr_phdr,
        size_t *out_vmem, size_t *out_rw)
{
    size_t vmem_max = 0;
    size_t rw_max = 0;

    for (int i = 0; i < nr_phdr; i++) {
        elf_phdr_t *phdr = (elf_phdr_t *)((char *)phdr_buf + (i * entsize));
        if (phdr->p_type != PT_LOAD) {
            continue;
        }

        size_t foffset = phdr->p_offset;
        size_t voffset = phdr->p_vaddr;
        PAGE_ALIGN_DOWN(foffset);
        PAGE_ALIGN_DOWN(voffset);

        size_t vend = phdr->p_vaddr + phdr->p_memsz;
        PAGE_ALIGN_UP(vend);

        size_t fsize = phdr->p_filesz;
        size_t vsize = vend - voffset;
        PAGE_ALIGN_UP(fsize);

        if (phdr->p_flags & PF_W) {
            rw_max += vsize;
        }

        if (vend > vmem_max) {
            vmem_max = vend;
        }
    }

    *out_rw = rw_max;
    *out_vmem = vmem_max;
}

static mx_status_t map_memory(struct elf_image *image,
        mx_off_t voffset, mx_off_t foffset, size_t sz, int flags, void **out_buf)
{
    mx_status_t err = MX_OK;
    mx_flags_t options = MX_VM_SPECIFIC;
    mx_handle_t vmo = image->image_vmo;

    if (flags & PF_R) {
        options |= MX_VM_PERM_READ;
    }

    if (flags & PF_W) {
        options |= MX_VM_PERM_WRITE;
        vmo = image->rw_vmo;
    }

    if (flags & PF_X) {
        options |= MX_VM_PERM_EXEC;
    }

    mx_vaddr_t local_addr = 0, remote_addr = 0;
    *out_buf = NULL;

    err = mx_vmar_map(image->local_exec_vmar, options, voffset, vmo, foffset, sz, &local_addr);
    if (err != MX_OK) {
        return err;
    }

    err = mx_vmar_map(image->remote_exec_vmar, options, voffset, vmo, foffset, sz, &remote_addr);
    if (err != MX_OK) {
        return err;
    }

    *out_buf = (void *)local_addr;
    return MX_OK;
}

static mx_status_t map_image(struct elf_image *image,
        void *phdr_buf, int phdr_entsize, int phdr_num)
{
    size_t rw_offset = 0;
    mx_status_t err = MX_OK;

    for (int i = 0; i < phdr_num; i++) {
        elf_phdr_t *phdr = (elf_phdr_t *)((char *)phdr_buf + (i * phdr_entsize));
        if (phdr->p_type == PT_DYNAMIC) {
            image->dynamic = *phdr;
        }

        if (phdr->p_type != PT_LOAD) {
            continue;
        }

        size_t foffset = phdr->p_offset;
        size_t voffset = phdr->p_vaddr;
        PAGE_ALIGN_DOWN(foffset);
        PAGE_ALIGN_DOWN(voffset);

        size_t vend = phdr->p_vaddr + phdr->p_memsz;
        PAGE_ALIGN_UP(vend);

        size_t fsize = phdr->p_filesz;
        size_t vsize = vend - voffset;
        PAGE_ALIGN_UP(fsize);

        if (phdr->p_flags & PF_W) {
            size_t rw_vmo_offset = rw_offset;
            size_t rw_vmo_size = vsize;

            size_t bss_offset = phdr->p_vaddr + phdr->p_filesz;
            size_t bss_size = phdr->p_memsz - phdr->p_filesz;

            void *segment_buf = NULL;
            err = map_memory(image, voffset, rw_vmo_offset, vsize, phdr->p_flags, &segment_buf);

            if (err != MX_OK) {
                return err;
            }

            void *file_dest = (char *)image->local_base_addr + phdr->p_vaddr;
            void *bss_dest = (char *)image->local_base_addr + bss_offset;

            mx_vmo_read(image->image_vmo, file_dest, phdr->p_filesz, phdr->p_offset);

            if (bss_size) {
                memset(bss_dest, 0x0, bss_size);
            }

            rw_offset += rw_vmo_size;
        } else {
            void *segment_buf = NULL;
            err = map_memory(image, voffset, foffset, vsize, phdr->p_flags, &segment_buf);

            if (err != MX_OK) {
                return err;
            }
        }
    }

    return MX_OK;
}

HIDDEN int __elf_image_load_image(struct elf_image *image, mx_handle_t image_vmo)
{
    image->image_vmo = image_vmo;

    int status = read_header(image);

    if (status != 0) {
        printf("read_header error\n");
        return status;
    }

    size_t phdr_bufsz = image->hdr.e_phentsize * image->hdr.e_phnum;
    unsigned char phdr_buf[phdr_bufsz];

    mx_vmo_read(image_vmo, phdr_buf, phdr_bufsz, image->hdr.e_phoff);

    size_t rw_size, vmem_size;
    get_memory_requirements(phdr_buf, image->hdr.e_phentsize, image->hdr.e_phnum, &vmem_size, &rw_size);

    if (allocate_vmar(image, vmem_size) != MX_OK) {
        printf("allocate_vmar error\n");
        return -1;
    }

    if (rw_size) {
        mx_vmo_create(rw_size, MX_VM_PERM_READ | MX_VM_PERM_WRITE, &image->rw_vmo);
    }

    mx_status_t err = map_image(image, phdr_buf, image->hdr.e_phentsize, image->hdr.e_phnum);
    if (err != MX_OK) {
        printf("map_image error\n");
        return -1;
    }

    read_dynamic_segment(image);

    return 0;
}

HIDDEN int __elf_get_interp_path(mx_handle_t image_vmo, char *out, size_t max)
{
    elf_ehdr_t hdr;
    mx_vmo_read(image_vmo, (uint8_t *)&hdr, sizeof(hdr), 0);

    if (elf_validate_ehdr(&hdr) != 0) {
        return -1;
    }

    unsigned char phdr_buf[hdr.e_phentsize * hdr.e_phnum];
    mx_vmo_read(image_vmo, phdr_buf, hdr.e_phentsize * hdr.e_phnum, hdr.e_phoff);

    uintptr_t dynamic_offset = 0, interp_offset = 0;
    size_t dynamic_sz = 0, interp_sz = 0;

    for (size_t i = 0; i < hdr.e_phnum; i++) {
        elf_phdr_t *phdr = (elf_phdr_t *)(phdr_buf + i * hdr.e_phentsize);

        switch (phdr->p_type) {
        case PT_DYNAMIC:
            dynamic_offset = phdr->p_offset;
            dynamic_sz = phdr->p_filesz;
        break;
        case PT_INTERP:
            interp_offset = phdr->p_offset;
            interp_sz = phdr->p_filesz;
            break;
        default:
            break;
        }
    }

    if (!dynamic_sz || !interp_offset) {
        return 0;
    }

    if (max > interp_sz) {
        max = interp_sz;
    }

    mx_vmo_read(image_vmo, (uint8_t *)out, max, interp_offset);
    return 1;
}

HIDDEN int __elf_check_dependencies(struct elf_image *image)
{
    elf_dyn_t *dyn_array = (elf_dyn_t *)(image->local_base_addr + image->dynamic.p_vaddr);

    for (int i = 0;; i++) {
        elf_dyn_t *dyn = &dyn_array[i];

        if (dyn->d_tag == DT_NULL) {
            break;
        }

        if (dyn->d_tag != DT_NEEDED) {
            continue;
        }

        const char *lib_name = (const char *)(image->local_base_addr + image->dynstr_offset + dyn->d_un.d_ptr);
        if (strcmp(lib_name, VDSO_SONAME) != 0) {
            /* We can't load executables that link to libraries other than
             * libmx */
            return -1;
        }
    }

    return 0;
}

static elf_sym_t *get_dynsym_entry(struct elf_image *image, unsigned int idx)
{
    return (elf_sym_t *)(image->local_base_addr + image->dynsym_offset + idx * image->dynsym_entsz);
}

static const char *get_dynstr(struct elf_image *image, unsigned int idx)
{
    return (const char *)(image->local_base_addr + image->dynstr_offset + idx);
}

static int do_rela(struct elf_image *image, struct elf_image *lib, uintptr_t ptr, size_t sz, size_t entsz)
{
    size_t entries = sz / entsz;

    for (size_t i = 0; i < entries; i++) {
        elf_rela_t *rela = (elf_rela_t *)(image->local_base_addr + ptr + i * entsz);
        int sym_idx = ELF64_R_SYM(rela->r_info);
        int type = ELF64_R_TYPE(rela->r_info);

        mx_vaddr_t sym_val = 0;
        if (type != R_X86_64_RELATIVE) {
            elf_sym_t *dynsym = get_dynsym_entry(image, sym_idx);
            const char *name = get_dynstr(image, dynsym->st_name);
            sym_val = __elf_image_get_symbol(image, name);
            if (!sym_val && lib) {
                sym_val = __elf_image_get_symbol(lib, name);
            }

            if (!sym_val) {
                return -1;
            }
        }

        int ok = 1;
        mx_status_t status = MX_OK;

        switch (type) {
        case R_X86_64_GLOB_DAT: {
            elf_xword_t val = sym_val;
            elf_xword_t *dest = (elf_xword_t *)((char *)image->local_base_addr + rela->r_offset);
            *dest = val;
            break;
        } case R_X86_64_64: {
            elf_xword_t val = sym_val + rela->r_addend;
            elf_xword_t *dest = (elf_xword_t *)((char *)image->local_base_addr + rela->r_offset);
            *dest = val;
            break;
        } case R_X86_64_JUMP_SLOT: {
            elf_xword_t val = sym_val;
            elf_xword_t *dest = (elf_xword_t *)((char *)image->local_base_addr + rela->r_offset);
            *dest = val;
            break;
        } case R_X86_64_RELATIVE: {
            elf_xword_t val = image->remote_base_addr + rela->r_addend;
            elf_xword_t *dest = (elf_xword_t *)((char *)image->local_base_addr + rela->r_offset);
            *dest = val;
            break;
        } default:
            ok = 0;
            break;
        }

        if (!ok || status != MX_OK) {
            return -1;
        }
    }

    return 0;
}

static int relocate(struct elf_image *image, struct elf_image *lib)
{
    elf_dyn_t *dyn_array = (elf_dyn_t *)(image->local_base_addr + image->dynamic.p_vaddr);

    int result = 0;

    uintptr_t rel_addr = 0, rela_addr = 0, plt_addr = 0;
    size_t rel_sz = 0, rel_entsz = 0;
    size_t rela_sz = 0, rela_entsz = 0;
    size_t plt_sz = 0, plt_entsz = 0;
    int plt_enttype;
    for (int i = 0;; i++) {
        elf_dyn_t *dyn = &dyn_array[i];

        if (dyn->d_tag == DT_NULL) {
            break;
        }

        switch (dyn->d_tag) {
        case DT_REL:
            rel_addr = dyn->d_un.d_ptr;
            break;
        case DT_RELSZ:
            rel_sz = dyn->d_un.d_val;
            break;
        case DT_RELENT:
            rel_entsz = dyn->d_un.d_val;
            break;
        case DT_RELA:
            rela_addr = dyn->d_un.d_ptr;
            break;
        case DT_RELASZ:
            rela_sz = dyn->d_un.d_val;
            break;
        case DT_RELAENT:
            rela_entsz = dyn->d_un.d_val;
            break;
        case DT_PLTRELSZ:
            plt_sz = dyn->d_un.d_val;
            break;
        case DT_JMPREL:
            plt_addr = dyn->d_un.d_ptr;
            break;
        case DT_PLTREL:
            plt_enttype = dyn->d_un.d_ptr;
            break;
        default:
            break;
        }

        if (dyn->d_tag != DT_NEEDED) {
            continue;
        }
    }

    if (rel_sz) {
        /* DT_REL is not supported */
        return -1;
    }

    if (plt_enttype == DT_RELA) {
        plt_entsz = rela_entsz ? rela_entsz : sizeof(elf_rela_t);
    } else if (plt_enttype == DT_REL) {
        return -1;
    }

    int r;
    if (rela_sz) {
        r = do_rela(image, lib, rela_addr, rela_sz, rela_entsz);
        if (r != 0) {
            return r;
        }
    }

    if (plt_sz) {
        r = do_rela(image, lib, plt_addr, plt_sz, plt_entsz);
        if (r != 0) {
            return r;
        }
    }

    return 0;
}

HIDDEN int __elf_image_link(struct elf_image *exec, struct elf_image *vdso)
{
    int status = relocate(vdso, NULL);
    if (status != 0) {
        return status;
    }

    return relocate(exec, vdso);
}

HIDDEN mx_vaddr_t __elf_image_entry_point(struct elf_image *image)
{
    return image->remote_base_addr + image->hdr.e_entry;
}

static uint32_t gnu_hash(const char *name)
{
    uint32_t h = 5381;
    for (; *name; name++) {
        h = (h << 5) + h + *name;
    }

    return h;
}

static mx_vaddr_t sym_gnu_hash_search(struct elf_image *image, const char *name)
{
    uint32_t hash = gnu_hash(name);
    uint32_t *hashtab = (uint32_t *)(image->local_base_addr + image->hash_offset);

    uint32_t nbuckets = hashtab[0];
    uint32_t symoffset = hashtab[1];
    uint32_t bloom_size = hashtab[2];
    uint32_t bloom_shift = hashtab[3];
    uint64_t *bloom = (uint64_t *)&hashtab[4];
    uint32_t *buckets = (uint32_t *)&bloom[bloom_size];
    uint32_t *chain = &buckets[nbuckets];

    uint64_t bloom_word =
        bloom[(hash / elf_class_bits(image)) % bloom_size];
    uint64_t bloom_mask =
        ((uint64_t)1 << (hash % elf_class_bits(image))) |
        ((uint64_t)1 << ((hash >> bloom_shift) % elf_class_bits(image)));

    if ((bloom_word & bloom_mask) != bloom_mask) {
        return 0;
    }

    uint32_t symix = buckets[hash % nbuckets];
    if (symix < symoffset) {
        return 0;
    }

    const char *strtab = (const char *)(image->local_base_addr + image->dynstr_offset);
    elf_sym_t *symtab = (elf_sym_t *)(image->local_base_addr + image->dynsym_offset);

    while (1) {
        const char *symname = strtab + symtab[symix].st_name;
        uint32_t symhash = chain[symix - symoffset];

        if ((symhash | 1) == (hash | 1) && !strcmp(symname, name)) {
            return image->remote_base_addr +  symtab[symix].st_value;
        }

        if (symhash & 1) {
            break;
        }

        symix++;
    }

    return 0;
}

HIDDEN mx_vaddr_t __elf_image_get_symbol(struct elf_image *image, const char *name)
{
    switch (image->hash_type) {
    case GNU_HASH:
        return sym_gnu_hash_search(image, name);
    default:
        return 0;
    }
}

HIDDEN void __elf_image_cleanup(struct elf_image *image)
{
    if (image->local_exec_vmar) {
        mx_vmar_unmap(image->local_exec_vmar, 0, image->exec_vmar_size);
        mx_vmar_destroy(image->local_exec_vmar);
        mx_handle_close(image->local_exec_vmar);
        image->local_exec_vmar = MX_NULL_HANDLE;
    }

    if (image->remote_exec_vmar) {
        mx_handle_close(image->remote_exec_vmar);
        image->remote_exec_vmar = MX_NULL_HANDLE;
    }

    if (image->rw_vmo) {
        mx_handle_close(image->rw_vmo);
        image->rw_vmo = MX_NULL_HANDLE;
    }
}