/*++

Module Name:

    image_protection_manager.c

Abstract:

    Exact-build MmChangeImageProtection resolver and managed test-page MDLs.

Environment:

    Kernel mode, PASSIVE_LEVEL control path only.

--*/

#include "image_protection_manager.h"
#include <ntimage.h>
#include "src/platform/kernel_module_identity.h"
#include "src/platform/pool_compat.h"

NTSYSAPI
PVOID
NTAPI
RtlPcToFileHeader(
    _In_ PVOID PcValue,
    _Outptr_ PVOID* BaseOfImage
    );

#define KSW_RXPF_BACKUP_TAG 'bPxR'
#define KSW_RXPF_CONTENT_TAG 'cPxR'
#define KSW_RXPF_EXPECTED_NT_BUILD 26220UL
#define KSW_RXPF_EXPECTED_TIMESTAMP 0xB292FCA9UL
#define KSW_RXPF_EXPECTED_IMAGE_SIZE 0x01450000UL
#define KSW_RXPF_EXPECTED_CHECKSUM 0x00C7D4D4UL
#define KSW_RXPF_MM_CHANGE_RVA 0x00A3B300UL
#define KSW_RXPF_MM_CHANGE_OPERATION_RX 1UL
#define KSW_RXPF_MM_CHANGE_OPERATION_READONLY_NX 2UL
#define KSW_RXPF_RSDS_SIGNATURE 0x53445352UL

typedef NTSTATUS
(NTAPI* KSW_RXPF_MM_CHANGE_IMAGE_PROTECTION)(
    _In_ PMDL Mdl,
    _In_ PVOID BaseAddress,
    _In_ ULONG NumberOfBytes,
    _In_ ULONG Operation
    );

#pragma pack(push, 1)
typedef struct _KSW_RXPF_RSDS_HEADER
{
    ULONG Signature;
    GUID Guid;
    ULONG Age;
} KSW_RXPF_RSDS_HEADER, *PKSW_RXPF_RSDS_HEADER;
#pragma pack(pop)

typedef struct _KSW_RXPF_BUILD_PROFILE
{
    ULONG NtBuildNumber;
    ULONG ImageTimeDateStamp;
    ULONG ImageSize;
    ULONG ImageCheckSum;
    ULONG FunctionRva;
    ULONG FunctionParameterCount;
    ULONG OperationRx;
    ULONG OperationReadOnlyNx;
    ULONG PdbAge;
    const GUID* PdbGuid;
    const UCHAR* Signature;
    const UCHAR* SignatureMask;
    ULONG SignatureLength;
} KSW_RXPF_BUILD_PROFILE, *PKSW_RXPF_BUILD_PROFILE;

typedef struct _KSW_RXPF_IMAGE_PROTECTION_STATE
{
    PDRIVER_OBJECT DriverObject;
    PVOID KernelBase;
    ULONG NtBuildNumber;
    ULONG ImageTimeDateStamp;
    ULONG ImageSize;
    ULONG ImageCheckSum;
    ULONG BuildStatus;
    GUID PdbGuid;
    ULONG PdbAge;
    KSW_RXPF_MM_CHANGE_IMAGE_PROTECTION MmChangeImageProtection;
    const KSW_RXPF_BUILD_PROFILE* ActiveProfile;
    volatile LONG Initialized;
    volatile LONG BuildSupported;
    NTSTATUS LastStatus;
} KSW_RXPF_IMAGE_PROTECTION_STATE;

static KSW_RXPF_IMAGE_PROTECTION_STATE g_KswRxpfImageProtection;

static const GUID g_KswRxpfExpectedPdbGuid = {
    0x3E5A9A8BUL,
    0x6B78U,
    0x281FU,
    { 0x3BU, 0xE2U, 0x15U, 0x01U, 0x10U, 0x32U, 0x52U, 0x15U }
};

static const UCHAR g_KswRxpfMmChangeSignature[64] = {
    0x48U, 0x8BU, 0xC4U, 0x48U, 0x89U, 0x58U, 0x08U, 0x48U,
    0x89U, 0x68U, 0x10U, 0x48U, 0x89U, 0x70U, 0x18U, 0x44U,
    0x89U, 0x48U, 0x20U, 0x57U, 0x41U, 0x54U, 0x41U, 0x55U,
    0x41U, 0x56U, 0x41U, 0x57U, 0x48U, 0x83U, 0xECU, 0x30U,
    0x83U, 0x60U, 0xC8U, 0x00U, 0x49U, 0x8BU, 0xD8U, 0x41U,
    0x8DU, 0x41U, 0xFFU, 0x4CU, 0x8BU, 0xE2U, 0x48U, 0x8BU,
    0xF9U, 0x83U, 0xF8U, 0x01U, 0x0FU, 0x87U, 0x41U, 0x02U,
    0x00U, 0x00U, 0x44U, 0x8BU, 0xEBU, 0x49U, 0x3BU, 0xDDU
};

static const UCHAR g_KswRxpfMmChangeSignatureMask[64] = {
    0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU,
    0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU,
    0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU,
    0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU,
    0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU,
    0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU,
    0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU,
    0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU, 0xFFU
};

static const KSW_RXPF_BUILD_PROFILE g_KswRxpfBuildProfiles[] = {
    {
        KSW_RXPF_EXPECTED_NT_BUILD,
        KSW_RXPF_EXPECTED_TIMESTAMP,
        KSW_RXPF_EXPECTED_IMAGE_SIZE,
        KSW_RXPF_EXPECTED_CHECKSUM,
        KSW_RXPF_MM_CHANGE_RVA,
        4UL,
        KSW_RXPF_MM_CHANGE_OPERATION_RX,
        KSW_RXPF_MM_CHANGE_OPERATION_READONLY_NX,
        1UL,
        &g_KswRxpfExpectedPdbGuid,
        g_KswRxpfMmChangeSignature,
        g_KswRxpfMmChangeSignatureMask,
        RTL_NUMBER_OF(g_KswRxpfMmChangeSignature)
    }
};

C_ASSERT(RTL_NUMBER_OF(g_KswRxpfBuildProfiles) == 1U);
C_ASSERT(sizeof(g_KswRxpfMmChangeSignature) == 64U);

/*
 * The linker gives this unique executable section its own page.  No handler or
 * control-path code resides here, so leaving it NX until image unload cannot
 * strand an unload callback.  Bytes encode MOV RAX,12345678h; ADD RAX,1; RET.
 */
#pragma section(".rxpftst", read, execute)
__declspec(allocate(".rxpftst"))
__declspec(align(PAGE_SIZE))
const UCHAR g_KswRxpfSelfImageTestPage[PAGE_SIZE] = {
    0x48U, 0xB8U, 0x78U, 0x56U, 0x34U, 0x12U, 0x00U, 0x00U,
    0x00U, 0x00U, 0x48U, 0x83U, 0xC0U, 0x01U, 0xC3U
};

static BOOLEAN
KswRxpfCanonicalKernelPage(
    _In_ ULONGLONG Address
    )
{
    ULONGLONG high = Address >> 48;

    /* Require canonical sign extension, kernel range, and page alignment. */
    return high == 0xFFFFULL &&
        Address >= (ULONGLONG)(ULONG_PTR)MmSystemRangeStart &&
        (Address & (PAGE_SIZE - 1ULL)) == 0ULL;
}

static BOOLEAN
KswRxpfValidateExecutableSection(
    _In_ PVOID ImageBase,
    _In_ ULONG Rva
    )
{
    BOOLEAN valid = FALSE;

    /* Parse the loaded PE under exception protection before trusting sections. */
    __try {
        PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)ImageBase;
        PIMAGE_NT_HEADERS64 nt = NULL;
        PIMAGE_SECTION_HEADER section = NULL;
        USHORT index = 0U;

        if (dos->e_magic != IMAGE_DOS_SIGNATURE ||
            dos->e_lfanew <= 0 ||
            (ULONG)dos->e_lfanew > 0x1000UL) {
            return FALSE;
        }
        nt = (PIMAGE_NT_HEADERS64)((PUCHAR)ImageBase + dos->e_lfanew);
        if (nt->Signature != IMAGE_NT_SIGNATURE ||
            nt->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
            return FALSE;
        }
        section = IMAGE_FIRST_SECTION(nt);
        for (index = 0U; index < nt->FileHeader.NumberOfSections; ++index) {
            ULONG sectionSize = max(
                section[index].Misc.VirtualSize,
                section[index].SizeOfRawData);
            ULONG sectionEnd = section[index].VirtualAddress + sectionSize;

            if (sectionEnd < section[index].VirtualAddress) {
                continue;
            }
            if (Rva >= section[index].VirtualAddress && Rva < sectionEnd) {
                valid =
                    (section[index].Characteristics &
                        IMAGE_SCN_MEM_EXECUTE) != 0UL &&
                    (section[index].Characteristics &
                        IMAGE_SCN_MEM_DISCARDABLE) == 0UL;
                break;
            }
        }
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        valid = FALSE;
    }
    return valid;
}

static BOOLEAN
KswRxpfReadRsds(
    _In_ PVOID ImageBase,
    _Out_ GUID* GuidOut,
    _Out_ ULONG* AgeOut
    )
{
    ULONG matchingEntries = 0UL;

    /* Accept exactly one well-formed RSDS record from the loaded image. */
    __try {
        PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)ImageBase;
        PIMAGE_NT_HEADERS64 nt = NULL;
        IMAGE_DATA_DIRECTORY directory;
        PIMAGE_DEBUG_DIRECTORY entries = NULL;
        ULONG entryCount = 0UL;
        ULONG index = 0UL;

        if (dos->e_magic != IMAGE_DOS_SIGNATURE ||
            dos->e_lfanew <= 0 ||
            (ULONG)dos->e_lfanew > 0x1000UL) {
            return FALSE;
        }
        nt = (PIMAGE_NT_HEADERS64)((PUCHAR)ImageBase + dos->e_lfanew);
        if (nt->Signature != IMAGE_NT_SIGNATURE ||
            nt->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC ||
            nt->OptionalHeader.NumberOfRvaAndSizes <=
                IMAGE_DIRECTORY_ENTRY_DEBUG) {
            return FALSE;
        }
        directory = nt->OptionalHeader.DataDirectory[
            IMAGE_DIRECTORY_ENTRY_DEBUG];
        if (directory.VirtualAddress == 0UL ||
            directory.Size < sizeof(IMAGE_DEBUG_DIRECTORY) ||
            directory.VirtualAddress + directory.Size <
                directory.VirtualAddress ||
            directory.VirtualAddress + directory.Size >
                nt->OptionalHeader.SizeOfImage) {
            return FALSE;
        }
        entries = (PIMAGE_DEBUG_DIRECTORY)(
            (PUCHAR)ImageBase + directory.VirtualAddress);
        entryCount = directory.Size / sizeof(IMAGE_DEBUG_DIRECTORY);
        for (index = 0UL; index < entryCount; ++index) {
            KSW_RXPF_RSDS_HEADER rsds;

            if (entries[index].Type != IMAGE_DEBUG_TYPE_CODEVIEW ||
                entries[index].AddressOfRawData == 0UL ||
                entries[index].SizeOfData < sizeof(rsds) ||
                entries[index].AddressOfRawData + sizeof(rsds) <
                    entries[index].AddressOfRawData ||
                entries[index].AddressOfRawData + sizeof(rsds) >
                    nt->OptionalHeader.SizeOfImage) {
                continue;
            }
            RtlCopyMemory(
                &rsds,
                (PUCHAR)ImageBase + entries[index].AddressOfRawData,
                sizeof(rsds));
            if (rsds.Signature != KSW_RXPF_RSDS_SIGNATURE) {
                continue;
            }
            *GuidOut = rsds.Guid;
            *AgeOut = rsds.Age;
            matchingEntries += 1UL;
        }
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        return FALSE;
    }
    return matchingEntries == 1UL;
}

static BOOLEAN
KswRxpfSignatureMatches(
    _In_ const KSW_RXPF_BUILD_PROFILE* Profile,
    _In_reads_(Profile->SignatureLength) const UCHAR* Address
    )
{
    ULONG index = 0UL;

    /* Compare every masked byte from the exact PDB/disassembly profile. */
    if (Profile == NULL || Address == NULL ||
        Profile->SignatureLength == 0UL ||
        Profile->SignatureLength > 64UL) {
        return FALSE;
    }
    __try {
        for (index = 0UL; index < Profile->SignatureLength; ++index) {
            if ((Address[index] & Profile->SignatureMask[index]) !=
                (Profile->Signature[index] &
                    Profile->SignatureMask[index])) {
                return FALSE;
            }
        }
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        return FALSE;
    }
    return TRUE;
}

static VOID
KswRxpfResolveExactBuild(
    VOID
    )
{
    static const KSW_KERNEL_MODULE_NAME_MATCH matches[] = {
        { "ntoskrnl.exe", 0UL },
        { "ntkrnlmp.exe", 0UL },
        { "ntkrnlpa.exe", 0UL },
        { "ntkrpamp.exe", 0UL }
    };
    KSW_DYN_MODULE_IDENTITY_PACKET identity;
    RTL_OSVERSIONINFOW version;
    GUID pdbGuid;
    ULONG pdbAge = 0UL;
    PVOID functionAddress = NULL;
    NTSTATUS status = STATUS_SUCCESS;
    const KSW_RXPF_BUILD_PROFILE* profile = NULL;
    ULONG profileIndex = 0UL;

    /* Start from an unsupported state and promote only after every check. */
    g_KswRxpfImageProtection.BuildStatus =
        KSWORD_ARK_RXPF_BUILD_STATUS_UNINITIALIZED;
    g_KswRxpfImageProtection.LastStatus = STATUS_NOT_SUPPORTED;
    RtlZeroMemory(&identity, sizeof(identity));
    RtlZeroMemory(&version, sizeof(version));
    RtlZeroMemory(&pdbGuid, sizeof(pdbGuid));
    version.dwOSVersionInfoSize = sizeof(version);
    status = RtlGetVersion(&version);
    if (!NT_SUCCESS(status)) {
        g_KswRxpfImageProtection.BuildStatus =
            KSWORD_ARK_RXPF_BUILD_STATUS_OS_BUILD_MISMATCH;
        g_KswRxpfImageProtection.LastStatus = status;
        return;
    }
    g_KswRxpfImageProtection.NtBuildNumber = version.dwBuildNumber;
    for (profileIndex = 0UL;
         profileIndex < RTL_NUMBER_OF(g_KswRxpfBuildProfiles);
         ++profileIndex) {
        if (g_KswRxpfBuildProfiles[profileIndex].NtBuildNumber ==
            version.dwBuildNumber) {
            profile = &g_KswRxpfBuildProfiles[profileIndex];
            break;
        }
    }
    if (profile == NULL) {
        g_KswRxpfImageProtection.BuildStatus =
            KSWORD_ARK_RXPF_BUILD_STATUS_OS_BUILD_MISMATCH;
        return;
    }

    /* Resolve the loaded nt image instead of trusting an on-disk path. */
    status = KswordARKQueryKernelModuleIdentity(
        matches,
        RTL_NUMBER_OF(matches),
        &identity);
    if (!NT_SUCCESS(status) || identity.present == 0UL) {
        g_KswRxpfImageProtection.BuildStatus =
            KSWORD_ARK_RXPF_BUILD_STATUS_IMAGE_MISMATCH;
        g_KswRxpfImageProtection.LastStatus = status;
        return;
    }
    g_KswRxpfImageProtection.KernelBase =
        (PVOID)(ULONG_PTR)identity.imageBase;
    g_KswRxpfImageProtection.ImageTimeDateStamp = identity.timeDateStamp;
    g_KswRxpfImageProtection.ImageSize = identity.sizeOfImage;
    if (identity.machine != IMAGE_FILE_MACHINE_AMD64 ||
        identity.timeDateStamp != profile->ImageTimeDateStamp ||
        identity.sizeOfImage != profile->ImageSize) {
        g_KswRxpfImageProtection.BuildStatus =
            KSWORD_ARK_RXPF_BUILD_STATUS_IMAGE_MISMATCH;
        return;
    }

    /* Read the checksum from the same loaded PE headers. */
    __try {
        PIMAGE_DOS_HEADER dos =
            (PIMAGE_DOS_HEADER)g_KswRxpfImageProtection.KernelBase;
        PIMAGE_NT_HEADERS64 nt = (PIMAGE_NT_HEADERS64)(
            (PUCHAR)g_KswRxpfImageProtection.KernelBase + dos->e_lfanew);

        g_KswRxpfImageProtection.ImageCheckSum =
            nt->OptionalHeader.CheckSum;
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        g_KswRxpfImageProtection.BuildStatus =
            KSWORD_ARK_RXPF_BUILD_STATUS_IMAGE_MISMATCH;
        g_KswRxpfImageProtection.LastStatus = GetExceptionCode();
        return;
    }
    if (g_KswRxpfImageProtection.ImageCheckSum !=
        profile->ImageCheckSum) {
        g_KswRxpfImageProtection.BuildStatus =
            KSWORD_ARK_RXPF_BUILD_STATUS_IMAGE_MISMATCH;
        return;
    }

    /* Bind the profile to the exact Microsoft RSDS GUID and age. */
    if (!KswRxpfReadRsds(
            g_KswRxpfImageProtection.KernelBase,
            &pdbGuid,
            &pdbAge) ||
        !RtlEqualMemory(
            &pdbGuid,
            profile->PdbGuid,
            sizeof(pdbGuid)) ||
        pdbAge != profile->PdbAge) {
        g_KswRxpfImageProtection.BuildStatus =
            KSWORD_ARK_RXPF_BUILD_STATUS_RSDS_MISMATCH;
        return;
    }
    g_KswRxpfImageProtection.PdbGuid = pdbGuid;
    g_KswRxpfImageProtection.PdbAge = pdbAge;

    /* Require the PDB RVA to land in a live executable PE section. */
    if (!KswRxpfValidateExecutableSection(
            g_KswRxpfImageProtection.KernelBase,
            profile->FunctionRva)) {
        g_KswRxpfImageProtection.BuildStatus =
            KSWORD_ARK_RXPF_BUILD_STATUS_SECTION_INVALID;
        return;
    }
    functionAddress =
        (PUCHAR)g_KswRxpfImageProtection.KernelBase +
        profile->FunctionRva;
    if (!KswRxpfSignatureMatches(
            profile,
            (const UCHAR*)functionAddress)) {
        g_KswRxpfImageProtection.BuildStatus =
            KSWORD_ARK_RXPF_BUILD_STATUS_SIGNATURE_MISMATCH;
        return;
    }

    /* Publish the internal call target only after every independent check. */
    g_KswRxpfImageProtection.ActiveProfile = profile;
    g_KswRxpfImageProtection.MmChangeImageProtection =
        (KSW_RXPF_MM_CHANGE_IMAGE_PROTECTION)functionAddress;
    g_KswRxpfImageProtection.BuildStatus =
        KSWORD_ARK_RXPF_BUILD_STATUS_SUPPORTED;
    g_KswRxpfImageProtection.LastStatus = STATUS_SUCCESS;
    InterlockedExchange(&g_KswRxpfImageProtection.BuildSupported, 1);
}

NTSTATUS
KswRxpfImageProtectionInitialize(
    _In_ PDRIVER_OBJECT DriverObject
    )
{
    /* Record driver identity and resolve but never call the internal routine. */
    if (DriverObject == NULL) {
        return STATUS_INVALID_PARAMETER;
    }
    if (InterlockedCompareExchange(
            &g_KswRxpfImageProtection.Initialized,
            1,
            1) != 0) {
        return STATUS_SUCCESS;
    }
    RtlZeroMemory(
        &g_KswRxpfImageProtection,
        sizeof(g_KswRxpfImageProtection));
    g_KswRxpfImageProtection.DriverObject = DriverObject;
    KswRxpfResolveExactBuild();
    InterlockedExchange(&g_KswRxpfImageProtection.Initialized, 1);
    return STATUS_SUCCESS;
}

VOID
KswRxpfImageProtectionUninitialize(
    VOID
    )
{
    /* Remove the private function pointer before driver-owned state disappears. */
    InterlockedExchange(&g_KswRxpfImageProtection.BuildSupported, 0);
    InterlockedExchange(&g_KswRxpfImageProtection.Initialized, 0);
    KeMemoryBarrier();
    RtlZeroMemory(
        &g_KswRxpfImageProtection,
        sizeof(g_KswRxpfImageProtection));
}

VOID
KswRxpfImageProtectionQuerySupport(
    _Out_ KSWORD_ARK_RXPF_QUERY_SUPPORT_RESPONSE* Response
    )
{
    ULONG supportFlags = 0UL;
    const KSW_RXPF_BUILD_PROFILE* profile =
        g_KswRxpfImageProtection.ActiveProfile;

    if (profile == NULL) {
        profile = &g_KswRxpfBuildProfiles[0];
    }

    /* Return the exact build profile and signature used by runtime gating. */
    RtlZeroMemory(Response, sizeof(*Response));
    Response->version = KSWORD_ARK_RXPF_PROTOCOL_VERSION;
    Response->size = sizeof(*Response);
    if (InterlockedCompareExchange(
            &g_KswRxpfImageProtection.Initialized,
            1,
            1) != 0) {
        supportFlags |= KSWORD_ARK_RXPF_SUPPORT_INITIALIZED |
            KSWORD_ARK_RXPF_SUPPORT_ALLOCATED_TEST_PAGE |
            KSWORD_ARK_RXPF_SUPPORT_IDT_SHADOW |
            KSWORD_ARK_RXPF_SUPPORT_EMULATOR;
    }
    if (KswRxpfImageProtectionBuildSupported()) {
        supportFlags |= KSWORD_ARK_RXPF_SUPPORT_BUILD_MATCH |
            KSWORD_ARK_RXPF_SUPPORT_ABI_VERIFIED |
            KSWORD_ARK_RXPF_SUPPORT_SELF_IMAGE_TEST_PAGE;
    }
#if KSW_RXPF_ENABLE_EXTERNAL_IMAGE_TARGETS
    supportFlags |= KSWORD_ARK_RXPF_SUPPORT_EXTERNAL_IMAGE_COMPILED;
#endif
    Response->supportFlags = supportFlags;
    Response->buildStatus = g_KswRxpfImageProtection.BuildStatus;
    Response->ntBuildNumber = g_KswRxpfImageProtection.NtBuildNumber;
    Response->imageTimeDateStamp =
        g_KswRxpfImageProtection.ImageTimeDateStamp;
    Response->imageSize = g_KswRxpfImageProtection.ImageSize;
    Response->imageCheckSum = g_KswRxpfImageProtection.ImageCheckSum;
    Response->functionRva = profile->FunctionRva;
    Response->functionOperationRx = profile->OperationRx;
    Response->functionOperationReadOnlyNx =
        profile->OperationReadOnlyNx;
    Response->functionParameterCount = profile->FunctionParameterCount;
    Response->pdbAge = g_KswRxpfImageProtection.PdbAge;
    Response->lastStatus = g_KswRxpfImageProtection.LastStatus;
    RtlCopyMemory(
        Response->pdbGuid,
        &g_KswRxpfImageProtection.PdbGuid,
        sizeof(Response->pdbGuid));
    RtlCopyMemory(
        Response->signature,
        profile->Signature,
        profile->SignatureLength);
    RtlCopyMemory(
        Response->signatureMask,
        profile->SignatureMask,
        profile->SignatureLength);
}

BOOLEAN
KswRxpfImageProtectionBuildSupported(
    VOID
    )
{
    /* The function pointer and exact identity are published as one gate. */
    return InterlockedCompareExchange(
        &g_KswRxpfImageProtection.BuildSupported,
        1,
        1) != 0 &&
        g_KswRxpfImageProtection.MmChangeImageProtection != NULL;
}

PVOID
KswRxpfImageProtectionSelfTestPage(
    VOID
    )
{
    /* The dedicated image section is page-aligned by declaration and linker. */
    return (PVOID)(ULONG_PTR)g_KswRxpfSelfImageTestPage;
}

static NTSTATUS
KswRxpfCreateAllocatedTestPage(
    _In_ ULONG Flags,
    _Out_ KSW_RXPF_PAGE_RECORD* RecordSource
    )
{
    PHYSICAL_ADDRESS lowAddress;
    PHYSICAL_ADDRESS highAddress;
    PHYSICAL_ADDRESS skipBytes;
    PMDL mdl = NULL;
    PVOID writableMapping = NULL;
    PVOID executeMapping = NULL;
    NTSTATUS status = STATUS_SUCCESS;

    /* Allocate one physical page described by an owned MDL. */
    lowAddress.QuadPart = 0LL;
    highAddress.QuadPart = MAXLONGLONG;
    skipBytes.QuadPart = 0LL;
    mdl = MmAllocatePagesForMdlEx(
        lowAddress,
        highAddress,
        skipBytes,
        PAGE_SIZE,
        MmCached,
        MM_ALLOCATE_FULLY_REQUIRED);
    if (mdl == NULL || MmGetMdlByteCount(mdl) != PAGE_SIZE) {
        if (mdl != NULL) {
            MmFreePagesFromMdl(mdl);
            ExFreePool(mdl);
        }
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    /* Seed bytes through a writable/NX mapping before exposing an RX mapping. */
    writableMapping = MmMapLockedPagesSpecifyCache(
        mdl,
        KernelMode,
        MmCached,
        NULL,
        FALSE,
        (MM_PAGE_PRIORITY)(NormalPagePriority | MdlMappingNoExecute));
    if (writableMapping == NULL) {
        MmFreePagesFromMdl(mdl);
        ExFreePool(mdl);
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    RtlCopyMemory(
        writableMapping,
        g_KswRxpfSelfImageTestPage,
        PAGE_SIZE);
    MmUnmapLockedPages(writableMapping, mdl);

    /* Map the same locked page read/execute for the initial registered state. */
    executeMapping = MmMapLockedPagesSpecifyCache(
        mdl,
        KernelMode,
        MmCached,
        NULL,
        FALSE,
        (MM_PAGE_PRIORITY)(NormalPagePriority | MdlMappingNoWrite));
    if (executeMapping == NULL) {
        MmFreePagesFromMdl(mdl);
        ExFreePool(mdl);
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    status = MmProtectMdlSystemAddress(mdl, PAGE_EXECUTE_READ);
    if (!NT_SUCCESS(status)) {
        MmUnmapLockedPages(executeMapping, mdl);
        MmFreePagesFromMdl(mdl);
        ExFreePool(mdl);
        return status;
    }
    if (!KswRxpfCanonicalKernelPage(
            (ULONGLONG)(ULONG_PTR)executeMapping)) {
        MmUnmapLockedPages(executeMapping, mdl);
        MmFreePagesFromMdl(mdl);
        ExFreePool(mdl);
        return STATUS_CONFLICTING_ADDRESSES;
    }

    /* Fill an unpublished source record for the fixed state table. */
    RtlZeroMemory(RecordSource, sizeof(*RecordSource));
    RecordSource->TargetKind = KSWORD_ARK_RXPF_TARGET_ALLOCATED_TEST;
    RecordSource->Flags = Flags;
    RecordSource->PageBase = (LONG64)(ULONG_PTR)executeMapping;
    RecordSource->OriginalMapping = executeMapping;
    RecordSource->OriginalProtection = PAGE_EXECUTE_READ;
    RecordSource->CurrentProtection = PAGE_EXECUTE_READ;
    RecordSource->WritableAliasProtection = 0UL;
    RecordSource->Pfn = (ULONGLONG)MmGetMdlPfnArray(mdl)[0];
    RecordSource->Mdl = mdl;
    RecordSource->OwnsMdlPages = TRUE;
    RecordSource->MappingIsAlias = FALSE;
    RecordSource->LastStatus = STATUS_SUCCESS;
    return STATUS_SUCCESS;
}

static NTSTATUS
KswRxpfCreateSelfImageTestPage(
    _In_ ULONGLONG RequestedAddress,
    _In_ ULONG Flags,
    _Out_ KSW_RXPF_PAGE_RECORD* RecordSource
    )
{
    PVOID page = KswRxpfImageProtectionSelfTestPage();
    PVOID ownerImageBase = NULL;
    PMDL mdl = NULL;
    BOOLEAN pagesLocked = FALSE;

    /* The internal ABI is unavailable on every non-exact Windows build. */
    if (!KswRxpfImageProtectionBuildSupported()) {
        return STATUS_NOT_SUPPORTED;
    }
    if (((ULONGLONG)(ULONG_PTR)page & (PAGE_SIZE - 1ULL)) != 0ULL ||
        (RequestedAddress != 0ULL &&
            RequestedAddress != (ULONGLONG)(ULONG_PTR)page)) {
        return STATUS_INVALID_ADDRESS;
    }
    if (RtlPcToFileHeader(page, &ownerImageBase) == NULL ||
        ownerImageBase == NULL) {
        return STATUS_NOT_FOUND;
    }

    /* Probe-and-lock creates the exact low-bit MDL flag state required by nt. */
    mdl = IoAllocateMdl(page, PAGE_SIZE, FALSE, FALSE, NULL);
    if (mdl == NULL) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    __try {
        MmProbeAndLockPages(mdl, KernelMode, IoReadAccess);
        pagesLocked = TRUE;
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        NTSTATUS exceptionStatus = GetExceptionCode();

        IoFreeMdl(mdl);
        return exceptionStatus;
    }
    if (!pagesLocked ||
        (mdl->MdlFlags &
            (MDL_MAPPED_TO_SYSTEM_VA |
             MDL_PAGES_LOCKED |
             MDL_SOURCE_IS_NONPAGED_POOL)) != MDL_PAGES_LOCKED ||
        MmGetMdlByteOffset(mdl) != 0UL ||
        MmGetMdlByteCount(mdl) != PAGE_SIZE ||
        MmGetMdlVirtualAddress(mdl) != page) {
        if (pagesLocked) {
            MmUnlockPages(mdl);
        }
        IoFreeMdl(mdl);
        return STATUS_INVALID_PARAMETER;
    }

    /* Fill the record while the original mapping is still RX and untouched. */
    RtlZeroMemory(RecordSource, sizeof(*RecordSource));
    RecordSource->TargetKind = KSWORD_ARK_RXPF_TARGET_SELF_IMAGE_TEST;
    RecordSource->Flags = Flags;
    RecordSource->PageBase = (LONG64)(ULONG_PTR)page;
    RecordSource->OriginalMapping = page;
    RecordSource->OriginalProtection = PAGE_EXECUTE_READ;
    RecordSource->CurrentProtection = PAGE_EXECUTE_READ;
    RecordSource->WritableAliasProtection = 0UL;
    RecordSource->Pfn = (ULONGLONG)MmGetMdlPfnArray(mdl)[0];
    RecordSource->OwnerImageBase =
        (ULONGLONG)(ULONG_PTR)ownerImageBase;
    RecordSource->Mdl = mdl;
    RecordSource->PagesLockedByProbe = TRUE;
    RecordSource->LastStatus = STATUS_SUCCESS;
    return STATUS_SUCCESS;
}

NTSTATUS
KswRxpfImageProtectionCreateRecord(
    _In_ ULONG TargetKind,
    _In_ ULONGLONG RequestedAddress,
    _In_ ULONG Flags,
    _Out_ KSW_RXPF_PAGE_RECORD* RecordSource
    )
{
    NTSTATUS status = STATUS_NOT_SUPPORTED;

    /* Only documented test kinds are accepted by the default build. */
    if (RecordSource == NULL ||
        (Flags & ~KSWORD_ARK_RXPF_FLAG_CAPTURE_BACKUP) != 0UL) {
        return STATUS_INVALID_PARAMETER;
    }
    RtlZeroMemory(RecordSource, sizeof(*RecordSource));
    if (TargetKind == KSWORD_ARK_RXPF_TARGET_ALLOCATED_TEST) {
        if (RequestedAddress != 0ULL) {
            return STATUS_INVALID_PARAMETER;
        }
        status = KswRxpfCreateAllocatedTestPage(Flags, RecordSource);
    } else if (TargetKind == KSWORD_ARK_RXPF_TARGET_SELF_IMAGE_TEST) {
        status = KswRxpfCreateSelfImageTestPage(
            RequestedAddress,
            Flags,
            RecordSource);
    } else if (TargetKind == KSWORD_ARK_RXPF_TARGET_EXTERNAL_IMAGE) {
#if KSW_RXPF_ENABLE_EXTERNAL_IMAGE_TARGETS
        /* External image allowlists are intentionally absent in protocol v1. */
        UNREFERENCED_PARAMETER(RequestedAddress);
        status = STATUS_NOT_SUPPORTED;
#else
        UNREFERENCED_PARAMETER(RequestedAddress);
        status = STATUS_NOT_SUPPORTED;
#endif
    } else {
        status = STATUS_INVALID_PARAMETER;
    }
    if (!NT_SUCCESS(status)) {
        return status;
    }

    /* Keep a complete nonpaged mirror of the page's current bytes. */
    RecordSource->CurrentContent = KswordARKAllocateNonPagedPool(
        PAGE_SIZE,
        KSW_RXPF_CONTENT_TAG);
    if (RecordSource->CurrentContent == NULL) {
        KswRxpfImageProtectionReleaseRecord(RecordSource);
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    /* Optional backup is allocated on the control path, never by vector 14. */
    if ((Flags & KSWORD_ARK_RXPF_FLAG_CAPTURE_BACKUP) != 0UL) {
        RecordSource->Backup = KswordARKAllocateNonPagedPool(
            PAGE_SIZE,
            KSW_RXPF_BACKUP_TAG);
        if (RecordSource->Backup == NULL) {
            KswRxpfImageProtectionReleaseRecord(RecordSource);
            return STATUS_INSUFFICIENT_RESOURCES;
        }
        RtlCopyMemory(
            RecordSource->Backup,
            RecordSource->OriginalMapping,
            PAGE_SIZE);
    }
    __try {
        RtlCopyMemory(
            RecordSource->CurrentContent,
            RecordSource->OriginalMapping,
            PAGE_SIZE);
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        status = GetExceptionCode();
        KswRxpfImageProtectionReleaseRecord(RecordSource);
        return status;
    }
    RecordSource->LastWriteOffset = 0UL;
    RecordSource->LastWriteLength = 0UL;
    RtlZeroMemory(RecordSource->LastWriteBytes, sizeof(RecordSource->LastWriteBytes));
    return STATUS_SUCCESS;
}

NTSTATUS
KswRxpfImageProtectionChangeToRwNx(
    _Inout_ PKSW_RXPF_PAGE_RECORD Record
    )
{
    NTSTATUS status = STATUS_NOT_SUPPORTED;

    /* A record transitions exactly once from RX to its persistent NX state. */
    if (Record == NULL || Record->Mdl == NULL ||
        Record->State != KSWORD_ARK_RXPF_PAGE_STATE_RX) {
        return STATUS_INVALID_DEVICE_STATE;
    }
    if (Record->TargetKind == KSWORD_ARK_RXPF_TARGET_ALLOCATED_TEST) {
        /* The public MDL API changes the existing mapping to PAGE_READWRITE/NX. */
        status = MmProtectMdlSystemAddress(Record->Mdl, PAGE_READWRITE);
        if (NT_SUCCESS(status)) {
            Record->WritableAlias =
                (ULONGLONG)(ULONG_PTR)Record->OriginalMapping;
            Record->MappingIsAlias = FALSE;
            Record->CurrentProtection = PAGE_READWRITE;
            Record->WritableAliasProtection = PAGE_READWRITE;
        }
    } else if (Record->TargetKind ==
        KSWORD_ARK_RXPF_TARGET_SELF_IMAGE_TEST) {
        PVOID writableAlias = NULL;

        /* Revalidate the exact MDL low-bit contract immediately before the call. */
        if (!KswRxpfImageProtectionBuildSupported() ||
            (Record->Mdl->MdlFlags &
                (MDL_MAPPED_TO_SYSTEM_VA |
                 MDL_PAGES_LOCKED |
                 MDL_SOURCE_IS_NONPAGED_POOL)) != MDL_PAGES_LOCKED ||
            MmGetMdlByteOffset(Record->Mdl) != 0UL ||
            MmGetMdlByteCount(Record->Mdl) != PAGE_SIZE ||
            MmGetMdlVirtualAddress(Record->Mdl) !=
                (PVOID)(ULONG_PTR)Record->PageBase) {
            return STATUS_INVALID_PARAMETER;
        }
        status = g_KswRxpfImageProtection.MmChangeImageProtection(
            Record->Mdl,
            (PVOID)(ULONG_PTR)Record->PageBase,
            PAGE_SIZE,
            KSW_RXPF_MM_CHANGE_OPERATION_READONLY_NX);
        if (!NT_SUCCESS(status)) {
            Record->LastStatus = status;
            return status;
        }
        Record->CurrentProtection = PAGE_READONLY;

        /* Map a separate writable/NX alias only after the image mapping is NX. */
        writableAlias = MmMapLockedPagesSpecifyCache(
            Record->Mdl,
            KernelMode,
            MmCached,
            NULL,
            FALSE,
            (MM_PAGE_PRIORITY)(NormalPagePriority | MdlMappingNoExecute));
        if (writableAlias == NULL) {
            Record->State = KSWORD_ARK_RXPF_PAGE_STATE_ERROR;
            Record->LastStatus = STATUS_INSUFFICIENT_RESOURCES;
            return STATUS_INSUFFICIENT_RESOURCES;
        }
        Record->WritableAlias = (ULONGLONG)(ULONG_PTR)writableAlias;
        Record->MappingIsAlias = TRUE;
        Record->WritableAliasProtection = PAGE_READWRITE;
    } else {
        return STATUS_NOT_SUPPORTED;
    }

    /* Publish the completed transition only after a writable alias exists. */
    if (NT_SUCCESS(status)) {
        KeMemoryBarrier();
        InterlockedExchange(
            &Record->State,
            KSWORD_ARK_RXPF_PAGE_STATE_RW_NX);
        InterlockedIncrement((volatile LONG*)&Record->Generation);
    } else {
        Record->LastStatus = status;
        Record->LastFailureReason =
            KSWORD_ARK_RXPF_EMULATION_INTERNAL_ERROR;
    }
    return status;
}

NTSTATUS
KswRxpfImageProtectionWrite(
    _Inout_ PKSW_RXPF_PAGE_RECORD Record,
    _In_ ULONG Offset,
    _In_reads_bytes_(Length) const UCHAR* Bytes,
    _In_ ULONG Length
    )
{
    ULONGLONG end = (ULONGLONG)Offset + (ULONGLONG)Length;
    NTSTATUS status = STATUS_SUCCESS;

    /* Writes are bounded to one transitioned page and its explicit alias. */
    if (Record == NULL || Bytes == NULL || Length == 0UL ||
        Record->CurrentContent == NULL ||
        Length > KSWORD_ARK_RXPF_MAX_WRITE_BYTES ||
        end > PAGE_SIZE || end < Offset ||
        Record->State != KSWORD_ARK_RXPF_PAGE_STATE_RW_NX ||
        Record->WritableAlias == 0ULL) {
        return STATUS_INVALID_PARAMETER;
    }
    __try {
        RtlCopyMemory(
            (PUCHAR)(ULONG_PTR)Record->WritableAlias + Offset,
            Bytes,
            Length);
        RtlCopyMemory(
            (PUCHAR)Record->CurrentContent + Offset,
            Bytes,
            Length);
        Record->LastWriteOffset = Offset;
        Record->LastWriteLength = Length;
        RtlZeroMemory(Record->LastWriteBytes, sizeof(Record->LastWriteBytes));
        RtlCopyMemory(Record->LastWriteBytes, Bytes, Length);

        KeMemoryBarrier();
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        status = GetExceptionCode();
    }
    Record->LastStatus = status;
    if (NT_SUCCESS(status)) {
        InterlockedIncrement((volatile LONG*)&Record->Generation);
    }
    return status;
}

VOID
KswRxpfImageProtectionReleaseRecord(
    _Inout_ PKSW_RXPF_PAGE_RECORD Record
    )
{
    /* Release only resources whose ownership flags were recorded at creation. */
    if (Record == NULL) {
        return;
    }
    if (Record->MappingIsAlias && Record->WritableAlias != 0ULL &&
        Record->Mdl != NULL) {
        MmUnmapLockedPages(
            (PVOID)(ULONG_PTR)Record->WritableAlias,
            Record->Mdl);
        Record->WritableAlias = 0ULL;
    } else if (Record->OwnsMdlPages &&
        Record->OriginalMapping != NULL && Record->Mdl != NULL) {
        MmUnmapLockedPages(Record->OriginalMapping, Record->Mdl);
        Record->OriginalMapping = NULL;
        Record->WritableAlias = 0ULL;
    }
    if (Record->PagesLockedByProbe && Record->Mdl != NULL) {
        MmUnlockPages(Record->Mdl);
        Record->PagesLockedByProbe = FALSE;
    }
    if (Record->OwnsMdlPages && Record->Mdl != NULL) {
        MmFreePagesFromMdl(Record->Mdl);
        ExFreePool(Record->Mdl);
        Record->Mdl = NULL;
    } else if (Record->Mdl != NULL) {
        IoFreeMdl(Record->Mdl);
        Record->Mdl = NULL;
    }
    if (Record->Backup != NULL) {
        ExFreePoolWithTag(Record->Backup, KSW_RXPF_BACKUP_TAG);
        Record->Backup = NULL;
    }
    if (Record->CurrentContent != NULL) {
        ExFreePoolWithTag(Record->CurrentContent, KSW_RXPF_CONTENT_TAG);
        Record->CurrentContent = NULL;
    }
}
