NewStack.cpp

#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <tchar.h>
#include <winnt.h>
#include <iostream>
#include <winternl.h>

typedef NTSTATUS(NTAPI* pfnNtQueryInformationThread)(
	HANDLE ThreadHandle,
	THREADINFOCLASS ThreadInformationClass,
	PVOID ThreadInformation,
	ULONG ThreadInformationLength,
	PULONG ReturnLength
	);

typedef struct _THREAD_BASIC_INFORMATION {
	NTSTATUS ExitStatus;
	PVOID TebBaseAddress;
	CLIENT_ID ClientId;
	KAFFINITY AffinityMask;
	LONG Priority;
	LONG BasePriority;
} THREAD_BASIC_INFORMATION, * PTHREAD_BASIC_INFORMATION;

#define STACK_INITIAL_PAGE_COUNT 2
#define STACK_GUARD_SLACK_PAGE_COUNT 1

struct ManagedStack 
{
	// Self-managed resource fields
	PBYTE StackLowerBound; // Initial commit address and guard lower boundary
	PBYTE GuardUpperBound; // End of the gaurd; available space
	PBYTE StackUpperBound; // End of the entire committed region
	SIZE_T PageSize;	   // System-derived page size
	PNT_TIB TIB;		   // Managed thread's thread information block
	
	// Kernel managed resource fields
	HANDLE ThreadHandle;   // Handle to a created thread from CreateThread()
	DWORD ThreadId;		   // Thread ID value from CreateThread()
} g_ManagedStack = { 0 };

// Checks if exception originates from managed range and dynamically grows stack
// if remaining guard boundary is available.
LONG WINAPI ManagedStackExceptionHandler(PEXCEPTION_POINTERS ExceptionPointers)
{
	PEXCEPTION_RECORD exceptionRecord = ExceptionPointers->ExceptionRecord;

	// If this isn't a guard page violation, immediately continue the search
	if (exceptionRecord->ExceptionCode != STATUS_GUARD_PAGE_VIOLATION)
	{
		return EXCEPTION_CONTINUE_SEARCH;
	}

	// If the faulting address is not within the guard range of the managed stack,
	// continue the search
	PBYTE faultAddress = (PBYTE)exceptionRecord->ExceptionInformation[1];
	if (
		faultAddress < g_ManagedStack.StackLowerBound ||
		faultAddress >= g_ManagedStack.GuardUpperBound)
	{
		return EXCEPTION_CONTINUE_SEARCH;
	}

	// Guard pages are cleared on access and it was allocated with RW, so just update
	// the manager's tracking iff not overflowed
	SIZE_T pageIndex = (faultAddress - g_ManagedStack.StackLowerBound) / g_ManagedStack.PageSize;
	PBYTE hitPage = g_ManagedStack.StackLowerBound + pageIndex * g_ManagedStack.PageSize;
	if (hitPage == g_ManagedStack.StackLowerBound)
	{
		ExitProcess(ERROR_STACK_OVERFLOW);
	}

	g_ManagedStack.GuardUpperBound = hitPage;
	g_ManagedStack.TIB->StackLimit = hitPage;

	return EXCEPTION_CONTINUE_EXECUTION;
}

DWORD AllocManagedStack()
{
	// Get the system page size and create the entire region at once
	SYSTEM_INFO systemInfo = { 0 };
	GetSystemInfo(&systemInfo);
	g_ManagedStack.PageSize = systemInfo.dwPageSize;
	SIZE_T stackSize = 
		g_ManagedStack.PageSize * (STACK_INITIAL_PAGE_COUNT + STACK_GUARD_SLACK_PAGE_COUNT);
	SIZE_T guardRegionSize =
		g_ManagedStack.PageSize * STACK_GUARD_SLACK_PAGE_COUNT;

	g_ManagedStack.StackLowerBound = (PBYTE)VirtualAlloc(
		NULL,
		stackSize,
		MEM_COMMIT | MEM_RESERVE,
		PAGE_READWRITE
	);
	if (!g_ManagedStack.StackLowerBound) {
		return GetLastError();
	}
	g_ManagedStack.StackUpperBound = g_ManagedStack.StackLowerBound + stackSize;

	// Apply guard page protections
	DWORD flOldProtect;
	BOOL bStatus = VirtualProtect(
		g_ManagedStack.StackLowerBound,
		guardRegionSize,
		PAGE_READWRITE | PAGE_GUARD,
		&flOldProtect
	);
	if (!bStatus) {
		return GetLastError();
	}
	g_ManagedStack.GuardUpperBound =
		g_ManagedStack.StackLowerBound + guardRegionSize;

	return ERROR_SUCCESS;
}

// Get the TIB address for a thread by its handle
DWORD GetTibByHandle(HANDLE hThread, PNT_TIB* ppNtTib)
{
	HMODULE hNtdll = GetModuleHandleA("C:\\Windows\\System32\\ntdll.dll");
	if (!hNtdll)
	{
		return GetLastError();
	}

	pfnNtQueryInformationThread NtQueryInformationThread =
		(pfnNtQueryInformationThread)GetProcAddress(hNtdll, "NtQueryInformationThread");
	if (!NtQueryInformationThread)
	{
		return GetLastError();
	}

	THREAD_BASIC_INFORMATION threadBasicInfo = { 0 };
	DWORD threadBasicInfoLength = 0;
	NTSTATUS ntstatus = NtQueryInformationThread(
		hThread,
		(THREADINFOCLASS)0,
		&threadBasicInfo,
		sizeof(threadBasicInfo),
		&threadBasicInfoLength
	);
	if (ntstatus != 0) // !STATUS_SUCCESS
	{
		return ERROR_NOT_FOUND;
	}
	if ((PNT_TIB)threadBasicInfo.TebBaseAddress == NULL)
	{
		return ERROR_NOT_FOUND;
	}

	// TIB is the first field in TEB
	*ppNtTib = (PNT_TIB)threadBasicInfo.TebBaseAddress;
	return ERROR_SUCCESS;
}

// Create a new thread with CreateThread(), updating its state to use ManagedStack iff
// initialized.
DWORD CreateThreadWithManagedStack(LPTHREAD_START_ROUTINE lpStartAddress, LPVOID lpParameter)
{
	// Check manager initialization and arguments
	if (
		!g_ManagedStack.StackLowerBound ||
		!g_ManagedStack.GuardUpperBound ||
		!g_ManagedStack.StackUpperBound)
	{
		return ERROR_INVALID_ADDRESS;
	}

	if (!g_ManagedStack.PageSize)
	{
		return ERROR_INVALID_DATA;
	}

	if (!lpStartAddress) {
		return ERROR_INVALID_PARAMETER;
	}

	// Create the new thread in a suspended state
	g_ManagedStack.ThreadHandle = CreateThread(
		NULL,
		0,
		lpStartAddress,
		lpParameter,
		CREATE_SUSPENDED,
		&g_ManagedStack.ThreadId
	);
	if (!g_ManagedStack.ThreadHandle || !g_ManagedStack.ThreadId) {
		return GetLastError();
	}

	// Get the thread's context
	CONTEXT context = { 0 };
	context.ContextFlags = CONTEXT_ALL;
	BOOL bStatus = GetThreadContext(g_ManagedStack.ThreadHandle, &context);
	if (!bStatus) {
		return GetLastError();
	}

	// Get the thread's TIB
	PNT_TIB ntTib = NULL;
	DWORD dwStatus = GetTibByHandle(g_ManagedStack.ThreadHandle, &ntTib);
	if (dwStatus != ERROR_SUCCESS)
	{
		return dwStatus;
	}

	// Copy over any existing data for complete unwinding
	// Note: values within may still refer to the old stack
	g_ManagedStack.TIB = ntTib;
	PBYTE originalStackLimit = (PBYTE)ntTib->StackLimit;
	PBYTE originalStackBase = (PBYTE)ntTib->StackBase;
	PBYTE originalDeallocStack = *(PBYTE*)((PBYTE)ntTib + 0x1478);
	PBYTE originalRsp = (PBYTE)context.Rsp;
	PBYTE originalRbp = (PBYTE)context.Rbp;
	SIZE_T originalStackSize = originalStackBase - originalStackLimit;

	// Update the thread's TIB with the new values
	ntTib->StackLimit = g_ManagedStack.GuardUpperBound;
	ntTib->StackBase = g_ManagedStack.StackUpperBound;
	// Update the DeallocationStack; field not present in Windows headers
	*(PVOID*)((PBYTE)ntTib + 0x1478) = g_ManagedStack.StackUpperBound;

	// Register the exception handler
	AddVectoredExceptionHandler(1, ManagedStackExceptionHandler);

	// Update the thread's stack
	DWORD64 top = (DWORD64)g_ManagedStack.StackUpperBound;
	context.Rsp = top - 0x40 + ((DWORD64)originalRsp & 0xF);
	bStatus = SetThreadContext(g_ManagedStack.ThreadHandle, &context);
	if (!bStatus)
	{
		return GetLastError();
	}

	// Free the original stack
	bStatus = VirtualFree(originalDeallocStack, 0, MEM_RELEASE);
	if (!bStatus)
	{
		return GetLastError();
	}
	
	dwStatus = ResumeThread(g_ManagedStack.ThreadHandle);
	if (dwStatus == (DWORD)-1)
	{
		return GetLastError();
	}

	return ERROR_SUCCESS;
}

DWORD WINAPI TestThreadFn() {
	UCHAR stackTester[4096];
	for (INT i = 0; i < 4096; i++) {
		stackTester[i] = 'A';
	}
	return 0;
}

INT main()
{
	DWORD status = AllocManagedStack();
	if (status != ERROR_SUCCESS)
	{
		return status;
	}

	status = CreateThreadWithManagedStack((LPTHREAD_START_ROUTINE)TestThreadFn, NULL);
	WaitForSingleObject(g_ManagedStack.ThreadHandle, INFINITE);
	return status;
}