#define WIN32_MEAN_AND_LEAN
#include <Windows.h>
#include <stdint.h>
#include <stdio.h>
#include <assert.h>
#include <sanitizer/asan_interface.h>

#define IsPowerOfTwo(value, pot) ((((size_t)value) & (((size_t)pot) - 1)) == 0)
#define AsanAlignment 8

#define PrintByteSpacing(space) \
    if (index) \
        printf(space); \
    if (index && (index % 8 == 0)) \
        printf("| ")

static void PrintByteArray(char const *array, size_t array_size)
{
    printf("   Byte Array                    ");
    for (size_t index = 0 ; index < array_size; index++) {
        PrintByteSpacing(" ");
        printf("%02zu", index);
    }
    printf("\n");
}

static void PrintPoisonedBytes(int step, char const *array, size_t array_size)
{
    printf("%d. __asan_address_is_poisoned    ", step);
    for (size_t index = 0 ; index < array_size; index++) {
        PrintByteSpacing("  ");
        printf("%c", __asan_address_is_poisoned(array + index) ? 'x' : ' ');
    }
    printf("\n");
}

static void PrintPoisonMemoryIndex(int step, char const *array, size_t array_size, size_t poison_index)
{
    printf("%d. __asan_poison_memory_region   ", step);
    for (size_t index = 0; index < array_size; index++) {
        PrintByteSpacing(" ");
        printf("%c ", index == poison_index ? 'x' : ' ');
    }
    printf("\n");
}

static void PrintPoisonMemoryRegion(int step, char const *array, size_t array_size, size_t poison_start_index, size_t poison_end_index)
{
    printf("%d. __asan_poison_memory_region   ", step);
    for (size_t index = 0; index < array_size; index++) {
        PrintByteSpacing(" ");
        printf("%c ", index >= poison_start_index && index <= poison_end_index ? 'x' : ' ');
    }
    printf("\n");
}

int main()
{
    printf("Here we demonstrate that ASAN poison-ing will only poison the\n"
           "byte region if the region meets an 8 byte boundary. It will only\n"
           "poison bytes upto the 8 byte boundary, any bytes that straddle\n"
           "the boundary that do not hit the next 8 byte boundary are not\n"
           "poison-ed.\n\n");

    uint32_t const ASAN_ALIGNMENT = 8;
    uint32_t const REGION_WINDOW  = 7;
    char array[16]                = {};
    for (size_t poison_index = 0; poison_index < sizeof(array) - (REGION_WINDOW - 1); poison_index++) {
        assert(IsPowerOfTwo(array, ASAN_ALIGNMENT));

        // NOTE: Print byte array ==================================================================

        PrintByteArray(array, sizeof(array));

        // NOTE: Poison the array ==================================================================

        __asan_poison_memory_region(&array[poison_index], REGION_WINDOW);
        PrintPoisonMemoryRegion(1 /*step*/, array, sizeof(array), poison_index, poison_index + (REGION_WINDOW - 1));

        // NOTE: Print the poison-ed bytes =========================================================

        PrintPoisonedBytes(2 /*step*/, array, sizeof(array));

        // NOTE: Cleanup ===========================================================================

        __asan_unpoison_memory_region(array, sizeof(array));
        printf("\n");
    }

    printf(
        "Now we demonstrate that unpoison-ing 1 byte in the 8 byte window\n"
        "will unpoison all bytes prior to it up until the start of the window\n"
        "until the previous 8 byte boundary.\n\n");

    for (size_t unpoison_index = 0; unpoison_index < sizeof(array); unpoison_index++) {
        assert(IsPowerOfTwo(array, ASAN_ALIGNMENT));

        // NOTE: Print byte array ==================================================================

        PrintByteArray(array, sizeof(array));

        // NOTE: Poison the array ==================================================================

        __asan_poison_memory_region(array, sizeof(array));
        PrintPoisonMemoryRegion(1 /*step*/, array, sizeof(array), 0, sizeof(array) - 1);

        // NOTE: Print the poison-ed bytes =========================================================

        PrintPoisonedBytes(2 /*step*/, array, sizeof(array));

        // NOTE: Unpoison byte at each position separately =========================================

        printf("3. __asan_unpoison_memory_region ");
        __asan_unpoison_memory_region(array + unpoison_index, 1);

        for (size_t index = 0 ; index < sizeof(array); index++) {
            PrintByteSpacing("  ");
            printf("%c", index == unpoison_index ? 'x' : ' ');
        }
        printf("\n");

        // NOTE: Print the poison-ed bytes =========================================================

        printf("4. __asan_address_is_poisoned    ");
        for (size_t index = 0 ; index < sizeof(array); index++) {
            PrintByteSpacing("  ");
            printf("%c", __asan_address_is_poisoned(array + index) ? 'x' : ' ');
        }
        printf("\n\n");

        // NOTE: Cleanup ===========================================================================

        __asan_unpoison_memory_region(array, sizeof(array));
    }

    printf(
        "Unpoison-ing across 8 byte boundaries may lead to undesired\n"
        "behaviour, with all bytes on the left side of the boundary being\n"
        "unpoisoned\n\n");
    {
        // NOTE: Print byte array ==================================================================

        char array[16] = {};
        PrintByteArray(array, sizeof(array));

        // NOTE: Poison the array ==================================================================

        __asan_poison_memory_region(array, sizeof(array));
        PrintPoisonMemoryRegion(1 /*step*/, array, sizeof(array), 0, sizeof(array) - 1);

        // NOTE: Print the poison-ed bytes =========================================================

        PrintPoisonedBytes(2 /*step*/, array, sizeof(array));

        // NOTE: Unpoison across the 8 byte boundary ===============================================

        printf("3. __asan_unpoison_memory_region ");
        uint32_t const start_poison_index = 7;
        uint32_t const bytes_to_unpoison  = 2;
        uint32_t const end_poison_index   = start_poison_index + (bytes_to_unpoison - 1);
        __asan_unpoison_memory_region(array + start_poison_index, bytes_to_unpoison);

        for (size_t index = 0 ; index < sizeof(array); index++) {
            PrintByteSpacing("  ");
            printf("%c", index >= start_poison_index && index <= end_poison_index ? 'x' : ' ');
        }
        printf("\n");

        // NOTE: Print the poison-ed bytes =========================================================

        PrintPoisonedBytes(4 /*step*/, array, sizeof(array));
    }

    printf("\n\nTLDR:\n"
           "  __asan_poison_memory_region(ptr, size)\n"
           "  Poisons the byte region [ptr, AlignDown(ptr+size, 8))\n\n"
           "  __asan_unpoison_memory_region(ptr, size)\n"
           "  Unpoisons the byte region [AlignDown(ptr, 8), ptr+size)\n\n"
          );
    return 0;
}