// Copyright (c) Microsoft Corporation.  All rights reserved.
// Copyright (c) Lysann Tranvouez.  All rights reserved.

#include "mach_detours.h"

#include "detours_internal.h"
#include "arm64/detours_arm64.h"

#include <stdatomic.h>
#include <stdint.h>
#include <stdlib.h>

#include <mach/mach.h>
#include <mach/mach_vm.h>

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Trampoline Memory Management

typedef struct detour_region
{
    uint32_t signature;
    struct detour_region* next;  // Next region in list of regions.
    detour_trampoline* free_list_head;  // List of free trampolines in this region.
} detour_region;

static const uint32_t DETOUR_REGION_SIGNATURE = 'Rrtd';
static const uint32_t DETOUR_REGION_SIZE = 0x10000;
static const uint32_t DETOUR_TRAMPOLINES_PER_REGION = (DETOUR_REGION_SIZE / sizeof(detour_trampoline)) - 1;
static detour_region* s_regions_head = nullptr;
static detour_region* s_default_region = nullptr;

static mach_error_t internal_detour_writable_trampoline_regions()
{
    // Mark all the regions as writable.
    const mach_port_t port = mach_task_self();
    for (detour_region* pRegion = s_regions_head; pRegion != NULL; pRegion = pRegion->next) {
        const mach_error_t error = mach_vm_protect(port, (mach_vm_address_t)pRegion, DETOUR_REGION_SIZE, false, VM_PROT_READ | VM_PROT_WRITE);
        if (error != err_none) {
            return error;
        }
    }
    return err_none;
}

static void internal_detour_runnable_trampoline_regions()
{
    // Mark all the regions as executable.
    const mach_port_t port = mach_task_self();
    for (detour_region* pRegion = s_regions_head; pRegion != NULL; pRegion = pRegion->next) {
        const mach_error_t error = mach_vm_protect(port, (mach_vm_address_t)pRegion, DETOUR_REGION_SIZE, false, VM_PROT_READ | VM_PROT_EXECUTE);
        if (error != err_none) {
            DETOUR_BREAK();
        }
    }
}

static void internal_detour_free_trampoline(detour_trampoline* trampoline)
{
    detour_region* region = (detour_region*)((uintptr_t)trampoline & ~(uintptr_t)0xffff);

    memset(trampoline, 0, sizeof(*trampoline));
    trampoline->ptr_remain = (uint8_t*)region->free_list_head;
    region->free_list_head = trampoline;
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Transactions

typedef enum detour_operation_kind {
    detour_operation_kind_attach,
    detour_operation_kind_detach,
} detour_operation_kind;

typedef struct detour_operation
{
    struct detour_operation* next;
    detour_operation_kind kind;
    uint8_t** pointer;
    uint8_t* target;
    detour_trampoline* trampoline;
    vm_prot_t perm;
} detour_operation;

typedef struct detour_pending_thread
{
    struct detour_pending_thread* next;
    thread_t thread;
} detour_pending_thread;

static _Atomic(thread_t) s_transaction_thread = THREAD_NULL;
static detour_operation* s_pending_operations_head = nullptr;
static detour_pending_thread* s_pending_threads_head = nullptr;

static mach_error_t s_pending_error = err_none;
static void** s_pending_error_pointer = nullptr;

mach_error_t detour_transaction_begin()
{
    // Make sure only one thread can start a transaction.
    thread_t expected = THREAD_NULL;
    // ReSharper disable once CppIncompatiblePointerConversion
    if (!atomic_compare_exchange_strong(&s_transaction_thread, &expected, mach_thread_self())) {
        return detour_err_in_progress;
    }

    s_pending_operations_head = nullptr;
    s_pending_threads_head = nullptr;
    s_pending_error_pointer = nullptr;

    // Make sure the trampoline pages are writable.
    s_pending_error = internal_detour_writable_trampoline_regions();

    return s_pending_error;
}

mach_error_t detour_transaction_abort()
{
    if (s_transaction_thread != mach_thread_self()) {
        return detour_err_wrong_thread;
    }

    // Restore all the page permissions.
    const mach_port_t port = mach_task_self();
    for (detour_operation* operation = s_pending_operations_head; operation != nullptr;) {
        DETOUR_CHECK(mach_vm_protect(port, (mach_vm_address_t)operation->target, operation->trampoline->restore_code_size, false, operation->perm));

        if (operation->kind == detour_operation_kind_attach) {
            if (operation->trampoline) {
                internal_detour_free_trampoline(operation->trampoline);
                operation->trampoline = nullptr;
            }
        }

        detour_operation* next = operation->next;
        free(operation);
        operation = next;
    }
    s_pending_operations_head = nullptr;

    // Make sure the trampoline pages are no longer writable.
    internal_detour_runnable_trampoline_regions();

    // Resume any suspended threads.
    for (detour_pending_thread* thread = s_pending_threads_head; thread != nullptr;) {
        // There is nothing we can do if this fails.
        DETOUR_CHECK(thread_resume(thread->thread));

        detour_pending_thread* next = thread->next;
        free(thread);
        thread = next;
    }
    s_pending_threads_head = nullptr;
    s_transaction_thread = THREAD_NULL;

    return err_none;
}