/*
 * Copyright 2008, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Distributed under the terms of the MIT License.
 */

#include <stdio.h>
#include <stdlib.h>

#include <new>

#include <OS.h>

#include <runtime_loader.h>
#include <util/OpenHashTable.h>

#include "arch/ltrace_stub.h"


//#define TRACE_PRINTF debug_printf
#define TRACE_PRINTF ktrace_printf


static void* function_call_callback(const void* stub, const void* args);


struct PatchEntry {
        PatchEntry*     originalTableLink;
        PatchEntry*     patchedTableLink;

        void*           originalFunction;
        void*           patchedFunction;
        const char*     functionName;

        static PatchEntry* Create(const char* name, void* function)
        {
                // TODO memory should be executable, use mmap with PROT_EXEC
                void* memory = malloc(_ALIGN(sizeof(PatchEntry))
                        + arch_call_stub_size());
                if (memory == NULL)
                        return NULL;

                PatchEntry* entry = new(memory) PatchEntry;

                void* stub = (uint8*)memory + _ALIGN(sizeof(PatchEntry));
                arch_init_call_stub(stub, &function_call_callback, function);

                entry->originalFunction = function;
                entry->patchedFunction = stub;
                entry->functionName = name;

                return entry;
        }
};


struct OriginalTableDefinition {
        typedef const void*     KeyType;
        typedef PatchEntry      ValueType;

        size_t HashKey(const void* key) const
        {
                return (addr_t)key >> 2;
        }

        size_t Hash(PatchEntry* value) const
        {
                return HashKey(value->originalFunction);
        }

        bool Compare(const void* key, PatchEntry* value) const
        {
                return value->originalFunction == key;
        }

        PatchEntry*& GetLink(PatchEntry* value) const
        {
                return value->originalTableLink;
        }
};


struct PatchedTableDefinition {
        typedef const void*     KeyType;
        typedef PatchEntry      ValueType;

        size_t HashKey(const void* key) const
        {
                return (addr_t)key >> 2;
        }

        size_t Hash(PatchEntry* value) const
        {
                return HashKey(value->patchedFunction);
        }

        bool Compare(const void* key, PatchEntry* value) const
        {
                return value->patchedFunction == key;
        }

        PatchEntry*& GetLink(PatchEntry* value) const
        {
                return value->patchedTableLink;
        }
};


static rld_export* sRuntimeLoaderInterface;
static runtime_loader_add_on_export* sRuntimeLoaderAddOnInterface;

static BOpenHashTable<OriginalTableDefinition> sOriginalTable;
static BOpenHashTable<PatchedTableDefinition> sPatchedTable;


static void*
function_call_callback(const void* stub, const void* _args)
{
        PatchEntry* entry = sPatchedTable.Lookup(stub);
        if (entry == NULL)
{
TRACE_PRINTF("function_call_callback(): CALLED FOR UNKNOWN FUNCTION!\n");
                return NULL;
}

        char buffer[1024];
        size_t bufferSize = sizeof(buffer);
        size_t written = 0;

        const ulong* args = (const ulong*)_args;
        written += snprintf(buffer, bufferSize, "ltrace: %s(",
                entry->functionName);
        for (int32 i = 0; i < 5; i++) {
                written += snprintf(buffer + written, bufferSize - written, "%s%#lx",
                        i == 0 ? "" : ", ", args[i]);
        }
        written += snprintf(buffer + written, bufferSize - written, ")");
        TRACE_PRINTF("%s\n", buffer);

        return entry->originalFunction;
}


static void
symbol_patcher(void* cookie, image_t* rootImage, image_t* image,
        const char* name, image_t** foundInImage, void** symbol, int32* type)
{
        TRACE_PRINTF("symbol_patcher(%p, %p, %p, \"%s\", %p, %p, %" B_PRId32 ")\n",
                cookie, rootImage, image, name, *foundInImage, *symbol, *type);

        // patch functions only
        if (*type != B_SYMBOL_TYPE_TEXT)
                return;

        // already patched?
        PatchEntry* entry = sOriginalTable.Lookup(*symbol);
        if (entry != NULL) {
                *foundInImage = NULL;
                *symbol = entry->patchedFunction;
                return;
        }

        entry = PatchEntry::Create(name, *symbol);
        if (entry == NULL)
                return;

        sOriginalTable.Insert(entry);
        sPatchedTable.Insert(entry);

        TRACE_PRINTF("  -> patching to %p\n", entry->patchedFunction);

        *foundInImage = NULL;
        *symbol = entry->patchedFunction;
}


static void
ltrace_stub_init(rld_export* standardInterface,
        runtime_loader_add_on_export* addOnInterface)
{
        TRACE_PRINTF("ltrace_stub_init(%p, %p)\n", standardInterface, addOnInterface);
        sRuntimeLoaderInterface = standardInterface;
        sRuntimeLoaderAddOnInterface = addOnInterface;

        sOriginalTable.Init();
        sPatchedTable.Init();
}


static void
ltrace_stub_image_loaded(image_t* image)
{
        TRACE_PRINTF("ltrace_stub_image_loaded(%p): \"%s\" (%" B_PRId32 ")\n",
                image, image->path, image->id);

        if (sRuntimeLoaderAddOnInterface->register_undefined_symbol_patcher(image,
                        symbol_patcher, (void*)(addr_t)0xc0011eaf) != B_OK) {
                TRACE_PRINTF("  failed to install symbol patcher\n");
        }
}


static void
ltrace_stub_image_relocated(image_t* image)
{
        TRACE_PRINTF("ltrace_stub_image_relocated(%p): \"%s\" (%" B_PRId32 ")\n",
                image, image->path, image->id);
}


static void
ltrace_stub_image_initialized(image_t* image)
{
        TRACE_PRINTF("ltrace_stub_image_initialized(%p): \"%s\" (%" B_PRId32 ")\n",
                image, image->path, image->id);
}


static void
ltrace_stub_image_uninitializing(image_t* image)
{
        TRACE_PRINTF("ltrace_stub_image_uninitializing(%p): \"%s\" (%" B_PRId32
                ")\n",image, image->path, image->id);
}


static void
ltrace_stub_image_unloading(image_t* image)
{
        TRACE_PRINTF("ltrace_stub_image_unloading(%p): \"%s\" (%" B_PRId32 ")\n",
                image, image->path, image->id);
}


// interface for the runtime loader
runtime_loader_add_on __gRuntimeLoaderAddOn = {
        RUNTIME_LOADER_ADD_ON_VERSION,  // version
        0,                                                              // flags

        ltrace_stub_init,

        ltrace_stub_image_loaded,
        ltrace_stub_image_relocated,
        ltrace_stub_image_initialized,
        ltrace_stub_image_uninitializing,
        ltrace_stub_image_unloading
};