/* load shared from memory, sort of (x) 2021 by herm1t.vx@gmail.com / https://twitter.com/vx_herm1t */
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <unistd.h>
#include <elf.h>
#include <assert.h>
#include <link.h>
#include <fcntl.h>

static int load(unsigned char *lib, int size)
{
    uint8_t *map;
    Elf64_Ehdr *ehdr = (Elf64_Ehdr*)lib;
    Elf64_Phdr *phdr = (Elf64_Phdr*)(lib + ehdr->e_phoff);
    Elf64_Addr lo = ~0, hi = 0, t;
    int i;
    struct elf {
        int relent, relpltsz, symcnt, relasz, gnuhash;
        char *dynstr;
        Elf64_Sym * dynsym;
        void *relplt, *rela;
        void *delta;
        void (*init)(void);
        uint32_t *hash;
    } lm, lc;
    bzero(&lm, sizeof(lm));
    bzero(&lc, sizeof(lc));

    /* get ELF size */
    for (i = 0; i < ehdr->e_phnum; i++)
        if (phdr[i].p_type == PT_LOAD) {
            if (phdr[i].p_vaddr < lo)
                lo = phdr[i].p_vaddr;
            t = phdr[i].p_vaddr + phdr[i].p_memsz;
            if (t > hi)
                hi = t;
        }

    /* allocate memory */
    map = mmap(NULL, hi - lo, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
    if (map == MAP_FAILED) {
        printf("failed to allocate memory\n");
        return 2;
    }

    /* copy segments and find DYNAMIC */
    Elf64_Dyn *dyn = NULL;
    int f2p(int f) {
        int p = PROT_READ;
        if (f & PF_X)
            p |= PROT_EXEC;
        if (f & PF_W)
            p |= PROT_WRITE;
        return p;
    }
    for (i = 0; i < ehdr->e_phnum; i++) {
        if (phdr[i].p_type == PT_LOAD) {
            if (phdr[i].p_offset == 0)
                lm.delta = map - lo;
            Elf64_Addr off = phdr[i].p_vaddr & 4095ULL;
            void *seg = phdr[i].p_vaddr - off + lm.delta;
            memcpy(seg + off, lib + phdr[i].p_offset, phdr[i].p_filesz);
            mprotect(seg, phdr[i].p_memsz + off, f2p(phdr[i].p_flags));
        }
        if (phdr[i].p_type == PT_DYNAMIC)
            dyn = (Elf64_Dyn*)(lm.delta + phdr[i].p_vaddr);
    }
    assert(dyn != NULL);

    /* get init(), relocations and symbols */
    void parse_dyn(Elf64_Dyn *dyn, struct elf *elf) {
#define GET_PTR(tag, var)   if (dyn->d_tag == tag) elf->var = (void*)(dyn->d_un.d_val + elf->delta);
#define GET_VAL(tag, var)   if (dyn->d_tag == tag) elf->var = dyn->d_un.d_val;
        for ( ; dyn->d_tag != DT_NULL; dyn++) {
            GET_PTR(DT_INIT, init)
            GET_PTR(DT_JMPREL, relplt)
            GET_VAL(DT_RELENT, relent)
            GET_VAL(DT_RELAENT, relent)
            GET_VAL(DT_PLTRELSZ, relpltsz)
            GET_PTR(DT_REL, rela)
            GET_PTR(DT_RELA, rela)
            GET_VAL(DT_RELSZ, relasz)
            GET_VAL(DT_RELASZ, relasz)
            GET_PTR(DT_SYMTAB, dynsym)
            GET_PTR(DT_STRTAB, dynstr)
            GET_PTR(DT_HASH, hash)
            GET_PTR(DT_GNU_HASH, hash)
            if (dyn->d_tag == DT_GNU_HASH)
                elf->gnuhash = 1;
        }
        /* get symbol count */
        uint32_t *hash = elf->hash;
        if (elf->gnuhash) {
            int i;
            uint32_t nbucket = hash[0];
            uint32_t base = hash[1];
            uint32_t bloom_size = hash[2];
            uint64_t *bloom = (void*)&hash[4];
            uint32_t *bucket = (void*)&bloom[bloom_size];
            uint32_t *chain = &bucket[nbucket];
                unsigned int last_sym = 0;
                for (i = 0; i < nbucket; i++)
                        if (bucket[i] > last_sym)
                                last_sym = bucket[i];
            for (i = last_sym - base; !(chain[i] & 1); i++)
                last_sym++;
            elf->symcnt= last_sym + 1;
        } else
        if (hash) {
            elf->symcnt = hash[1];
        } else {
            /* it's outdated, but... */
            elf->symcnt = (elf->dynstr - (char*)elf->dynsym) / sizeof(Elf64_Sym);
        }
    }
    parse_dyn(dyn, &lm);

    /* find libc */
        struct link_map *l = _r_debug.r_map;
    for ( ; l != NULL; l = l->l_next)
        if (l->l_name && *l->l_name && strstr(l->l_name, "libc")) {
            printf("Found %s %p\n", l->l_name, l->l_ld);
            parse_dyn(l->l_ld, &lc);
            lc.delta = (void*)l->l_addr;
            break;
    }
    assert(lc.dynsym != NULL);

    /* relocate, sort of */
    void relocate(void *rela, int relasz, int relaent) {
        void *r;
        for (r = rela; (r - rela) < relasz; r += relaent) {
            Elf64_Rela *rela = (Elf64_Rela*)r;
            int index = ELF64_R_SYM(rela->r_info);
            void *value = 0;
            char *name = NULL;
            if (index) {
                if (lm.dynsym[index].st_value) {
                    value = lm.dynsym[index].st_value + lm.delta;
                } else {
                    /* FIXME: use hash, search all libraries */
                    int found = 0;
                    name = lm.dynsym[index].st_name + lm.dynstr;
                    for (i = 0; i < lc.symcnt; i++)
                        if (! strcmp(name, lc.dynsym[i].st_name + lc.dynstr)) {
                            value = lc.dynsym[i].st_value + lc.delta;
                            found = 1;
                        }
                    if (! found)
                        printf("Failed to resolve %d %s\n", index, name);
                }
            }
            if (lm.relent == sizeof(Elf64_Rela))
                value += rela->r_addend;
            printf("%p (%s + %x) %p\n", rela->r_offset, name ? name : "",
                rela->r_addend, value - lm.delta);
            *(void**)(rela->r_offset + lm.delta) = value;
        }
    }
    relocate(lm.relplt, lm.relpltsz, lm.relent);
    relocate(lm.rela, lm.relasz, lm.relent);

    printf("Done! Call init\n");
    if (lm.init)
        lm.init();
}

int main(int argc, char **argv)
{
    int h = open(argv[1], 0);
    int l = lseek(h, 0, 2);
    lseek(h, 0, 0);
    unsigned char lib[l];
    if (h > 0 && l > 0 && read(h, lib, l) == l)
        load(lib, l);
    else
        printf("Failed to read %s\n", argv[1]);
    close(h);
    return 0;
}