contextjail.cpp

//
// contextjail.cpp
// Copyright (c) 2025 sixtyvividtails; provided under the MIT license.
//
// Contextjail test code. Compile using msvc.
// 👉 thread: https://x.com/sixtyvividtails/status/1919833709080879298
// 👉 compiled binary: https://pixeldrain.com/u/1pCYhms4
//
//
//
// TLDR: just spam NtGetContextThread(targetThread).
// As a result, targetThread will jailed, repeatedly running nt!PspGetSetContextSpecialApc in ring0 at APC_LEVEL.
//
// Usecases:
// 1. Pause thread midway in exploit races (even in ring0, since it's context retrieval APC is kernel APC).
// 2. Or block entire CPU core. Since kernel APCs run at APC_LEVEL (🤯), thread scheduling would be kinda
//    disabled (think priority = ∞). We use "witness" thread here to test that assumption, and it appears
//    you'd actually need more than 1 prisoner latched to a particular core to really block that core.
// 3. And more...
//
//
//
// Our setup (with default parameters):
// 1. Master thread, pinned to cpu 0. It's the main thread, it just sets up and monitors stuff, prints stats.
// 2. Prisoner thread, pinned to cpu 1. Has priority 15. Tries exec a few instructions, then sleeps for ~17 ms.
//    Each successful execution of that couple of instructions treated as prison escape.
//    Ideally thread won't execute user code at all until we let it.
// 3. Witness thread, pinned to same cpu as the Prisoner thread. It has priority 31 (or 15, if not privileged).
//    This thread is optional (/w 0 to disable). It's almost like the prisoner, but it sleeps for longer periods,
//    counts execs as "witness moves" rather than escapes, and is not target of the NtGetContextThread spam.
//    Main purpose of the witness thread is to check possibility to hijack cpu core.
// 4. Jailer threads, by default we use 99 of them (mainly for win11 24H2), otherwise 0x10 can be enough.
//    These threads just spam NtGetContextThread against the prisoner thread.
//    We set their priority to 15 and affinity ~3 (so, any processors other than Master's cpu 0 and Prisoner's cpu 1).
//    Their count is important, but their affinity mask is not so much - you can limit all of them to just 4 to 6
//    processors, and that's often will be enough.
//
#include <stdio.h>
#include <conio.h>
#include <stdlib.h>
#include <intrin.h>
#include <stdint.h>
#include "phnt.h"   // https://github.com/mrexodia/phnt-single-header



//----------------------------------------------------------------------------------------------------------------------
// Defines.
//----------------------------------------------------------------------------------------------------------------------

typedef unsigned char uchar, *puchar;
typedef unsigned short ushort, *pushort;
typedef unsigned int uint, *puint;
typedef unsigned long ulong, *pulong;
typedef long long int64, *pint64;
typedef unsigned long long uint64, *puint64;
typedef wchar_t wchar;
typedef _Null_terminated_ char* pstr;
typedef _Null_terminated_ wchar_t* pwstr;
typedef _Null_terminated_ const char* pcstr;
typedef _Null_terminated_ const wchar_t* pcwstr;

#define PAGE_SIZE                       0x1000
#define PAGE_MASK                       0xFFF

#define ATTR_INLINE msvc::forceinline

static bool g_shouldRundown;


//----------------------------------------------------------------------------------------------------------------------
// Some helper routines.
//----------------------------------------------------------------------------------------------------------------------
// Copypasted/edited; refer to umt, don't copypaste from here.

/// <summary>
/// NtSetInformationThread wrapper.
/// </summary>
/// <param name="infoclass"></param>
/// <param name="data">We allow const; caller's responsibility is awareness of infoclasses which do modify data.</param>
/// <param name="thread"></param>
[[ATTR_INLINE]] DECLSPEC_SAFEBUFFERS
inline NTSTATUS set_thread_info(THREADINFOCLASS infoclass, const auto& data, HANDLE thread = NtCurrentThread())
{
    return NtSetInformationThread(thread, infoclass, (PVOID)&data, sizeof(data));
}


/// <summary>
/// Sets both ETHREAD.BasePriority and ETHREAD.Priority to the specified value.
/// If new priority is not the same as the old ETHREAD.Priority, thread quantum is replenished.
/// </summary>
/// <param name="priority">Value in range [1;15] (or up to 31 if SeIncreaseBasePriority privilege enabled).</param>
/// <param name="thread">Requires THREAD_SET_INFORMATION right.</param>
[[ATTR_INLINE]] DECLSPEC_SAFEBUFFERS
inline NTSTATUS set_thread_base_priority(uint priority, HANDLE thread = NtCurrentThread())
{
    return set_thread_info(ThreadActualBasePriority, priority, thread);
}


/// <summary>
/// Sets thread boost state (reciprocal flag: ETHREAD.DisableBoost).
/// </summary>
/// <param name="enableBoost"></param>
/// <param name="thread">Requires THREAD_SET_LIMITED_INFORMATION right.</param>
[[ATTR_INLINE]] DECLSPEC_SAFEBUFFERS
inline NTSTATUS set_thread_priority_boost(bool enableBoost, HANDLE thread = NtCurrentThread())
{
    BOOL disableBoost = !enableBoost;   // monkeys designed some apis
    return set_thread_info(ThreadPriorityBoost, disableBoost, thread);
}


/// <summary>
/// Sets thread affinity to a specified cpu group and mask.
/// </summary>
/// <param name="group"></param>
/// <param name="thread">Requires THREAD_SET_INFORMATION right.</param>
[[ATTR_INLINE]] DECLSPEC_SAFEBUFFERS
inline NTSTATUS set_thread_affinity(const GROUP_AFFINITY& group, HANDLE thread = NtCurrentThread())
{
    return set_thread_info(ThreadGroupInformation, group, thread);
}


/// <summary>
/// Sets thread name (for debugging).
/// </summary>
/// <param name="threadName"></param>
/// <param name="thread"></param>
[[ATTR_INLINE]] DECLSPEC_SAFEBUFFERS
inline NTSTATUS set_thread_name(pcwstr threadName, HANDLE thread = NtCurrentThread())
{
    size_t len = wcslen(threadName);
    UNICODE_STRING ustr{(ushort)(sizeof(wchar)*len), (ushort)(sizeof(wchar)*len), (pwstr)threadName};
    return set_thread_info(ThreadNameInformation, ustr, thread);
}


inline bool is_wow64()
{
#ifdef _WIN64
    return false;
#else
    return (uchar)USER_SHARED_DATA->ImageNumberLow == 0x64;
#endif
}

// lower 6 bits: cpu, higher bits: group
static uint get_cpu_enumber()
{
#if defined (_M_ARM64)
    uint id = (uint)_ReadStatusReg(ARM64_TPIDRRO_EL0);
    return ((id & 0xFF00) | ((uchar)id << 2)) >> 2;
#else
#ifdef _WIN64
    constexpr uint selector = 0x53;     // KGDT64_R3_CMTEB|RPL_MASK   vs   KGDT_R3_TEB|RPL_MASK
#else
    uint selector = is_wow64()? 0x53: 0x3B;
#endif
    uint segLimit = __segmentlimit(selector);
    return ((segLimit >> 14) & 0x3F) | ((uchar)segLimit << 6);
#endif
}


// retrieves valid cpu affinity mask for the current cpu group, in the laziest possible way (destroys current affinity)
static GROUP_AFFINITY get_valid_cpu_affinities()
{
    ushort group = (ushort)(get_cpu_enumber() >> 6);
    size_t mask = 0;
    for (uint i = 0; i < 0x40; ++i)
        if (SUCCEEDED(set_thread_affinity({.Mask{(size_t)(1ui64 << i)}, .Group{group}})))
            mask |= (size_t)(1ui64 << i);
    return {.Mask{mask}, .Group{group}};
}


static uchar popcount64(uint64 v)
{
    v = v - ((v >> 1) & 0x5555'5555'5555'5555);             // put count of each 2 bits into these 2 bits
    v = (v & 0x3333'3333'3333'3333) + ((v >> 2) & 0x3333'3333'3333'3333);     // put count of each 4 bits
    v = (v + (v >> 4)) & 0x0F0F'0F0F'0F0F'0F0F;             // put count of each 8 bits into these 8 bits
    v = (v * 0x0101'0101'0101'0101) >> 0x38;                // left 8 bits of v + (v << 8) + (v << 0x10) + ...
    return (uchar)v;
}


//----------------------------------------------------------------------------------------------------------------------
// Actual test.
//----------------------------------------------------------------------------------------------------------------------

struct TestData
{
    NTSTATUS status;
    volatile bool prisonLocked;     // when enough jailers started jailing
    volatile bool shouldRundown;

    uchar cpuGroup;
    uchar jailersAffinityCount;
    uint jailersCount;
    uint witnessPriority;           // 0: no witness
    int jailTimeSeconds;            // -1: until first escape
    uint contextFlags;
    size_t validAffinity;
    size_t prisonerAffinity;        // also for witness
    size_t jailersAffinity;

    HANDLE signalEvent;
    HANDLE masterThread;
    HANDLE prisonerThread;

    uint prisonerEscapes;
    uint witnessMoves;

    uint tickPrisonLocked;
    uint tickAllJailersJailing;
    uint tickFirstPrisonerEscape;

    DECLSPEC_CACHEALIGN volatile uint jailersIndex;
    DECLSPEC_CACHEALIGN volatile uint threadsReady;
    DECLSPEC_CACHEALIGN volatile uint64 getThreadContextCalls;
};


// Returns exactly STATUS_SUCCESS to indicate thread should continue.
// In that case caller must decrement data.threadsReady on return, otherwise it must not touch data.
static NTSTATUS init_thread(pcwstr name, size_t affinity, uint priority, _Inout_ TestData* data)
{
    NTSTATUS st = set_thread_name(name);
    st = set_thread_priority_boost(false);
    st = set_thread_affinity({.Mask{affinity}, .Group{data->cpuGroup}});
    if (SUCCEEDED(st))
        st = set_thread_base_priority(priority);
    if (FAILED(st))
    {
        data->status = st;
        return st;
    }
    if (_InterlockedIncrement(&data->threadsReady) == 1 + data->jailersCount + !!data->witnessPriority)
        NtAlertThread(data->masterThread);
    if (data->shouldRundown)
        return STATUS_ABANDONED;    // success status
    st = NtWaitForSingleObject(data->signalEvent, false, {});
    if (st == STATUS_WAIT_0 && data->shouldRundown)
        st = STATUS_ABANDONED;    // success status
    if (st != STATUS_WAIT_0)
    {
        static_assert(STATUS_WAIT_0 == STATUS_SUCCESS);
        if (FAILED(st))
            data->status = st;
        _InterlockedDecrement(&data->threadsReady);
        return st;
    }
    return STATUS_SUCCESS;
}


static NTSTATUS prisoner_or_witness_thread(_Inout_ TestData* data, bool isWitness = false)
{
    // note witness affinity by definition is same as prisoner affinity
    NTSTATUS st = init_thread(isWitness? L"witness": L"prisoner", data->prisonerAffinity,
        isWitness? data->witnessPriority: 15, data);
    if (st != STATUS_SUCCESS)   // need to be exactly that to continue; otherwise may not touch data
        return st;

    uint escapes = 0;
    uint escapesLimit = isWitness || data->jailTimeSeconds >= 0? UINT_MAX: 0;
    volatile uint* escapesPtr = isWitness? &data->witnessMoves: &data->prisonerEscapes;

    // allow to run while there's not enough jailers yet
    volatile int spinner = 0;
    while (!data->prisonLocked && !data->shouldRundown)
        ++spinner;
    if (spinner == INT_MAX && __rdtsc() == 0)   // never true, just to ensure spinner won't be optimized out
        st = STATUS_INTERNAL_ERROR;
    LARGE_INTEGER sleepAfterEscape{.QuadPart{isWitness? -250'000'0: -17'000'0}};

    while (!data->shouldRundown)
    {
#if defined(_M_IX86) || defined(_M_X64)
        _mm_pause();
#else
        __yield();
#endif
        if (data->shouldRundown)
            break;
        *escapesPtr = ++escapes;
        if (escapes == 1 && !isWitness)
            data->tickFirstPrisonerEscape = NtGetTickCount();
        if (escapes > escapesLimit)
        {
            st = STATUS_LOCK_NOT_GRANTED;
            break;
        }
        NtDelayExecution(false, &sleepAfterEscape);
    }

    NtAlertThread(data->masterThread);
    _InterlockedDecrement(&data->threadsReady);
    return st;
}


static NTSTATUS jailer_thread(_Inout_ TestData* data)
{
    // prepare local data
    HANDLE prisoner = data->prisonerThread;
    CONTEXT ctx;
    ctx.ContextFlags = data->contextFlags;

    // prepare name and init this jailer thread
    NTSTATUS st;
    uint jailerId = _InterlockedIncrement(&data->jailersIndex) - 1;
    {
        wchar threadName[] = L"jailer AAAAAA";
        for (pwstr p = threadName + 7; jailerId && *p; ++p)
        {
            *p += jailerId % 26;
            jailerId /= 26;
        }
        st = init_thread(threadName, data->jailersAffinity, 15, data);
    }
    if (st != STATUS_SUCCESS)   // need to be exactly that to continue; otherwise may not touch data
        return st;

    // try initial context retrieval
    st = NtGetContextThread(prisoner, &ctx);
    _InterlockedIncrement(&data->getThreadContextCalls);    // stats
    if (SUCCEEDED(st))  [[likely]]
    {
        // assume prison is locked once at least 6 jailers started jailing (or less if there's less than 6 jailers)
        uint jailersCount = data->jailersCount;
        uint prisonLockWhen = jailersCount > 6? 6: jailersCount;
        uint jailersJailing = _InterlockedIncrement(&data->jailersIndex) - jailersCount;
        uint tickNow = NtGetTickCount();
        if (jailersJailing == prisonLockWhen)
        {
            data->prisonLocked = true;      // enough jailers at work, prisoner shouldn't run anymore
            data->tickPrisonLocked = tickNow;
        }
        if (jailersJailing == jailersCount)
            data->tickAllJailersJailing = tickNow;
    }
    // now just spam NtGetContextThread
    while (!data->shouldRundown && SUCCEEDED(st))
    {
        st = NtGetContextThread(prisoner, &ctx);
        _InterlockedIncrement(&data->getThreadContextCalls);
    }
    if (FAILED(st) && (st != STATUS_UNSUCCESSFUL || !data->shouldRundown))
    {
        // we get STATUS_UNSUCCESSFUL when prisoner thread is running down (apc queueuing gets disabled)
        data->status = st;
        NtAlertThread(data->masterThread);
    }

    _InterlockedDecrement(&data->threadsReady);
    return st;
}


static NTSTATUS test_context_jail2(_Inout_ TestData* data)
{
    printf("[ ] initializing threads...\n");

    uint threadFlags = THREAD_CREATE_FLAGS_SKIP_THREAD_ATTACH;
    if (USER_SHARED_DATA->NtMajorVersion >= 10)
        threadFlags |= THREAD_CREATE_FLAGS_SKIP_LOADER_INIT;

    for (uint i = 0; i < 2; ++i)
    {
        HANDLE thread{};
        NTSTATUS st = NtCreateThreadEx(&thread, THREAD_GET_CONTEXT, {}, NtCurrentProcess(), [](PVOID arg)
        {
            return prisoner_or_witness_thread((TestData*)((size_t)arg & ~1), !!((size_t)arg & 1));
        }, (PVOID)((size_t)data | i), threadFlags, 0, 0x4000, 0x8000, {});
        if (FAILED(st))
            return st;
        if (i == 0)
            data->prisonerThread = thread;
        else
            NtClose(thread);
        if (!data->witnessPriority)
            break;
    }
    uint threadsReadyTarget = 1 + !!data->witnessPriority + data->jailersCount;
    for (uint i = 0; i < data->jailersCount; ++i)
    {
        if (FAILED(data->status))
        {
            printf("[x] some child thread failed: %08X, threads ready: %u/%u\n",
                data->status, data->threadsReady, threadsReadyTarget);
            return STATUS_THREAD_NOT_RUNNING;
        }
        if (g_shouldRundown)
        {
            printf("[!] cancelled\n");
            return STATUS_CANCELLED;
        }
        HANDLE thread{};
        NTSTATUS st = NtCreateThreadEx(&thread, SYNCHRONIZE, {}, NtCurrentProcess(), [](PVOID arg)
        {
            return jailer_thread((TestData*)arg);
        }, data, threadFlags, 0, 0x4000, 0x8000, {});
        if (FAILED(st))
            return st;
        NtClose(thread);
    }

    printf("[ ] waiting for threads to finish init...\n");
    NTSTATUS st = STATUS_SUCCESS;
    for (uint retries = 0; data->threadsReady < threadsReadyTarget && retries < 60; ++retries)
    {
        if (retries == 5)
            printf("[!] still waiting, threads done: %u/%u\n", data->threadsReady, threadsReadyTarget);
        if (FAILED(data->status))
        {
            printf("[x] some child thread failed: %08X, threads ready: %u/%u\n",
                data->status, data->threadsReady, threadsReadyTarget);
            return STATUS_THREAD_NOT_RUNNING;
        }
        if (g_shouldRundown)
        {
            printf("[!] cancelled\n");
            return STATUS_CANCELLED;
        }
        LARGE_INTEGER timeout{.QuadPart{-1'000'000'0}};
        st = NtDelayExecution(true, &timeout);
    }
    if (data->threadsReady < threadsReadyTarget)
    {
        printf("[x] failed to init threads, threads ready: %u/%u", data->threadsReady, threadsReadyTarget);
        return STATUS_THREAD_NOT_IN_PROCESS;
    }

    if (st != STATUS_ALERTED)
    {
        // to minimize chance of spurious wakeup
        printf("[ ] waiting for the final alert...\n");
        LARGE_INTEGER timeout{.QuadPart{-1'000'000'0}};
        st = NtDelayExecution(true, &timeout);      // and ignore if alert never comes
    }
    printf("[+] all threads ready\n");
    if (g_shouldRundown)
    {
        printf("[!] cancelled\n");
        return STATUS_CANCELLED;
    }

    int jailTime = data->jailTimeSeconds;
    if (jailTime < 0)
        printf("[ ] signaling threads and waiting until jail escape...\n");
    else
        printf("[ ] signaling threads and waiting for %u seconds...\n", data->jailTimeSeconds);

    uint64 ticks0 = NtGetTickCount64();
    uint64 ticksTarget = jailTime >= 0? ticks0 + 1000ui64 * jailTime: UINT64_MAX;
    LARGE_INTEGER timeout{.QuadPart{-1'000'000'0i64 * jailTime}};
    if (jailTime < 0 || jailTime > 10)
        timeout.QuadPart = -4'500'000'0;
    st = NtSignalAndWaitForSingleObject(data->signalEvent, NtCurrentThread(), true, &timeout);
    while (!g_shouldRundown && SUCCEEDED(data->status) && (data->jailTimeSeconds >= 0 || data->prisonerEscapes == 0))
    {
        uint64 ticksNow = NtGetTickCount64();
        if (ticksNow > ticksTarget - 250)
            break;
        double elapsed = (ticksNow - ticks0) / 1000.0;
        printf("[*] elapsed %.01f/%d seconds, prisoner escapes: %u, witness moves: %u, getContext calls: %I64u\n",
            elapsed, jailTime, data->prisonerEscapes, data->witnessMoves, data->getThreadContextCalls);
        uint64 delta = ticksTarget - ticksNow;
        timeout.QuadPart = delta < 12'000? -1'000'0i64 * delta: -9'500'000'0;
        NtDelayExecution(true, &timeout);
    }

    double elapsed = (NtGetTickCount64() - ticks0) / 1000.0;
    int allJailersJailingDelta = data->tickAllJailersJailing?
        data->tickAllJailersJailing - data->tickPrisonLocked: 0;    // might still get negative coz preemption
    int firstEscapeDelta = (data->tickPrisonLocked && data->tickFirstPrisonerEscape)?
        data->tickFirstPrisonerEscape - data->tickPrisonLocked: 0;  // again, might get negative here
    printf("[+] done; elapsed %.01f/%d seconds, threads status: %08X, rundown enforced: %u\n"
        "    prisoner escapes: %u, witness moves: %u, getContext calls: %I64u\n"
        "    ticks delta: prison locked: 0, all jailers working: %d ms, first escape: %d ms\n\n",
        elapsed, jailTime, data->status, g_shouldRundown,
        data->prisonerEscapes, data->witnessMoves, data->getThreadContextCalls,
        allJailersJailingDelta, firstEscapeDelta);
    if (data->prisonerEscapes)
    {
        printf("[x] >>>>>>>> PRISONER ESCAPED, remediation needed\n");
        if (data->jailersAffinityCount < 6 || data->jailersCount < 99)
            printf("    possible reason: jailers processors below 6, or jailers count below 99\n\n");
        else
            printf("    possible reason: unknown; but try using more jailers or more jailer processors\n\n");
        return STATUS_SYSTEM_NEEDS_REMEDIATION;
    }
    st = data->status;
    if (FAILED(st))
    {
        printf("[*] >>>>>>>> PRISONER status unknown: %08X, remediation needed\n\n", st);
        return st;
    }
    printf("[+] >>>>>>>> PRISONER STILL JAILED, all good\n\n");
    if (g_shouldRundown && jailTime >= 0)
        return STATUS_CANCELLED;
    return STATUS_SUCCESS;
}


static NTSTATUS test_context_jail(uint jailersCount, int jailTimeSeconds, uint witnessPriority,
    uint64 jailersAffinity, uint64 prisonerAffinity, bool fullContext)
{
    if (is_wow64())
    {
        printf("[x] wow64 is not supported\n"); // technically it's fine, but requires some tinkering with native api
        return STATUS_WOW_ASSERTION;
    }
    GROUP_AFFINITY validAffinity = get_valid_cpu_affinities();
    printf("[ ] affinity mask for cpu group %u: %016I64X\n", validAffinity.Group, (uint64)validAffinity.Mask);

    // we're the master thread, set name and affinity
    set_thread_name(L"master");
    NTSTATUS st = set_thread_affinity({.Mask{1}, .Group{validAffinity.Group}});
    if (FAILED(st) || (validAffinity.Mask & (validAffinity.Mask - 1)) == 0)
    {
        // we need at least 2 bits set
        printf("[x] valid affinity mask unsuitable for test\n");
        return STATUS_CPU_SET_INVALID;
    }

    TestData data
    {
        .cpuGroup = (uchar)validAffinity.Group,
        .jailersAffinityCount = popcount64(jailersAffinity & validAffinity.Mask),
        .jailersCount = jailersCount,
        .witnessPriority = witnessPriority,
        .jailTimeSeconds = jailTimeSeconds,
        .contextFlags = (uint)(fullContext? CONTEXT_ALL: CONTEXT_CONTROL),  // LATER: CONTEXT_XSTATE?
        .validAffinity = validAffinity.Mask,
        .prisonerAffinity = prisonerAffinity & validAffinity.Mask,
        .jailersAffinity = jailersAffinity & validAffinity.Mask,
    };

    if ((data.prisonerAffinity & (data.prisonerAffinity - 1)) != 0)
        data.witnessPriority = 0;       // don't use witness when prisoner affinity spans multiple cpus
    if (!data.prisonerAffinity || !data.jailersAffinity)
    {
        printf("[x] updated prisoner affinity %016I64X or jailers affinity %016I64X unsuitable for test\n",
            (uint64)data.prisonerAffinity, (uint64)data.jailersAffinity);
        return STATUS_CPU_SET_INVALID;
    }
    BOOLEAN wasEnabled{true};
    if (data.witnessPriority > 15)
    {
        if (FAILED(RtlAdjustPrivilege(SE_INC_BASE_PRIORITY_PRIVILEGE, true, false, &wasEnabled)))
        {
            printf("[!] unable to adjust SeIncreaseBasePriority privilege, witness will be limited to priority 15\n");
            data.witnessPriority = 15;
        }
    }

    st = NtCreateEvent(&data.signalEvent, EVENT_ALL_ACCESS, {}, NotificationEvent, false);
    if (FAILED(st))
        return st;
    st = NtDuplicateObject(NtCurrentProcess(), NtCurrentThread(), NtCurrentProcess(), &data.masterThread,
        THREAD_ALL_ACCESS, 0, 0);
    if (FAILED(st))
    {
        NtClose(data.signalEvent);
        return st;
    }

    if (data.jailersAffinityCount < 4)
        printf("[!] WARNING: jailers processors count is %u, which is quite low, consider using more CPUs\n",
            data.jailersAffinityCount);
    if (data.jailersCount < 12)
        printf("[!] WARNING: jailers count is %u, which is quite low, consider using more jailers\n",
            data.jailersCount);

    char witnessInfo[0x20];
    witnessInfo[0] = {};
    if (data.witnessPriority)
        sprintf_s(witnessInfo, ", priority %u", data.witnessPriority);
    char jailTimeString[0x20];
    strcpy_s(jailTimeString, "until jail escape");
    if (data.jailTimeSeconds >= 0)
        sprintf_s(jailTimeString, "%u seconds", data.jailTimeSeconds);
    printf("[+] ready to test, parameters:\n"
        "    prisoner affinity: %016I64X, witness: %s%s\n"
        "    jailers  affinity: %016I64X, %u processors; jailers count: %u\n"
        "    timelimit: %s\n",
        (uint64)data.prisonerAffinity, data.witnessPriority? "enabled": "disabled", witnessInfo,
        (uint64)data.jailersAffinity, data.jailersAffinityCount, data.jailersCount,
        jailTimeString);
    st = test_context_jail2(&data);
    data.shouldRundown = true;
    NtSetEvent(data.signalEvent, {});
    printf("[ ] waiting for threads to rundown...\n");
    for (uint retries = 0; data.threadsReady && retries < 300; ++retries)
    {
        if (retries == 30)
            printf("[!] still waiting, threads left: %u\n", data.threadsReady);
        LARGE_INTEGER timeout{.QuadPart{-100'000'0}};
        NtDelayExecution(true, &timeout);
    }
    if (data.threadsReady)
        printf("[x] gave up waiting, threads left: %u; that's potentially fatal\n", data.threadsReady);
    NtClose(data.signalEvent);
    NtClose(data.masterThread);
    if (data.prisonerThread)
        NtClose(data.prisonerThread);
    return st;
}


//----------------------------------------------------------------------------------------------------------------------
// Entry/diag/args/etc.
//----------------------------------------------------------------------------------------------------------------------
// Parts just copypasted in from umt; don't use it from here, refer to the lib instead.

static void print_help()
{
    printf(R"(
contextjail.exe - ContextJail test.

Arguments:
  /j     JAILERS_COUNT          - specifies number of jailers (default 99)
  /t     JAIL_TIME              - time to spend in jail, seconds (default 30); -1: till escape
  /f                            - retrieve CONTEXT_ALL (otherwise just CONTEXT_CONTROL)
  /b                            - batch mode (never pause at the end)
  /w     WITNESS_PRIORITY       - priority of the witness thread (default 31); -1: no witness
  /jaff  JAILERS_AFFINITY       - uint64 affinity mask for jailers (default ~3)
  /paff  PRISONER_AFFINITY      - uint64 affinity mask for the prisoner (default 2)

Note for simplicity we mostly use currently assigned cpu group for all threads, and cpu 0 for the master thread.
We won't add witness if prisoner's affinity spans more than 1 cpu.
)");
}


struct Args
{
    uint jailersCount{};
    int jailTimeSeconds{};
    int witnessPriority{};
    uint64 jailersAffinity{};
    uint64 prisonerAffinity{};
    bool fullContext: 1{};
    bool batch: 1{};
    bool help: 1{};

    bool parse(int argc, _In_ pcwstr argv[])
    {
        for (int i = 0; i < argc; ++i)
        {
            auto arg = argv[i];
            if (arg[0] != '/' && arg[0] != '-')
                continue;
            ++arg;
            if (_wcsicmp(arg, L"?") == 0)
                help = true;
            if (_wcsicmp(arg, L"b") == 0)
                batch = true;
        }
        for (int i = 0; i < argc; ++i)
        {
            auto arg = argv[i];
            if (arg[0] != '/' && arg[0] != '-')     // might(!) be app name
            {
                if (i == 0)
                    continue;
                printf("[x] unexpected arg %d: %S\n", i, arg);
                return false;
            }
            ++arg;
            if (_wcsicmp(arg, L"?") == 0 || _wcsicmp(arg, L"b") == 0)
                continue;
            if (_wcsicmp(arg, L"f") == 0)
            {
                fullContext = true;
                continue;
            }
            int what = 1 * !!!_wcsicmp(arg, L"j") + 2 * !!!_wcsicmp(arg, L"t") + 4 * !!!_wcsicmp(arg, L"w")
                + 8 * !!!_wcsicmp(arg, L"jaff") + 0x10 * !!!_wcsicmp(arg, L"paff");     // "!!!" is for linter, lol
            if (!what)
            {
                printf("[x] unexpected arg %d: %S\n", i, arg);
                return false;
            }

            if (++i >= argc)
            {
                printf("[x] value absent for arg %d %S\n", i-1, arg-1);
                return false;
            }
            bool notop{};
            pcwstr p = argv[i];
            while (*p == ' ')
                ++p;
            notop = (*p == '~');
            if (notop)
                ++p;
            wchar clear[0x80];
            pwstr d = clear;
            for (; *p && d < clear + _countof(clear) - 1; ++p)
                if (*p != '\'' && *p != '`' && *p != '_')
                    *d++ = *p;
            while (*p == ' ')
                ++p;
            *d = {};
            wchar* end{};
            int64 v = _wcstoi64(clear, &end, 0);
            if (notop)
                v = ~v;
            if (*p
                || (!v && (!end || end == clear))
                || ((v < -1 || v > INT_MAX) && what <= 4)
                || (v < 0 && what == 1))
            {
                printf("[x] invalid value for arg %d %S: %S\n", i-1, arg-1, argv[i]);
                return false;
            }
            if (what == 1)
                jailersCount = (uint)v;
            else if (what == 2)
                jailTimeSeconds = (int)v;
            else if (what == 4)
                witnessPriority = (int)v;
            else if (what == 8)
                jailersAffinity = v;
            else
                prisonerAffinity = v;
        }
        normalize();
        return true;
    }


    void normalize()
    {
        if (!jailersCount)
            jailersCount = 99;      // JAILERS_COUNT default
        if (!jailTimeSeconds)
            jailTimeSeconds = 30;   // JAIL_TIME default
        if (witnessPriority < 0)
            witnessPriority = 0;    // no witness
        else if (witnessPriority == 0 || witnessPriority > 31)
            witnessPriority = 31;   // WITNESS_PRIORITY default, would become 15 if no SeIncreaseBasePriority privilege
        if (!prisonerAffinity)
            prisonerAffinity = 2;   // PRISONER_AFFINITY default
        if (!jailersAffinity)
        {
            jailersAffinity = ~3ui64;       // JAILERS_AFFINITY default
            if ((prisonerAffinity & (prisonerAffinity - 1)) == 0)   // JAILERS_AFFINITY when prisoner signle bit set
                jailersAffinity = ~1ui64 & ~prisonerAffinity;       // all but bit0 [master thread], and prisoner
            if (!jailersAffinity)
                jailersAffinity = ~3ui64;   // whatever, go default
        }
    }
};


DECLSPEC_SAFEBUFFERS
static inline NTSTATUS get_cpu_vendor(
    _Always_(_Post_z_) _Out_z_cap_post_count_(vendorBufSize, 13) char* vendor,
    size_t vendorBufSize)
{
    if (vendorBufSize < 13) [[unlikely]]
    {
        if (vendorBufSize)
            *vendor = {};
        return STATUS_BUFFER_TOO_SMALL;
    }
    int regs[4];
    __cpuidex(regs, 0, 0);
    (puint(vendor))[0] = regs[1];
    (puint(vendor))[1] = regs[3];
    (puint(vendor))[2] = regs[2];
    vendor[12] = {};
    return STATUS_SUCCESS;
}


DECLSPEC_SAFEBUFFERS
static inline NTSTATUS get_cpu_brand(
    _Always_(_Post_z_) _Out_z_cap_post_count_(brandBufSize, 0x30) char* brand,
    size_t brandBufSize)
{
    if (brandBufSize < 0x30) [[unlikely]]
    {
        if (brandBufSize)
            *brand = {};
        return STATUS_BUFFER_TOO_SMALL;
    }
    for (uint i = 0; i < 3; ++i)
    {
        int regs[4];
        __cpuidex(regs, 0x8000'0002 + i, 0);
        (puint(brand))[i*4 + 0] = regs[0];
        (puint(brand))[i*4 + 1] = regs[1];
        (puint(brand))[i*4 + 2] = regs[2];
        (puint(brand))[i*4 + 3] = regs[3];
    }
    brand[0x2F] = {};     // should be null-terminated, but we enforce that
    return STATUS_SUCCESS;
}


DECLSPEC_SAFEBUFFERS
static inline NTSTATUS get_cpu_brand_trimmed(
    _Always_(_Post_z_) _Out_z_cap_(brandBufSize) char* brand,
    size_t brandBufSize)
{
    if (brandBufSize)
        *brand = {};
    char buf[0x30];
    NTSTATUS st = get_cpu_brand(buf, _countof(buf));
    if (FAILED(st))
        return st;
    char* p = buf;
    while (*p == ' ')
        ++p;
    char* e = &buf[0x2E];
    while (e >= p && (!*e || *e == ' '))
        --e;
    ++e;
    size_t toCopy = e - p + 1;
    if (brandBufSize < toCopy) [[unlikely]]
        return STATUS_BUFFER_TOO_SMALL;
    *e = {};
    for (uint i = 0; i < toCopy; ++i)
        *brand++ = *p++;
    return STATUS_SUCCESS;
}


DECLSPEC_SAFEBUFFERS
_Success_(return >= 0)
static inline int sprintf_sx(
    _Out_writes_z_(bufferCount) _Always_(_Post_z_) char* buffer,
    _In_ size_t bufferCount,
    _In_z_ _Printf_format_string_ pcstr format,
    ...
)
{
    va_list args;
    va_start(args, format);
    return vsprintf_s(buffer, bufferCount, format, args);
}


static NTSTATUS FASTCALL format_info_banner(
    _In_opt_ const char* componentName,
    _Out_writes_z_(bufferSize) _Always_(_Post_z_) char* buffer,
    size_t bufferSize)
{
    char cpuVendor[0x10];
    get_cpu_vendor(cpuVendor, _countof(cpuVendor));
    char cpuBrand[0x30];
    get_cpu_brand_trimmed(cpuBrand, _countof(cpuBrand));

    // basic hypervisor info; refer to hv/mc for proper
    char hvBrand[0x14]{};
    int regs[4]{};
    __cpuidex(regs, 0x4000'0000, 0);    // retrieves max leaf for this range, and string like 'Microsoft Hv'
    uint hyperMaxLeaf = regs[0];
    if (hyperMaxLeaf < 0x8000'0000u && (hyperMaxLeaf > 0x4000'0000 || (hyperMaxLeaf == 0x4000'0000 && regs[2])))
    {
        // condition above should cover us in case there's no proper hypervisor (and cpu ignored this leaf properly)
        (puint(hvBrand))[0] = regs[1];      // ebx
        (puint(hvBrand))[1] = regs[2];      // ecx
        (puint(hvBrand))[2] = regs[3];      // edx
        if (hyperMaxLeaf >= 0x4000'0001)
        {
            __cpuidex(regs, 0x4000'0001, 0);
            size_t len = strlen(hvBrand);
            if (len && regs[0])             // e.g. "Hv#1"
                hvBrand[len++] = '/';
            *puint(&hvBrand[len]) = regs[0];
        }
        for (char& c: hvBrand)
            if (c && (c < ' ' || c > 0x7F))
                c = '.';
    }
    if (!hvBrand[0])
        *puint(hvBrand) = '-';
    __cpuidex(regs, 1, 0);
    uint ecx = regs[2];
    bool hypervisorPresent = (ecx >> 0x1F) != 0;

    SYSTEM_CODEINTEGRITY_INFORMATION ciInfo{.Length{sizeof(ciInfo)}};
    if (FAILED(NtQuerySystemInformation(SystemCodeIntegrityInformation, &ciInfo, sizeof(ciInfo), {})))
        ciInfo.CodeIntegrityOptions = 0;

    uint timestamp = ((IMAGE_NT_HEADERS*)((SIZE_T)&__ImageBase + __ImageBase.e_lfanew))->FileHeader.TimeDateStamp;
    int rez = sprintf_sx(buffer, bufferSize,
        "%s started, brid %08X x%02u\n"
        "    os %u.%u.%05u, img %04X, server: %u; bootId: %u\n"
        "    cpus: %u/%u, groups: %u, ram: %u MB; cpu: %s, %s\n"
        "    hypervisor: %s, present: %u, CodeIntegrity: %08X, SecureBoot: %u, KernelDebugger: %02X",
        componentName? componentName: "exe", timestamp, (uint)(sizeof(void*) * CHAR_BIT),
        USER_SHARED_DATA->NtMajorVersion, USER_SHARED_DATA->NtMinorVersion, USER_SHARED_DATA->NtBuildNumber,
        USER_SHARED_DATA->ImageNumberLow, USER_SHARED_DATA->NtProductType > NtProductWinNt, USER_SHARED_DATA->BootId,
        USER_SHARED_DATA->UnparkedProcessorCount, USER_SHARED_DATA->ActiveProcessorCount,
        USER_SHARED_DATA->ActiveGroupCount, USER_SHARED_DATA->NumberOfPhysicalPages/(1024*1024/0x1000),
        cpuVendor, cpuBrand,
        hvBrand, hypervisorPresent,
        ciInfo.CodeIntegrityOptions, USER_SHARED_DATA->DbgSecureBootEnabled, USER_SHARED_DATA->KdDebuggerEnabled);
    if (rez <= 0)
    {
        *buffer = {};
        return STATUS_BUFFER_TOO_SMALL;
    }
    return STATUS_SUCCESS;
}


static void FASTCALL print_info_banner(_In_opt_ pcstr componentName)
{
    char banner[0x800];
    if (FAILED(format_info_banner(componentName, banner, _countof(banner))))
        strcpy_s(banner, "<info_error>");
    printf("%s\n", banner);
}


static void FASTCALL wait_console_keypress_if_needed(_In_opt_ const char* message)
{
    ULONG dummy;
    bool needPause =
        GetConsoleMode(GetStdHandle(STD_OUTPUT_HANDLE), &dummy)     // success, this means output is not redirected
        && GetConsoleProcessList(&dummy, 1) <= 1;                   // single console process: it's our process, no cmd
    if (!needPause)
        return;

    if (message)
        printf("%s\n", message);
    _flushall();
    int c = _getch();
    if (!c || c == 0xE0)    // arrow or function key, read one more char to drain queue
        (void)_getch();
}


static void init()
{
    SetErrorMode(SEM_NOOPENFILEERRORBOX);
    _set_invalid_parameter_handler([](pcwstr expr, pcwstr func, pcwstr, uint line, uintptr_t)
    {
        printf("WARNING: invalid parameter handler invoked; func: '%S', line: %u, expr: '%S'\n", func, line, expr);
    });

    static HANDLE mainThread;   // for Ctrl+C, won't be released
    NtDuplicateObject(NtCurrentProcess(), NtCurrentThread(), NtCurrentProcess(), &mainThread, THREAD_ALERT, 0, 0);

    SetConsoleCtrlHandler([](ulong ctrl) -> BOOL
    {
        if (ctrl != CTRL_C_EVENT && ctrl != CTRL_BREAK_EVENT && ctrl != CTRL_CLOSE_EVENT
            && ctrl != CTRL_LOGOFF_EVENT && ctrl != CTRL_SHUTDOWN_EVENT)
        {
            return false;
        }
        printf("[!] got ctrl+xxx signal: %u...\n", ctrl);
        g_shouldRundown = true;
        NtAlertThread(mainThread);
        return true;    // handled
    }, TRUE);
}


int __cdecl wmain(int argc, const wchar* argv[])
{
    init();
    print_info_banner("ContextJail test");

    Args args{};
    bool parseOk = args.parse(argc, argv);
    if (!parseOk || args.help)
    {
        print_help();
        if (!args.batch)
            wait_console_keypress_if_needed("Press any key to exit...");
        return parseOk? 0: -1;
    }

    uint64 ticks0 = NtGetTickCount64();
    NTSTATUS st = test_context_jail(args.jailersCount, args.jailTimeSeconds, args.witnessPriority,
        args.jailersAffinity, args.prisonerAffinity, args.fullContext);
    printf("ContextJail status: %08X; spent %u seconds\n", st, (uint)((NtGetTickCount64() - ticks0)/1000u));

    if (!args.batch)
        wait_console_keypress_if_needed("Press any key to exit...");
    return (int)st;
}