#include <string.h>
#include "common.h"
#include "tizenasanenv.h"


template <typename Type, int STACK_SIZE>
class StaticStack {
    // We don't create constructor because
    // this class is used in a zeroed memory area
public:
    void push(Type addr)
    {
        _ASSERTE(m_pos < STACK_SIZE);

        m_data[m_pos++] = addr;
    }

    void pop()
    {
        _ASSERTE(m_pos > 0);
        --m_pos;
    }

    Type top()
    {
        _ASSERTE(m_pos > 0);

        return m_data[m_pos - 1];
    }

    bool empty()
    {
        return m_pos == 0;
    }

private:
    int m_pos;
    Type m_data[STACK_SIZE];
};

#include <pshpack1.h>
struct AuxiliaryCalls {
    LPVOID target;
    void (*pushAddr)(LPVOID addr);
    LPVOID (*popAddr)();
};

struct ReturnInfo {
    LPVOID addr;
    bool isSanitized;
};

extern "C" void __sanitizer_disable_interceptors();
extern "C" void __sanitizer_enable_interceptors();
extern "C" bool __sanitizer_interceptors_are_enabled();

extern LPVOID tizenASanWrapper;
extern UINT32 tizenASanWrapperSize;
extern UINT32 tizenASanWrapperEntryOffset;

// The maximum nesting of transitions between managed and unmanaged code that we support.
// This number is estimated from the common sense. We think this is enough to check any
// sane code (if it is not recursive) and it won't bloat TLS.  We do not use dynamic
// allocation because it complicates the process of memory management in TLS variables.
// It is used only for firmware with ASan and will not affect the release version.
#define MAX_STACK_DEPTH 128
static __thread StaticStack<ReturnInfo, MAX_STACK_DEPTH> s_retInfoStack;


static void DoEnable()
{
    _ASSERTE(__sanitizer_interceptors_are_enabled() == false);
    __sanitizer_enable_interceptors();
}

static void DoDisable()
{
    _ASSERTE(__sanitizer_interceptors_are_enabled() == true);
    __sanitizer_disable_interceptors();
}

static void PushAndEnableASan(LPVOID addr)
{
    _ASSERTE(__sanitizer_interceptors_are_enabled() == false);

    ReturnInfo retInfo = {
        .addr = addr,
        .isSanitized = false,
    };

    s_retInfoStack.push(retInfo);
    DoEnable();
}

static LPVOID PopAndDisableASan()
{
    _ASSERTE(__sanitizer_interceptors_are_enabled() == true);

    ReturnInfo retInfo = s_retInfoStack.top();
    s_retInfoStack.pop();

    _ASSERTE(retInfo.isSanitized == false);
    DoDisable();

    return retInfo.addr;
}

static void PushAndMayBeDisableASan(LPVOID addr)
{
    ReturnInfo retInfo = {
        .addr = addr,
        .isSanitized = __sanitizer_interceptors_are_enabled(),
    };

    if (retInfo.isSanitized)
        DoDisable();

    s_retInfoStack.push(retInfo);
}

static LPVOID PopAndMayBeEnableASan()
{
    _ASSERTE(__sanitizer_interceptors_are_enabled() == false);

    ReturnInfo retInfo = s_retInfoStack.top();
    s_retInfoStack.pop();

    if (retInfo.isSanitized)
        DoEnable();

    return retInfo.addr;
}

static LPVOID CreateWrapper(LPVOID target, void (*pushAddr)(LPVOID addr), LPVOID (*popAddr)())
{
    _ASSERTE(tizenASanWrapperEntryOffset == sizeof(AuxiliaryCalls));

    LPVOID wrapperSpace = (LPVOID)SystemDomain::GetGlobalLoaderAllocator()->GetExecutableHeap()->AllocMem(S_SIZE_T(tizenASanWrapperSize));

    AuxiliaryCalls calls = {
        .target = target,
	.pushAddr = pushAddr,
	.popAddr = popAddr,
    };

    // copy auxiliary calls
    memcpy(wrapperSpace, &calls, sizeof(calls));

    LPVOID entryPointer = (LPVOID)((UINT_PTR)wrapperSpace + tizenASanWrapperEntryOffset);
    LPVOID wrapperEntryPointer = (LPVOID)((UINT_PTR)&tizenASanWrapper + tizenASanWrapperEntryOffset);
    UINT32 wrapperCodeSize = tizenASanWrapperSize - tizenASanWrapperEntryOffset;

    // copy executable code wrapper
    memcpy(entryPointer, wrapperEntryPointer, wrapperCodeSize);

    FlushInstructionCache(GetCurrentProcess(), wrapperSpace, tizenASanWrapperSize);

    return entryPointer;
}


namespace TizenASanEnv {

LPVOID CreateWrapperSanitizedEntryPoint(LPVOID target)
{
    return CreateWrapper(target, PushAndEnableASan, PopAndDisableASan);
}

LPVOID CreateWrapperILCode(LPVOID target)
{
    return CreateWrapper(target, PushAndMayBeDisableASan, PopAndMayBeEnableASan);
}

} // namespace TizenASanEnv