/* -*- indent-tabs-mode: nil -*-
 *
 * This file is part of Funchook.
 * https://github.com/kubo/funchook
 *
 * Funchook is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the
 * Free Software Foundation, either version 2 of the License, or (at your
 * option) any later version.
 *
 * As a special exception, the copyright holders of this library give you
 * permission to link this library with independent modules to produce an
 * executable, regardless of the license terms of these independent
 * modules, and to copy and distribute the resulting executable under
 * terms of your choice, provided that you also meet, for each linked
 * independent module, the terms and conditions of the license of that
 * module. An independent module is a module which is not derived from or
 * based on this library. If you modify this library, you may extend this
 * exception to your version of the library, but you are not obliged to
 * do so. If you do not wish to do so, delete this exception statement
 * from your version.
 *
 * Funchook 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 General Public License
 * for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Funchook. If not, see <http://www.gnu.org/licenses/>.
 */
#define PSAPI_VERSION 1
#include <stdint.h>
#include <stdio.h>
#include <windows.h>
#include <psapi.h>
#include "funchook_internal.h"

typedef struct page_info {
    struct page_info *next;
    struct page_info *prev;
    int num_used;
    char used[1];
} page_list_t;

const size_t page_size = PAGE_SIZE; /* 4K */
const size_t allocation_unit = ALLOCATION_UNIT; /* 64K */

static size_t max_num_pages = ALLOCATION_UNIT / PAGE_SIZE - 1; /* 15 */
static page_list_t page_list = {
    &page_list,
    &page_list,
};

static const char *to_errmsg(DWORD err, char *buf, size_t bufsiz)
{
    size_t len = FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                                NULL, err, MAKELANGID(LANG_ENGLISH, SUBLANG_ENGLISH_US),
                                buf, (DWORD)bufsiz, NULL);
    if (len == 0) {
        return "Unknown Error";
    }
    if (len >= bufsiz) {
        len = bufsiz - 1;
    }
    while (len > 0 && (buf[len - 1] == '\r' || buf[len - 1] == '\n')) {
        len--;
    }
    buf[len] = '\0';
    return buf;
}

funchook_t *funchook_alloc(void)
{
    size_t size = ROUND_UP(funchook_size, page_size);
    return VirtualAlloc(NULL, size, MEM_COMMIT, PAGE_READWRITE);
}

int funchook_free(funchook_t *funchook)
{
    VirtualFree(funchook, 0, MEM_RELEASE);
    return 0;
}

/* Reserve 64K bytes (allocation_unit) and use the first
 * 4K bytes (1 page) as the control page.
 */
static int alloc_page_info(funchook_t *funchook, page_list_t **pl_out, void *hint)
{
    void *addr;
    page_list_t *pl;
#ifdef CPU_64BIT
    void *old_hint = hint;
    while (1) {
        MEMORY_BASIC_INFORMATION mbi;
        if (VirtualQuery(hint, &mbi, sizeof(mbi)) == 0) {
            DWORD err = GetLastError();
            char errbuf[128];

            funchook_set_error_message(funchook, "Failed to execute VirtualQuery (addr=%p, error=%lu(%s))",
                                       hint,
                                       err, to_errmsg(err, errbuf, sizeof(errbuf)));
            return FUNCHOOK_ERROR_MEMORY_FUNCTION;
        }
        funchook_log(funchook, "  process map: %016I64x-%016I64x %s\n",
                     (size_t)mbi.BaseAddress, (size_t)mbi.BaseAddress + mbi.RegionSize,
                     (mbi.State == MEM_FREE) ? "free" : "used");
        if (mbi.State == MEM_FREE) {
            size_t addr = ROUND_UP((size_t)mbi.BaseAddress, allocation_unit);
            intptr_t diff = addr - (size_t)mbi.BaseAddress;
            if (diff >= 0) {
                if (mbi.RegionSize - diff >= allocation_unit) {
                    hint = (void*)addr;
                    funchook_log(funchook, "  change hint address from %p to %p\n",
                                 old_hint, hint);
                    break;
                }
            }
        }
        hint = (void*)((size_t)mbi.BaseAddress + mbi.RegionSize);
    }
#else
    hint = NULL;
#endif
    pl = VirtualAlloc(hint, allocation_unit, MEM_RESERVE, PAGE_NOACCESS);
    if (pl == NULL) {
        DWORD err = GetLastError();
        char errbuf[128];

        funchook_set_error_message(funchook, "Failed to reserve memory %p (hint=%p, size=%"PRIuPTR", errro=%lu(%s))",
                                   pl, hint, allocation_unit,
                                   err, to_errmsg(err, errbuf, sizeof(errbuf)));
        return FUNCHOOK_ERROR_MEMORY_ALLOCATION;
    }
    funchook_log(funchook, "  reserve memory %p (hint=%p, size=%"PRIuPTR")\n", pl, hint, allocation_unit);
    addr = VirtualAlloc(pl, page_size, MEM_COMMIT, PAGE_READWRITE);
    if (addr == NULL) {
        DWORD err = GetLastError();
        char errbuf[128];

        funchook_set_error_message(funchook, "Failed to commit memory %p for read-write (hint=%p, size=%"PRIuPTR", error=%lu(%s))",
                                   addr, pl, page_size,
                                   err, to_errmsg(err, errbuf, sizeof(errbuf)));
        VirtualFree(pl, 0, MEM_RELEASE);
        return FUNCHOOK_ERROR_MEMORY_FUNCTION;
    }
    funchook_log(funchook, "  commit memory %p for read-write (hint=%p, size=%"PRIuPTR")\n", addr, pl, page_size);
    pl->next = page_list.next;
    pl->prev = &page_list;
    page_list.next->prev = pl;
    page_list.next = pl;
    *pl_out = pl;
    return 0;
}

/*
 * Get one page from page_list, commit it and return it.
 */
int funchook_page_alloc(funchook_t *funchook, funchook_page_t **page_out, uint8_t *func, ip_displacement_t *disp)
{
    page_list_t *pl;
    funchook_page_t *page = NULL;
    size_t i;

    for (pl = page_list.next; pl != &page_list; pl = pl->next) {
        for (i = 0; i < max_num_pages; i++) {
            if (!pl->used[i]) {
                funchook_page_t *p = (funchook_page_t *)((size_t)pl + (i + 1) * page_size);
                if (funchook_page_avail(funchook, p, 0, func, disp)) {
                    page = p;
                    goto exit_loop;
                }
            }
        }
    }
exit_loop:
    if (page == NULL) {
        /* no page_list is available. */
        int rv = alloc_page_info(funchook, &pl, func);
        if (rv != 0) {
            return rv;
        }
        i = 0;
        page = (funchook_page_t *)((size_t)pl + page_size);
    }
    if (VirtualAlloc(page, page_size, MEM_COMMIT, PAGE_READWRITE) == NULL) {
        DWORD err = GetLastError();
        char errbuf[128];

        funchook_set_error_message(funchook, "Failed to commit page %p (base=%p(used=%d), idx=%"PRIuPTR", size=%"PRIuPTR", error=%lu(%s))",
                                   page, pl, pl->num_used, i, page_size,
                                   err, to_errmsg(err, errbuf, sizeof(errbuf)));
        return FUNCHOOK_ERROR_MEMORY_FUNCTION;
    }
    pl->used[i] = 1;
    pl->num_used++;
    funchook_log(funchook, "  commit page %p (base=%p(used=%d), idx=%"PRIuPTR", size=%"PRIuPTR")\n",
                 page, pl, pl->num_used, i, page_size);
    *page_out = page;
    return 0;
}

/*
 * Back to one page to page_list.
 */
int funchook_page_free(funchook_t *funchook, funchook_page_t *page)
{
    page_list_t *pl = (page_list_t *)((size_t)page & ~(allocation_unit - 1));
    size_t idx = ((size_t)page - (size_t)pl) / page_size - 1;
    BOOL ok;

    ok = VirtualFree(page, page_size, MEM_DECOMMIT);
    if (!ok) {
        DWORD err = GetLastError();
        char errbuf[128];

        funchook_set_error_message(funchook, "Failed to decommit page %p (base=%p(used=%d), idx=%"PRIuPTR", size=%"PRIuPTR", error=%lu(%s))",
                                   page, pl, pl->num_used, idx, page_size,
                                   err, to_errmsg(err, errbuf, sizeof(errbuf)));
        return FUNCHOOK_ERROR_MEMORY_FUNCTION;
    }
    funchook_log(funchook, "  decommit page %p (base=%p(used=%d), idx=%"PRIuPTR", size=%"PRIuPTR")\n",
                 page, pl, pl->num_used, idx, page_size);
    pl->num_used--;
    pl->used[idx] = 0;
    if (pl->num_used != 0) {
        return 0;
    }
    /* all pages in this allocation unit are decommitted. delete this page_list */
    pl->next->prev = pl->prev;
    pl->prev->next = pl->next;
    ok = VirtualFree(pl, 0, MEM_RELEASE);
    if (!ok) {
        DWORD err = GetLastError();
        char errbuf[128];

        funchook_set_error_message(funchook, "Failed to release memory %p (size=%"PRIuPTR", error=%lu(%s))",
                                   pl, allocation_unit,
                                   err, to_errmsg(err, errbuf, sizeof(errbuf)));
        return FUNCHOOK_ERROR_MEMORY_FUNCTION;
    }
    funchook_log(funchook, "  release memory %p (size=%"PRIuPTR")\n",
                 pl, allocation_unit);
    return 0;
}

int funchook_page_protect(funchook_t *funchook, funchook_page_t *page)
{
    char errbuf[128];
    DWORD oldprot;
    BOOL ok = VirtualProtect(page, page_size, PAGE_EXECUTE_READ, &oldprot);

    if (ok) {
        funchook_log(funchook, "  protect page %p (size=%"PRIuPTR", prot=read,exec)\n",
                     page, page_size);
        return 0;
    }
    funchook_set_error_message(funchook, "Failed to protect page %p (size=%"PRIuPTR", prot=read,exec, error=%lu(%s))",
                               page, page_size,
                               GetLastError(), to_errmsg(GetLastError(), errbuf, sizeof(errbuf)));
    return FUNCHOOK_ERROR_MEMORY_FUNCTION;
}

int funchook_page_unprotect(funchook_t *funchook, funchook_page_t *page)
{
    char errbuf[128];
    DWORD oldprot;
    BOOL ok = VirtualProtect(page, page_size, PAGE_READWRITE, &oldprot);

    if (ok) {
        funchook_log(funchook, "  unprotect page %p (size=%"PRIuPTR", prot=read,write)\n",
                     page, page_size);
        return 0;
    }
    funchook_set_error_message(funchook, "Failed to unprotect page %p (size=%"PRIuPTR", prot=read,write, error=%lu(%s))",
                               page, page_size,
                               GetLastError(), to_errmsg(GetLastError(), errbuf, sizeof(errbuf)));
    return FUNCHOOK_ERROR_MEMORY_FUNCTION;
}

int funchook_unprotect_begin(funchook_t *funchook, mem_state_t *mstate, void *start, size_t len)
{
    char errbuf[128];
    size_t saddr = ROUND_DOWN((size_t)start, page_size);
    BOOL ok;

    mstate->addr = (void*)saddr;
    mstate->size = len + (size_t)start - saddr;
    mstate->size = ROUND_UP(mstate->size, page_size);
    ok = VirtualProtect(mstate->addr, mstate->size, PAGE_EXECUTE_READWRITE, &mstate->protect);
    if (ok) {
        funchook_log(funchook, "  unprotect memory %p (size=%"PRIuPTR") <- %p (size=%"PRIuPTR")\n",
                     mstate->addr, mstate->size, start, len);
        return 0;
    }
    funchook_set_error_message(funchook, "Failed to unprotect memory %p (size=%"PRIuPTR") <- %p (size=%"PRIuPTR", error=%lu(%s))",
                               mstate->addr, mstate->size, start, len,
                               GetLastError(), to_errmsg(GetLastError(), errbuf, sizeof(errbuf)));
    return FUNCHOOK_ERROR_MEMORY_FUNCTION;
}

int funchook_unprotect_end(funchook_t *funchook, const mem_state_t *mstate)
{
    char errbuf[128];
    DWORD oldprot;
    BOOL ok = VirtualProtect(mstate->addr, mstate->size, mstate->protect, &oldprot);

    if (ok) {
        funchook_log(funchook, "  protect memory %p (size=%"PRIuPTR")\n",
                     mstate->addr, mstate->size);
        return 0;
    }
    funchook_set_error_message(funchook, "Failed to protect memory %p (size=%"PRIuPTR", error=%lu(%s))",
                               mstate->addr, mstate->size,
                               GetLastError(), to_errmsg(GetLastError(), errbuf, sizeof(errbuf)));
    return FUNCHOOK_ERROR_MEMORY_FUNCTION;
}

static IMAGE_IMPORT_DESCRIPTOR *get_image_import_descriptor(HMODULE hMod, DWORD *cnt)
{
    IMAGE_DOS_HEADER *doshdr;
    IMAGE_NT_HEADERS *nthdr;
    IMAGE_DATA_DIRECTORY *dir;

    if (memcmp(hMod, "MZ", 2) != 0) {
        return NULL;
    }
    doshdr = (IMAGE_DOS_HEADER*)hMod;
    nthdr = (PIMAGE_NT_HEADERS)((size_t)hMod + doshdr->e_lfanew);
    dir = &nthdr->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (dir->VirtualAddress == 0) {
        return NULL;
    }
    *cnt = dir->Size / sizeof(IMAGE_IMPORT_DESCRIPTOR);
    return (IMAGE_IMPORT_DESCRIPTOR*)((size_t)hMod + dir->VirtualAddress);
}

void *funchook_resolve_func(funchook_t *funchook, void *func)
{
    char path[MAX_PATH];
    HMODULE hMod;
    BOOL ok;
    IMAGE_IMPORT_DESCRIPTOR *desc_head, *desc;
    insn_t *fn = (insn_t*)func;
    size_t pos = 0;
    DWORD cnt;

    if (*funchook_debug_file != '\0') {
        DWORD len = GetMappedFileNameA(GetCurrentProcess(), func, path, sizeof(path));
        if (len > 0) {
            funchook_log(funchook, "  func %p is in %.*s\n", func, (int)len, path);
        }
    }
#if defined CPU_X86 || defined CPU_X86_64
    if (fn[0] == 0xe9) {
        fn = (fn + 5) + *(int*)(fn + 1);
        funchook_log(funchook, "  relative jump to %p\n", fn);
    }
    if (fn[0] == 0xff && fn[1] == 0x25) {
#ifdef CPU_X86_64
        pos = (size_t)(fn + 6) + *(int*)(fn + 2);
#else
        pos = *(size_t*)(fn + 2);
#endif
        funchook_log(funchook, "  indirect jump to addresss at %p\n", (void*)pos);
    }
#endif
#if defined CPU_ARM64
#define ADRP_XIP0       0x90000010
#define ADRP_XIP0_MASK  0x9F00001F
#define ADRP_XIP0_IMMLO 0x60000000
#define ADRP_XIP0_IMMHI 0x00FFFFE0
#define LDR_XIP0        0xF9400210
#define LDR_XIP0_MASK   0xFFC003FF
#define LDR_XIP0_IMM12  0x003FFC00
#define BR_XIP0         0xD61F0200
    if ((fn[0] & ADRP_XIP0_MASK) == ADRP_XIP0 &&
        (fn[1] & LDR_XIP0_MASK) == LDR_XIP0 &&
        fn[2] == BR_XIP0) {
        // fn[0]: addrp xip0, immhi&immlo
        // fn[1]: ldr   xip0, [xip0,imm12]
        // fn[2]: br    xip0
        size_t addr = (size_t)fn & ~((1 << 12) - 1);
        size_t immhi = ((size_t)(fn[0] & ADRP_XIP0_IMMHI) >> 5) << (12 + 2);
        size_t immlo = ((size_t)(fn[0] & ADRP_XIP0_IMMLO) >> 29) << 12;
        size_t imm12 = ((size_t)(fn[1] & LDR_XIP0_IMM12) >> 10) << 3;
        pos = addr + immhi + immlo + imm12;
        // fprintf(stderr, "%016I64x: %08x %08x %08x : %I64x %I64x %I64x %I64x\n", (size_t)fn, fn[0], fn[1], fn[2], addr, immhi, immlo, imm12);
        funchook_log(funchook, "  indirect jump to addresss at %p\n", (void*)pos);
    }
#endif
    if (pos == 0) {
        return func;
    }
    ok = GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, func, &hMod);
    if (!ok) {
        return func;
    }

    desc_head = get_image_import_descriptor(hMod, &cnt);
    if (desc_head == NULL) {
        return func;
    }

    for (desc = desc_head; desc->Name != 0; desc++) {
        IMAGE_THUNK_DATA *addr_thunk = (IMAGE_THUNK_DATA*)((char*)hMod + desc->FirstThunk);

        while (addr_thunk->u1.Function != 0) {
            if (pos == (size_t)&addr_thunk->u1.Function) {
                func = (void*)addr_thunk->u1.Function;
                if (*funchook_debug_file != '\0') {
                    DWORD len = GetMappedFileNameA(GetCurrentProcess(), func, path, sizeof(path));
                    if (len > 0) {
                        funchook_log(funchook, "  -> func %p in %.*s\n", func, (int)len, path);
                    } else {
                        funchook_log(funchook, "  -> func %p\n", func);
                    }
                }
                return func;
            }
            addr_thunk++;
        }
    }
    return func;
}