Quarkys awesome ASAN

/****************************************************************
 * This file is part of NUSspli.                                *
 * (c) 2020 V10lator <v10lator@myway.de>                        *
 *                                                              *
 * Licensed under the MIT license (see LICENSE.txt for details) *
 ****************************************************************/

// This is based on Quarkys awesome ASAN from
// https://github.com/QuarkTheAwesome/RetroArch/blob/asan/wiiu/system/memory.c
// Quarky gave permission to relicense these codes under the MIT license.
// Copyright of Quarky applies to this file, too.

#include <wut-fixups.h>

#include <coreinit/memdefaultheap.h>

#include <memdebug.h>
#include <utils.h>

#ifdef NUSSPLI_DEBUG

#include <string.h>

#include <coreinit/debug.h>
#include <coreinit/dynload.h>
#include <coreinit/memory.h>
#include <coreinit/memexpheap.h>

typedef struct
{
    bool allocated;
    void *addr;
    size_t size;
    void *owner;
} lsan_allocation;

typedef struct
{
    char *name;
    void *codeStart;
    uint32_t unk01;
    uint32_t codeSize;
    void *dataStart;
    uint32_t unk02;
    uint32_t dataSize;
    uint32_t unk03;
    uint32_t unk04;
    uint32_t unk05;
} RPX_Info;

extern int OSDynLoad_GetRPLInfo(uint32_t unk01, uint32_t size, RPX_Info *out);

#define LSAN_ALLOCS_SZ 0x20000
static lsan_allocation lsan_allocs[LSAN_ALLOCS_SZ];
static MEMHeapHandle asan_heap;
static void *asan_heap_mem;
static unsigned int asan_sz;

static void *asan_shadow;
static unsigned int asan_shadow_off;
static bool asan_ready = false;

static void *codeStart;
static void *codeEnd;

#define MEM_TO_SHADOW(ptr) (unsigned char*)((((void*)ptr - asan_heap_mem) >> 3) + asan_shadow)
#define SET_SHADOW(shadow, ptr) (*shadow) |= 1 << (((unsigned int)ptr) & 0x7)
#define CLR_SHADOW(shadow, ptr) (*shadow) &= ~(1 << (((unsigned int)ptr) & 0x7))
#define GET_SHADOW(shadow, ptr) (*shadow >> (((unsigned int)ptr) & 0x7)) & 0x1

bool __asan_checkrange(void* ptr, size_t sz)
{
    if(ptr >= asan_heap_mem && ptr <= asan_shadow)
        for(int i = 0; i < sz; i++)
            if(!GET_SHADOW(MEM_TO_SHADOW(ptr + i), ptr + i))
                return false;
    return true;
}

void getSymbolName(void *ptr, char *out, int outLen)
{
    OSGetSymbolName((unsigned int)ptr, out, outLen);
    if(ptr >= codeStart && ptr < codeEnd)
    {
        char* tmp = strchr(out, '|') + sizeof(char);
        strcpy(out, tmp);
    }
}

void __asan_doreport(void* ptr, void* caller, const char* type)
{
    debugPrintf("======================== ASAN: out-of-bounds %s", type);

    char symbolName[128];
    getSymbolName(caller, symbolName, 127);
    debugPrintf("Source: 0x%08X: %s", caller, symbolName);

    //Thanks, LSAN!
    int closest_over = -1, closest_under = -1, uaf = -1;
    for(uint32_t i = 0; i < LSAN_ALLOCS_SZ; i++)
    {
        if(lsan_allocs[i].allocated)
        {
            if(ptr > lsan_allocs[i].addr + lsan_allocs[i].size)
            {
                if(closest_over == -1 || ptr - (lsan_allocs[i].addr + lsan_allocs[i].size) < ptr - (lsan_allocs[closest_over].addr + lsan_allocs[closest_over].size))
                    closest_over = i;
            }
            else if(ptr < lsan_allocs[i].addr)
            {
                if (closest_under == -1 || lsan_allocs[i].addr - ptr < lsan_allocs[closest_under].addr - ptr)
                    closest_under = i;
            }
        }
        else if(ptr > lsan_allocs[i].addr && ptr < lsan_allocs[i].addr + lsan_allocs[i].size && uaf == -1)
            uaf = i;
    }

    debugPrintf("Bad %s was on address 0x%08X. Guessing possible issues:", type, ptr);
    debugPrintf("Guess     Chunk      ChunkSz    Dist       ChunkOwner");
    if(closest_over >= 0)
    {
        getSymbolName(lsan_allocs[closest_over].owner, symbolName, 127);
        debugPrintf("Overflow  0x%08X 0x%08X 0x%08X %s", lsan_allocs[closest_over].addr, lsan_allocs[closest_over].size, ptr - (lsan_allocs[closest_over].addr + lsan_allocs[closest_over].size), symbolName);
    }
    if(closest_under >= 0)
    {
        getSymbolName(lsan_allocs[closest_under].owner, symbolName, 127);
        debugPrintf("Underflow 0x%08X 0x%08X 0x%08X %s", lsan_allocs[closest_under].addr, lsan_allocs[closest_under].size, ptr - (lsan_allocs[closest_under].addr + lsan_allocs[closest_under].size), symbolName);
    }
    if(uaf >= 0)
    {
        getSymbolName(lsan_allocs[uaf].owner, symbolName, 127);
        debugPrintf("UaF       0x%08X 0x%08X 0x%08X %s", lsan_allocs[uaf].addr, lsan_allocs[uaf].size, ptr - (lsan_allocs[uaf].addr + lsan_allocs[uaf].size), symbolName);
    }
}

void __asan_before_dynamic_init(const char *module_name)
{
    debugPrintf("STUB: __asan_before_dynamic_init(%s)", module_name);
}

void __asan_after_dynamic_init()
{
    debugPrintf("STUB: __asan_after_dynamic_init()");
}

void __asan_loadN_noabort(void* ptr, unsigned int size)
{
    if(asan_ready && !__asan_checkrange(ptr, (size_t)size))
        __asan_doreport(ptr, __builtin_return_address(0), "load");
}

void __asan_storeN_noabort(void* ptr, unsigned int size)
{
    if(asan_ready && !__asan_checkrange(ptr, (size_t)size))
        __asan_doreport(ptr, __builtin_return_address(0), "store");
}

#define ASAN_GENFUNC(type, num) void __asan_##type##num##_noabort(void* ptr) \
{ \
    if(asan_ready && !__asan_checkrange(ptr, num)) \
        __asan_doreport(ptr, __builtin_return_address(0), #type); \
}
ASAN_GENFUNC(load, 1);
ASAN_GENFUNC(store, 1);
ASAN_GENFUNC(load, 2);
ASAN_GENFUNC(store, 2);
ASAN_GENFUNC(load, 4);
ASAN_GENFUNC(store, 4);
ASAN_GENFUNC(load, 8);
ASAN_GENFUNC(store, 8);
ASAN_GENFUNC(load, 16);
ASAN_GENFUNC(store, 16);

void __asan_handle_no_return()
{
    //We only do heap checking, so no need to fiddle here
    debugPrintf("__asan_handle_no_return()");
}

static inline void setShadow(void *ptr, uint32_t size) // inline so we don't have to use __builtin_return_address(1)
{
    if(ptr != NULL)
    {
        bool lsan_ok = false;
        for(int i = 0; i < LSAN_ALLOCS_SZ; i++)
            if(!lsan_allocs[i].allocated)
            {
                lsan_allocs[i].allocated = true;
                lsan_allocs[i].addr = ptr;
                lsan_allocs[i].size = size;
                lsan_allocs[i].owner = __builtin_return_address(0);
                lsan_ok = true;
                break;
            }

        if(!lsan_ok)
            debugPrintf("[A/LSAN] WARNING: Too many allocs!");
        else
            for(size_t i = 0; i < size; i++)
                SET_SHADOW(MEM_TO_SHADOW(ptr + i), ptr + i);

    }
}

void *ASANAllocAligned(uint32_t size, int align)
{
    void *ptr = MEMAllocFromExpHeapEx(asan_heap, size, align);
    setShadow(ptr, size);
    return ptr;
}

void *ASANAlloc(uint32_t size)
{
    void *ptr = MEMAllocFromExpHeapEx(asan_heap, size, 4);
    setShadow(ptr, size);
    return ptr;
}

void ASANFree(void *ptr)
{
    if(ptr == NULL)
        return;

    size_t sz = 0;

    bool lsan_ok = false;
    for(int i = 0; i < LSAN_ALLOCS_SZ; i++)
    {
        if(lsan_allocs[i].allocated && lsan_allocs[i].addr == ptr)
        {
            lsan_allocs[i].allocated = false;
            sz = lsan_allocs[i].size; //Thanks, LSAN!
            lsan_ok = true;
            break;
        }
    }

    if(!lsan_ok)
    {
        debugPrintf("[LSAN] WARNING: attempted free 0x%08X; not in table!", ptr);
        return;
    }

    for(size_t i = 0; i < sz; i++)
        CLR_SHADOW(MEM_TO_SHADOW(ptr + i), ptr + i);

    MEMFreeToExpHeap(asan_heap, ptr);
}

MEMAllocFromDefaultHeapFn defaultAllocFn;
MEMAllocFromDefaultHeapExFn defaultAllocAlignedFn;
MEMFreeToDefaultHeapFn defaultFreeFn;

void injectASAN()
{
    MEMAllocFromDefaultHeap = ASANAlloc;
    MEMAllocFromDefaultHeapEx = ASANAllocAligned;
    MEMFreeToDefaultHeap = ASANFree;
}

void initASAN()
{
    RPX_Info rpxInfo[1];
    OSDynLoad_GetRPLInfo(0, 1, rpxInfo);
    codeStart = rpxInfo[0].codeStart;
    codeEnd = codeStart + rpxInfo[0].codeSize;

    for(uint32_t i = 0; i < LSAN_ALLOCS_SZ; i++)
        lsan_allocs[i].allocated = false;

    debugPrintf("[LSAN] LSAN initialized. Stack from 0x%08X to 0x%08X.", codeStart, codeEnd);

    MEMHeapHandle mem2 = MEMGetBaseHeapHandle(MEM_BASE_HEAP_MEM2);
    asan_sz = (MEMGetAllocatableSizeForExpHeapEx(mem2, 4) >> 2) & ~0x3;
    asan_heap_mem = MEMAllocFromDefaultHeap(asan_sz << 1);

    asan_heap = MEMCreateExpHeapEx(asan_heap_mem, asan_sz, 0);
    uint32_t asan_heap_sz = MEMGetTotalFreeSizeForExpHeap(asan_heap);

    char *multiplierName[2];
    if(asan_heap_sz < 1024)
        multiplierName[0] = "B";
    else if(asan_heap_sz < 1024 * 1024)
    {
        asan_heap_sz >>= 10;
        multiplierName[0] = "KB";
    }
    else
    {
        asan_heap_sz >>= 20;
        multiplierName[0] = "MB";
    }

    uint32_t wrapsize = asan_sz;
    if(wrapsize < 1024)
        multiplierName[1] = "B";
    else if(wrapsize < 1024 * 1024)
    {
        wrapsize >>= 10;
        multiplierName[1] = "KB";
    }
    else
    {
        wrapsize >>= 20;
        multiplierName[1] = "MB";
    }

    debugPrintf("[ASAN] Allocated wrapheap at 0x%08X, size %d %s. Avail RAM is %d %s.", asan_heap, wrapsize, multiplierName[1], asan_heap_sz, multiplierName[0]);

    asan_shadow = (void*)(((unsigned int)asan_heap_mem) + asan_sz);
    asan_shadow_off = (unsigned int)asan_heap - (unsigned int)asan_shadow;

    OSBlockSet(asan_shadow, 0, asan_sz);

    debugPrintf("[ASAN] Shadow at 0x%08X. Final shadow address at 0x%08X, final alloced at 0x%08X", asan_shadow, asan_shadow + asan_sz, asan_heap_mem + (asan_sz << 1));
    asan_ready = true;

    defaultAllocFn = MEMAllocFromDefaultHeap;
    defaultAllocAlignedFn = MEMAllocFromDefaultHeapEx;
    defaultFreeFn = MEMFreeToDefaultHeap;

    injectASAN();
}

void deinitASAN()
{
    unsigned long long alc = 0;
    for(uint32_t i = 0; i < LSAN_ALLOCS_SZ; i++)
        if(lsan_allocs[i].allocated)
            alc++;
    debugPrintf("[LSAN] Currently allocated: %llu", alc);

    MEMAllocFromDefaultHeap = defaultAllocFn;
    MEMAllocFromDefaultHeapEx = defaultAllocAlignedFn;
    MEMFreeToDefaultHeap = defaultFreeFn;

    asan_ready = false;

    MEMDestroyExpHeap(asan_heap);
    MEMFreeToDefaultHeap(asan_heap_mem);
}

#endif // ifdef NUSSPLI_DEBUG