Hook.cpp

#include "pch.h"
#include <Detours/detours.h>

#include <string>
#include <string_view>
#include <sstream>
#include <array>
#include <vector>

#define NtCurrentProcess() ((HANDLE)(LONG_PTR)-1)
#define NtCurrentThread() ((HANDLE)(LONG_PTR)-2)
#define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)
#define NT_INFORMATION(Status) ((((ULONG)(Status)) >> 30) == 1)
#define NT_WARNING(Status) ((((ULONG)(Status)) >> 30) == 2)
#define NT_ERROR(Status) ((((ULONG)(Status)) >> 30) == 3)

typedef struct _PEB_PARTIAL
{
	BOOLEAN InheritedAddressSpace;
	BOOLEAN ReadImageFileExecOptions;
	BOOLEAN BeingDebugged;
	union
	{
		BOOLEAN BitField;
		struct
		{
			BOOLEAN ImageUsesLargePages : 1;
			BOOLEAN IsProtectedProcess : 1;
			BOOLEAN IsImageDynamicallyRelocated : 1;
			BOOLEAN SkipPatchingUser32Forwarders : 1;
			BOOLEAN IsPackagedProcess : 1;
			BOOLEAN IsAppContainer : 1;
			BOOLEAN IsProtectedProcessLight : 1;
			BOOLEAN IsLongPathAwareProcess : 1;
		};
	};

	HANDLE Mutant;

	PVOID ImageBaseAddress;
	PPEB_LDR_DATA Ldr;
	PRTL_USER_PROCESS_PARAMETERS ProcessParameters;
	PVOID SubSystemData;
	PVOID ProcessHeap;
} PEB_PARTIAL, * PPEB_PARTIAL;

#define NtCurrentPeb() ((PPEB_PARTIAL)(NtCurrentTeb()->ProcessEnvironmentBlock))
#define RtlProcessHeap() (NtCurrentPeb()->ProcessHeap)


typedef NTSTATUS (__fastcall *tfpBcdEditStringToDevice)(unsigned int DataType, const wchar_t* String, unsigned int* Flags, void** Data, DWORD* DataSize);
typedef NTSTATUS (__fastcall *tfpBcdEditGetDeviceStringOld)(void* ObjectHandle, _BCDE_DEVICE* Device, DWORD DeviceSize, wchar_t** DeviceString);
typedef NTSTATUS (__fastcall *tfpBcdEditGetDeviceStringNew)(void* ObjectHandle, _BCDE_DEVICE* Device, DWORD DeviceSize, bool Unknown, wchar_t** DeviceString);
typedef NTSTATUS (__fastcall *tfpBcdEditParseCommandLine)(DWORD ArgumentCount, wchar_t** Arguments);
typedef NTSTATUS (__fastcall* tfpBiConvertNtDeviceToBootEnvironment)(_BCDE_DEVICE* Device, DWORD DeviceSize, DWORD Flags, _BOOT_ENVIRONMENT_DEVICE** BootDevice);
typedef size_t (WINAPI *tfpEntryPoint)(VOID);

static tfpBcdEditStringToDevice g_BcdEditStringToDevice = nullptr;
static tfpBcdEditGetDeviceStringOld g_BcdEditGetDeviceStringOld = nullptr;
static tfpBcdEditGetDeviceStringNew g_BcdEditGetDeviceStringNew = nullptr;
static tfpBcdEditParseCommandLine g_BcdEditParseCommandLine = nullptr;
static tfpBiConvertNtDeviceToBootEnvironment g_BiConvertNtDeviceToBootEnvironment = nullptr;
static tfpEntryPoint g_EntryPoint = nullptr;

static NTSTATUS StringToDevice(unsigned int DataType, const wchar_t* String, unsigned int* Flags, void** Data, DWORD* DataSize) {
	try {
		std::wstring_view WString = &String[5];

		std::vector<BYTE> device;
		auto BcdeDevice = Extend<BCDE_DEVICE>(device, offsetof(BCDE_DEVICE, Body));
		BcdeDevice->DeviceType = BCDE_TYPE_UNKNOWN;
		auto BootDeviceOffset = device.size();
		if (!DeviceParse(WString, device)) return STATUS_INVALID_PARAMETER;
		// optional extra unknown[] followed by optional AdditionalOptions GUID
		if (!WString.empty()) {
			if (!CharParse(WString, L',')) return STATUS_INVALID_PARAMETER;
			bool ExpectingGuid = true;
			if (WString._Starts_with(TypeUnknown)) {
				WString = WString.substr(TypeUnknown.length());
				if (!WString.empty() && WString[0] == L':') {
					// Instead of a hex string, this is a path to a file.
					if (!CharParse(WString, L':')) return STATUS_INVALID_PARAMETER;
					auto PathString = SubStringUntil(WString, L':', 1 + (WString.size() >= 2 && WString[1] == L':'));
					if (PathString.empty()) return false;
					WString = WString.substr(PathString.size() + 1);
					// Open the file.
					FILE* File = nullptr;
					if (_wfopen_s(&File, std::wstring(PathString).c_str(), L"rb") != 0) return STATUS_INVALID_PARAMETER;
					if (File == nullptr) return STATUS_INVALID_PARAMETER;
					// Get the size of the file.
					fseek(File, 0, SEEK_END);
					auto FileSize = (SIZE_T)_ftelli64(File);
					fseek(File, 0, SEEK_SET);
					// Allocate space for the extra data.
					auto ExtraData = Extend<BYTE>(device, FileSize);
					// Read the file.
					FileSize = fread(ExtraData, sizeof(*ExtraData), FileSize, File);
					// Close the file.
					fclose(File);
					if (!CharParse(WString, L']')) return STATUS_INVALID_PARAMETER;
				}
				else {
					if (!HexToBytesUntil(WString, L']', device)) return STATUS_INVALID_PARAMETER;
				}
				// Fix the length of the boot environment device.
				auto BootDevice = (PBOOT_ENVIRONMENT_DEVICE_HEADER)&device[BootDeviceOffset];
				BootDevice->Size = (DWORD)(device.size() - BootDeviceOffset);
				// If there's characters after this, then expect a comma followed by a GUID.
				ExpectingGuid = !WString.empty();
				if (ExpectingGuid && !CharParse(WString, L',')) return STATUS_INVALID_PARAMETER;
			}
			BcdeDevice = (PBCDE_DEVICE)&device[0];
			if (ExpectingGuid) {
				if (!GuidParse(WString, BcdeDevice->AdditionalOptions)) return STATUS_INVALID_PARAMETER;
				WString = WString.substr(ExampleGuidString.length());
				if (!WString.empty()) return STATUS_INVALID_PARAMETER;
			}
		}
		*DataSize = (DWORD)device.size();
		PVOID allocatedBuffer = RtlAllocateHeap(RtlProcessHeap(), 0, device.size());
		if (allocatedBuffer == nullptr) return STATUS_INSUFFICIENT_RESOURCES;
		memcpy(allocatedBuffer, &device[0], device.size());
		*Data = allocatedBuffer;
		return STATUS_SUCCESS;
	}
	catch (...) {
		return STATUS_UNSUCCESSFUL;
	}
}


static NTSTATUS __fastcall HookStringToDevice(unsigned int DataType, const wchar_t* String, unsigned int* Flags, void** Data, DWORD* DataSize) {
	if (wcsncmp(String, L"!raw:", 5) != 0) {
		return g_BcdEditStringToDevice(DataType, String, Flags, Data, DataSize);
	}

	return StringToDevice(DataType, String, Flags, Data, DataSize);
}



static NTSTATUS GetDeviceStringForUnknown(void* ObjectHandle, _BCDE_DEVICE* BcdeDevice, DWORD DeviceSize, wchar_t** DeviceString) {
	// Length check, device size must be at least header size.
	if (DeviceSize < offsetof(BCDE_DEVICE, Body.Unknown.Data) + sizeof(BOOT_ENVIRONMENT_DEVICE_HEADER)) return STATUS_INVALID_PARAMETER;

	PBOOT_ENVIRONMENT_DEVICE Device = (PBOOT_ENVIRONMENT_DEVICE)BcdeDevice->Body.Unknown.Data;
	// Can't trust the size from the structure.
	Device->Size = DeviceSize - offsetof(BCDE_DEVICE, Body.Unknown.Data);

	std::wstringstream ss;
	ss << L"!raw:" << DeviceToString(Device);
	if (InlineIsEqualGUID(NullGuid, BcdeDevice->AdditionalOptions) == FALSE) {
		ss << L',';
		wchar_t szGUID[64] = { 0 };
#pragma warning(suppress : 6031) // 64 wchars enough for a GUID, in fact all that is needed is 35 wchars.
		StringFromGUID2(BcdeDevice->AdditionalOptions, szGUID, sizeof(szGUID) / sizeof(*szGUID));
		ss << szGUID;
	}
	auto outDevice = ss.str();
	auto stringSize = outDevice.size() + 1;
	auto AllocatedString = (wchar_t*)RtlAllocateHeap(RtlProcessHeap(), 0, stringSize * sizeof(**DeviceString));
	if (AllocatedString == nullptr) return STATUS_INSUFFICIENT_RESOURCES;
	wcscpy_s(AllocatedString, stringSize, outDevice.c_str());
	*DeviceString = AllocatedString;
	return STATUS_SUCCESS;
}

static NTSTATUS __fastcall HookGetDeviceStringNew(void* ObjectHandle, _BCDE_DEVICE* BcdeDevice, DWORD DeviceSize, bool UnknownBool, wchar_t** DeviceString) {
	if (BcdeDevice->DeviceType != BCDE_TYPE_UNKNOWN) {
		return g_BcdEditGetDeviceStringNew(ObjectHandle, BcdeDevice, DeviceSize, UnknownBool, DeviceString);
	}

	return GetDeviceStringForUnknown(ObjectHandle, BcdeDevice, DeviceSize, DeviceString);
}

static NTSTATUS __fastcall HookGetDeviceStringOld(void* ObjectHandle, _BCDE_DEVICE* BcdeDevice, DWORD DeviceSize, wchar_t** DeviceString) {
	if (BcdeDevice->DeviceType != BCDE_TYPE_UNKNOWN) {
		return g_BcdEditGetDeviceStringOld(ObjectHandle, BcdeDevice, DeviceSize, DeviceString);
	}

	return GetDeviceStringForUnknown(ObjectHandle, BcdeDevice, DeviceSize, DeviceString);
}

static std::wstring_view IntegerParameterFileContents = L"";

static NTSTATUS ConvertToWString(UINT CodePage, PBYTE Data, SIZE_T Length, wchar_t*& Converted, SIZE_T& ConvertedLength) {
	if (Length > MAXINT32) return STATUS_INTEGER_OVERFLOW;
	int LengthS32 = (int)Length;
	auto result = MultiByteToWideChar(CodePage, 0, (LPCCH)Data, LengthS32, nullptr, 0);
	if (result <= 0) return STATUS_INVALID_PARAMETER;

	auto WString = (wchar_t*)malloc(((SIZE_T)result) * sizeof(wchar_t) + sizeof(L"x "));
	if (WString == nullptr) return STATUS_NO_MEMORY;

	memcpy(WString, L"x ", sizeof(L"x ") - sizeof(wchar_t));

	result = MultiByteToWideChar(CodePage, 0, (LPCCH)Data, LengthS32, &WString[2], result);
	if (result <= 0) {
		free(WString);
		return STATUS_INVALID_PARAMETER;
	}
	Converted = WString;
	ConvertedLength = (result - sizeof(*Converted)) / sizeof(*Converted);
	return STATUS_SUCCESS;
}

// Hooked to allow using @file argument like MS developer tools.
static NTSTATUS __fastcall HookParseCommandLine(DWORD ArgumentCount, wchar_t* Arguments[]) {
	if (ArgumentCount == 2 && Arguments[1][0] == L'@') {
		// Open the file.
		FILE* File = nullptr;
		if (_wfopen_s(&File, &Arguments[1][1], L"rb") != 0) return STATUS_INVALID_PARAMETER;
		if (File == nullptr) return STATUS_INVALID_PARAMETER;
		// Get the size of the file.
		fseek(File, 0, SEEK_END);
		auto FileSize = (SIZE_T)_ftelli64(File);
		fseek(File, 0, SEEK_SET);
		// Allocate space for the extra data.
		auto IntegerParameterFile = (BYTE*)malloc(FileSize + sizeof(L"x "));
		if (IntegerParameterFile == nullptr) return STATUS_NO_MEMORY;
		// Read the file.
		FileSize = fread(IntegerParameterFile, 1, FileSize, File);
		IntegerParameterFile[FileSize] = 0;
		IntegerParameterFile[FileSize + 1] = 0;
		// Close the file.
		fclose(File);
		// Starts with a BOM?
		bool NeedsConvert = !(FileSize >= 2 && IntegerParameterFile[0] == 0xFF && IntegerParameterFile[1] == 0xFE);
		if (!NeedsConvert) {
			memmove(&IntegerParameterFile[4], IntegerParameterFile, FileSize + 2);
			memcpy(IntegerParameterFile, L"x ", sizeof(L"x ") - sizeof(wchar_t));
			Arguments = CommandLineToArgvW((LPCWSTR)&IntegerParameterFile[2], (int*)&ArgumentCount);
			free(IntegerParameterFile);
		}
		else {
			wchar_t* IntegerParameterFileW = nullptr;
			SIZE_T FileSizeW = 0;
			NTSTATUS Status = STATUS_SUCCESS;
			if (FileSize >= 3 && IntegerParameterFile[0] == 0xEF && IntegerParameterFile[1] == 0xBB && IntegerParameterFile[2] == 0xBF) {
				// UTF-8
				Status = ConvertToWString(CP_UTF8, IntegerParameterFile, FileSize, IntegerParameterFileW, FileSizeW);
			}
			else if (FileSize >= 2 && IntegerParameterFile[0] == 0xFE && IntegerParameterFile[1] == 0xFF) {
				// UTF-16 big endian, not supported
				Status = STATUS_NOT_SUPPORTED;
			}
			else {
				Status = ConvertToWString(CP_ACP, IntegerParameterFile, FileSize, IntegerParameterFileW, FileSizeW);
			}
			free(IntegerParameterFile);
			if (!NT_SUCCESS(Status)) return Status;
			Arguments = CommandLineToArgvW(IntegerParameterFileW, (int*)&ArgumentCount);
		}
	}
	return g_BcdEditParseCommandLine(ArgumentCount, Arguments);
}

// Hooked to bypass the 1MB limit on BOOT_ENVIRONMENT_DEVICE size, which does not exist in bootmgr.
static NTSTATUS __fastcall HookConvertNtDeviceToBootEnvironment(_BCDE_DEVICE* Device, DWORD DeviceSize, DWORD Flags, _BOOT_ENVIRONMENT_DEVICE** BootDevice) {
	if (Device->DeviceType != BCDE_TYPE_UNKNOWN) return g_BiConvertNtDeviceToBootEnvironment(Device, DeviceSize, Flags, BootDevice);

	const DWORD ExpectedSize = offsetof(BCDE_DEVICE, Body) + offsetof(BOOT_ENVIRONMENT_DEVICE_HEADER, Pad);
	if (DeviceSize < ExpectedSize) return STATUS_INVALID_PARAMETER;

	auto UnknownDevice = (PBOOT_ENVIRONMENT_DEVICE)Device->Body.Unknown.Data;
	auto NewDevice = (PBOOT_ENVIRONMENT_DEVICE)RtlAllocateHeap(NtCurrentPeb()->ProcessHeap, 0, UnknownDevice->Size);
	if (NewDevice == nullptr) return STATUS_INSUFFICIENT_RESOURCES;
	memcpy(NewDevice, UnknownDevice, UnknownDevice->Size);
	*BootDevice = NewDevice;
	return STATUS_SUCCESS;
}

bool HookOnInit() {
	{
#pragma warning(suppress : 4996) // copying the env variable
		std::wstring path = L"srv*" + std::wstring(_wgetenv(L"TEMP")) + L"\\symbols*https://msdl.microsoft.com/download/symbols";
		SymSetOptions(SYMOPT_DEFERRED_LOADS | SYMOPT_EXACT_SYMBOLS);
		if (SymInitializeW(NtCurrentProcess(), path.c_str(), TRUE) == FALSE) return false;
	}

	ULONG64 buffersym[(sizeof(SYMBOL_INFOW) + MAX_SYM_NAME * sizeof(WCHAR) + sizeof(ULONG64) - 1) / sizeof(ULONG64)];
	PSYMBOL_INFOW sym = (PSYMBOL_INFOW)buffersym;

	memset(buffersym, 0, sizeof(buffersym));
	sym->SizeOfStruct = sizeof(SYMBOL_INFOW);
	sym->MaxNameLen = MAX_SYM_NAME;

	if (SymFromNameW(NtCurrentProcess(), L"bcdedit!BcdEditStringToDevice", sym) == FALSE) {
		return false;
	}
	g_BcdEditStringToDevice = (tfpBcdEditStringToDevice)(PVOID)sym->Address;

	memset(buffersym, 0, sizeof(buffersym));
	sym->SizeOfStruct = sizeof(SYMBOL_INFOW);
	sym->MaxNameLen = MAX_SYM_NAME;

	if (SymFromNameW(NtCurrentProcess(), L"bcdedit!BcdEditParseCommandLine", sym) == FALSE) {
		return false;
	}
	g_BcdEditParseCommandLine = (tfpBcdEditParseCommandLine)(PVOID)sym->Address;

	memset(buffersym, 0, sizeof(buffersym));
	sym->SizeOfStruct = sizeof(SYMBOL_INFOW);
	sym->MaxNameLen = MAX_SYM_NAME;

	if (SymFromNameW(NtCurrentProcess(), L"bcdedit!BiConvertNtDeviceToBootEnvironment", sym) == FALSE) {
		return false;
	}
	g_BiConvertNtDeviceToBootEnvironment = (tfpBiConvertNtDeviceToBootEnvironment)(PVOID)sym->Address;

	memset(buffersym, 0, sizeof(buffersym));
	sym->SizeOfStruct = sizeof(SYMBOL_INFOW);
	sym->MaxNameLen = MAX_SYM_NAME;

	if (SymFromNameW(NtCurrentProcess(), L"bcdedit!BcdEditGetDeviceString", sym) == FALSE) return false;

	// BcdEditGetDeviceString prototype changed at RS3.
	_Notnull_ auto hBcdedit = GetModuleHandleW(nullptr);
	auto hResInfo = FindResourceW(hBcdedit, MAKEINTRESOURCEW(100), RT_VERSION);
	if (hResInfo == nullptr) hResInfo = FindResourceW(hBcdedit, MAKEINTRESOURCEW(1), RT_VERSION); 
	if (hResInfo == nullptr) {
		printf("FindResource failed %d\n", GetLastError());
		return false;
	}
	auto dwSize = SizeofResource(hBcdedit, hResInfo);
	auto hResData = LoadResource(hBcdedit, hResInfo);
	if (hResData == nullptr) {
		printf("LoadResource failed %d\n", GetLastError());
		return false;
	}
	auto pRes = LockResource(hResData);
	auto pResCopy = std::make_unique<BYTE[]>(dwSize);
	memcpy(pResCopy.get(), pRes, dwSize);
	FreeResource(hResData);

	VS_FIXEDFILEINFO* lpFfi = (VS_FIXEDFILEINFO*)(pResCopy.get() + 0x28);
	WORD buildNumber = (lpFfi->dwProductVersionLS >> 16);
	if (buildNumber < 16299) {
		g_BcdEditGetDeviceStringOld = (tfpBcdEditGetDeviceStringOld)(PVOID)sym->Address;
	}
	else {
		g_BcdEditGetDeviceStringNew = (tfpBcdEditGetDeviceStringNew)(PVOID)sym->Address;
	}

	DetourTransactionBegin();
	DetourUpdateThread(NtCurrentThread());
	DetourAttach((PVOID*)&g_BcdEditStringToDevice, HookStringToDevice);
	DetourAttach((PVOID*)&g_BcdEditParseCommandLine, HookParseCommandLine);
	DetourAttach((PVOID*)&g_BiConvertNtDeviceToBootEnvironment, HookConvertNtDeviceToBootEnvironment);
	if (buildNumber < 16299) {
		DetourAttach((PVOID*)&g_BcdEditGetDeviceStringOld , HookGetDeviceStringOld);
	}
	else {
		DetourAttach((PVOID*)&g_BcdEditGetDeviceStringNew, HookGetDeviceStringNew);
	}
	DetourTransactionCommit();

#if 0
	MessageBoxW(nullptr, L"", L"Waiting for attach", 0);
#endif
	return true;
}

extern "C" __declspec(dllexport) void _() {}