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

#include <catch2/catch_test_macros.hpp>

#include <mach_detours.h>

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include "test_func.h"

static int testFunctionDetourCounter = 0;
int testFunctionDetour()
{
    testFunctionDetourCounter++;
    return 94;
}

TEST_CASE( "Overriding custom function in dylib" )
{
    int (*realTestFunction)() = testFunction;
    testFunctionCounter = 0;
    testFunctionDetourCounter = 0;

    REQUIRE( testFunction() == 42 );
    REQUIRE( testFunctionCounter == 1 );
    REQUIRE( testFunctionDetourCounter == 0 );

    CHECK( detour_transaction_begin() == err_none );
    CHECK( detour_attach(reinterpret_cast<detour_func_t*>(&realTestFunction), reinterpret_cast<detour_func_t>(testFunctionDetour)) == err_none );
    CHECK( detour_transaction_commit() == err_none );

    REQUIRE( realTestFunction != testFunction );

    REQUIRE( testFunctionCounter == 1 );
    REQUIRE( testFunctionDetourCounter == 0 );
    REQUIRE( testFunction() == 94 );
    REQUIRE( testFunctionCounter == 1 );
    REQUIRE( testFunctionDetourCounter == 1 );
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

static int localFunctionCounter = 0;
int localFunction()
{
    localFunctionCounter++;
    return 67;
}
static int localFunctionDetourCounter = 0;
int localFunctionDetour()
{
    localFunctionDetourCounter++;
    return 12;
}

TEST_CASE( "Overriding local function" )
{
    int (*realFunction)() = localFunction;
    localFunctionCounter = 0;
    localFunctionDetourCounter = 0;

    REQUIRE( localFunction() == 67 );
    REQUIRE( localFunctionCounter == 1 );
    REQUIRE( localFunctionDetourCounter == 0 );

    CHECK( detour_transaction_begin() == err_none );
    // Overriding a local function requires using detour_attach_and_commit instead of calling attach and commit individually.
    // Otherwise when we return from attach (and before commit), the code page with the local function is marked as read+write (but _not_ executable).
    // There is a good chance the code we return to (in this case the test function) is on the same memory page and can therefore not be executed (until we call commit).
    CHECK( detour_attach_and_commit(reinterpret_cast<detour_func_t*>(&realFunction), reinterpret_cast<detour_func_t>(localFunctionDetour)) == err_none );

    REQUIRE( realFunction != localFunction );

    REQUIRE( localFunctionCounter == 1 );
    REQUIRE( localFunctionDetourCounter == 0 );
    REQUIRE( localFunction() == 12 );
    REQUIRE( localFunctionCounter == 1 );
    REQUIRE( localFunctionDetourCounter == 1 );
}