/*
 * Copyright (C) 2025 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Defining private as public to create 'SkRegion::RunHead'
#define private public

#include <sys/mman.h>
#include <unistd.h>

#include "SkRegion.h"
#include "binary_loader.h"
#include "common.h"
#include "memutils.h"
#include "src/core/SkRegionPriv.h"

typedef bool (*oper_func)(const SkRegion&, const SkRegion&, SkRegion::Op, SkRegion*);
struct sigaction new_action, old_action;
static void* (*real_realloc_func)(void*, size_t) = nullptr;
bool isTestInProgress = false;
char* startPtr = nullptr;
int kElement = 0;
const int kRunArrayOffset = 7 /* Offset of RunArray in RunHead class */;

void sig_handler(int signum, siginfo_t* info, void* context) {
    if (!isTestInProgress) {
        (*old_action.sa_sigaction)(signum, info, context);
        return;
    }

    // Null check on startPtr.
    if (startPtr == nullptr) {
        exit(EXIT_FAILURE);
    }

    // Free the allocated buffer.
    ENABLE_MEM_ACCESS(startPtr, getpagesize());
    bool isExpectedAddress =
            (((uintptr_t)info->si_addr ^ (uintptr_t)startPtr) & (uintptr_t)info->si_addr) == 0;
    free(startPtr);
    startPtr = nullptr;

    // Check if signal is coming from expected address.
    if (isExpectedAddress) {
        (*old_action.sa_sigaction)(signum, info, context);
        exit(EXIT_VULNERABLE);
    }

    // Assumption failure for unexpected signals.
    exit(EXIT_FAILURE);
}

void* realloc(void* ptr, size_t size) {
    real_realloc_func = (void* (*)(void*, size_t))dlsym(RTLD_NEXT, "realloc");
    FAIL_CHECK(real_realloc_func);

    // Call real realloc() when test is not in progress.
    if (!isTestInProgress) {
        return real_realloc_func(ptr, size);
    }

    // Compute the result of addition operation on vulnerable function's input value.
    int vulnerableFuncInput = SK_MaxS32 - 2 * kRunArrayOffset;
    int additionResult = (int)(vulnerableFuncInput + (vulnerableFuncInput >> 1));

    // Check if 'size' contains expected value.
    if (size == ((size_t)additionResult * 4)) {
        // Create a guarded buffer.
        size_t page_size = getpagesize();
        startPtr = (char*)memalign(page_size, page_size);
        memset(startPtr, 0, page_size);
        DISABLE_MEM_ACCESS(startPtr, page_size);
        return startPtr;
    }
    return real_realloc_func(ptr, size);
}

// This function is referred from 'RunHead* Alloc()'.
SkRegion::RunHead* createRunHead(int count) {
    // Create RunHead pointer
    const int64_t size = sk_64_mul(count, 4 /* Size of RunType */) + sizeof(SkRegion::RunHead);
    SkRegion::RunHead* head = (SkRegion::RunHead*)sk_malloc_throw(size);

    // Configure the values.
    int* tmpPtr = (int*)(head);
    *(tmpPtr + kRunArrayOffset + 3 /* Second element of Run array */) = kElement++;
    *(tmpPtr + count - 1) = SkRegion_kRunTypeSentinel;
    *(tmpPtr + count - 2) = SkRegion_kRunTypeSentinel;
    return head;
}

int main(int /* argc */, char* argv[]) {
    // Setup signal handler.
    sigemptyset(&new_action.sa_mask);
    sigaddset(&new_action.sa_mask, SIGSEGV); // Add SIGSEGV to the signal mask
    sigaddset(&new_action.sa_mask, SIGBUS);  // Add SIGBUS to the signal mask
    new_action.sa_flags = SA_SIGINFO;
    new_action.sa_sigaction = sig_handler;
    sigaction(SIGSEGV, &new_action, &old_action);
    sigaction(SIGBUS, &new_action, &old_action);

    // Get the path to the shared library and offset from command-line arguments
    const char* libPath = argv[1];
    const uintptr_t functionOffset = strtoul(argv[2], nullptr, 0);

    // Get function address of 'oper()' from loaded library
    BinaryLoader binaryLoader(libPath);
    const uintptr_t functionAddress = binaryLoader.getFunctionAddress(functionOffset);
    FAIL_CHECK(functionAddress);

    // Create function pointer to 'oper()'
    oper_func oper_ptr = (oper_func)functionAddress;

    // Configure SkRegion's
    SkRegion region1, region2;
    region1.fRunHead = createRunHead(SK_MaxS32 / 2);
    region2.fRunHead = createRunHead(SK_MaxS32 / 2);

    // Call the vulnerable function Oper(). Without the fix, realloc() is called with overflowed
    // value. realloc() is overloaded to return a guarded buffer for overflowed value. Later,
    // 'memcpy' reads the guarded buffer, and a SIGSEGV is raised. With the fix, the integer
    // overflow does not occur, and realloc() returns a buffer of the expected size.
    isTestInProgress = true;
    SkRegion::Op operation = SkRegion::kXOR_Op;
    oper_ptr(region1, region2, operation, nullptr);
    return EXIT_SUCCESS;
}