Extremely hacky solution to close leaked handles in mscms.dll!OpenDisplay

icm_handle_leak_fix.cpp

// see: https://bugs.chromium.org/p/chromium/issues/detail?id=1201106
// see: https://twitter.com/gsuberland/status/1445547814965055488

#include <Windows.h>
#include <stdio.h>
#include <TlHelp32.h>
#include <memory>
#include <cassert>
#include <vector>

#pragma comment(lib, "ntdll")

#define NT_SUCCESS(status) (status >= 0)

#define STATUS_INFO_LENGTH_MISMATCH      ((NTSTATUS)0xC0000004L)

enum PROCESSINFOCLASS {
	ProcessHandleInformation = 51
};

typedef struct _PROCESS_HANDLE_TABLE_ENTRY_INFO
{
	HANDLE HandleValue;
	ULONG_PTR HandleCount;
	ULONG_PTR PointerCount;
	ULONG GrantedAccess;
	ULONG ObjectTypeIndex;
	ULONG HandleAttributes;
	ULONG Reserved;
} PROCESS_HANDLE_TABLE_ENTRY_INFO, * PPROCESS_HANDLE_TABLE_ENTRY_INFO;

typedef struct _PROCESS_HANDLE_SNAPSHOT_INFORMATION
{
	ULONG_PTR NumberOfHandles;
	ULONG_PTR Reserved;
	PROCESS_HANDLE_TABLE_ENTRY_INFO Handles[1];
} PROCESS_HANDLE_SNAPSHOT_INFORMATION, * PPROCESS_HANDLE_SNAPSHOT_INFORMATION;

extern "C" NTSTATUS NTAPI NtQueryInformationProcess(
	_In_ HANDLE ProcessHandle,
	_In_ PROCESSINFOCLASS ProcessInformationClass,
	_Out_writes_bytes_(ProcessInformationLength) PVOID ProcessInformation,
	_In_ ULONG ProcessInformationLength,
	_Out_opt_ PULONG ReturnLength);

typedef struct _UNICODE_STRING
{
	USHORT Length;
	USHORT MaximumLength;
	PWSTR  Buffer;
} UNICODE_STRING;

typedef enum _OBJECT_INFORMATION_CLASS
{
	ObjectBasicInformation = 0, // q: OBJECT_BASIC_INFORMATION
	ObjectNameInformation = 1, // q: OBJECT_NAME_INFORMATION
	ObjectTypeInformation = 2, // q: OBJECT_TYPE_INFORMATION
	ObjectTypesInformation = 3, // q: OBJECT_TYPES_INFORMATION
	ObjectHandleFlagInformation = 4, // qs: OBJECT_HANDLE_FLAG_INFORMATION
	ObjectSessionInformation = 5, // s: void // change object session // (requires SeTcbPrivilege)
	ObjectSessionObjectInformation = 6, // s: void // change object session // (requires SeTcbPrivilege)
} OBJECT_INFORMATION_CLASS;

typedef struct _OBJECT_TYPE_INFORMATION
{
	UNICODE_STRING TypeName;
	ULONG TotalNumberOfObjects;
	ULONG TotalNumberOfHandles;
	ULONG TotalPagedPoolUsage;
	ULONG TotalNonPagedPoolUsage;
	ULONG TotalNamePoolUsage;
	ULONG TotalHandleTableUsage;
	ULONG HighWaterNumberOfObjects;
	ULONG HighWaterNumberOfHandles;
	ULONG HighWaterPagedPoolUsage;
	ULONG HighWaterNonPagedPoolUsage;
	ULONG HighWaterNamePoolUsage;
	ULONG HighWaterHandleTableUsage;
	ULONG InvalidAttributes;
	GENERIC_MAPPING GenericMapping;
	ULONG ValidAccessMask;
	BOOLEAN SecurityRequired;
	BOOLEAN MaintainHandleCount;
	UCHAR TypeIndex; // since WINBLUE
	CHAR ReservedByte;
	ULONG PoolType;
	ULONG DefaultPagedPoolCharge;
	ULONG DefaultNonPagedPoolCharge;
} OBJECT_TYPE_INFORMATION, * POBJECT_TYPE_INFORMATION;

typedef struct _OBJECT_TYPES_INFORMATION
{
	ULONG NumberOfTypes;
	OBJECT_TYPE_INFORMATION Types[1];
} OBJECT_TYPES_INFORMATION, * POBJECT_TYPES_INFORMATION;


typedef struct _OBJECT_NAME_INFORMATION
{
	UNICODE_STRING Name;
} OBJECT_NAME_INFORMATION, * POBJECT_NAME_INFORMATION;

extern "C" NTSTATUS NTAPI NtQueryObject(
	_In_opt_ HANDLE Handle,
	_In_ OBJECT_INFORMATION_CLASS ObjectInformationClass,
	_Out_writes_bytes_opt_(ObjectInformationLength) PVOID ObjectInformation,
	_In_ ULONG ObjectInformationLength,
	_Out_opt_ PULONG ReturnLength);


UCHAR RegistryKeyHandleType = 0xFF;

void CleanupHandles(const DWORD pid)
{
	if (pid == GetCurrentProcessId())
	{
		printf("Skipping own process.\n");
		return;
	}

	// open a handle to the process. use a shared_ptr so that CloseHandle is always called when exiting
	const std::shared_ptr<void> hProcess(OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_DUP_HANDLE, FALSE, pid), CloseHandle);

	if (hProcess.get() == nullptr)
	{
		printf("CleanupHandles: Failed to open handle to process ID %u. Error: %u\n", pid, GetLastError());
		return;
	}

	// get the handle table for this process
	ULONG bufferSize = 0x1000;
	std::vector<BYTE> buffer(bufferSize);
	NTSTATUS status;
	while (status = NtQueryInformationProcess(hProcess.get(), ProcessHandleInformation, buffer.data(), bufferSize, &bufferSize) == STATUS_INFO_LENGTH_MISMATCH)
	{
		bufferSize *= 2;

		// fail if we're resizing the buffer to something very large.
		if (bufferSize > 1024 * 1024 * 4)
		{
			printf("CleanupHandles: NtQueryObject requested buffer exceeding 4MB. Buffer size: %u\n", bufferSize);
			return;
		}

		buffer.resize(bufferSize);
	}

	// if NtQueryInformationProcess failed with something other than STATUS_INFO_LENGTH_MISMATCH, something went wrong
	if (!NT_SUCCESS(status))
	{
		printf("CleanupHandles: NtQueryInformationProcess call failed for process %u. Status: %x. Last error: %u\n", pid, status, GetLastError());
		return;
	}

	// allocate a buffer that can store a OBJECT_NAME_INFORMATION structure with a path of the maximum allowed length
	// this is the size of the struct (effectively just a UNICODE_STRING) plus the max size of a UNICODE_STRING's contents in bytes, plus one wchar_t null terminator
	// this buffer used to store the name of the object that the handle is pointing to
	const ULONG nameBufferSize = sizeof(OBJECT_NAME_INFORMATION) + UNICODE_STRING_MAX_BYTES + sizeof(wchar_t);
	auto nameBuffer = std::vector<BYTE>(nameBufferSize);

	const auto info = reinterpret_cast<PROCESS_HANDLE_SNAPSHOT_INFORMATION*>(buffer.data());
	ULONG closeCount = 0;
	for (ULONG i = 0; i < info->NumberOfHandles; i++)
	{
		// zero the name buffer before use (will contain previous path on second loop)
		std::fill(nameBuffer.begin(), nameBuffer.end(), 0);

		const HANDLE h = info->Handles[i].HandleValue;

		// if we successfully found the registry key handle type, check that this handle matches that type
		if (RegistryKeyHandleType != 0xFF)
		{
			if (info->Handles[i].ObjectTypeIndex != RegistryKeyHandleType)
			{
				continue;
			}
		}

		// duplicate the handle into our process, so we can get its name information
		// we use DUPLICATE_SAME_ACCESS at this point because we're just checking the handle's path; we don't know that we want DUPLICATE_CLOSE_SOURCE yet
		HANDLE hTarget;
		if (!DuplicateHandle(hProcess.get(), h, GetCurrentProcess(), &hTarget, 0, FALSE, DUPLICATE_SAME_ACCESS))
			continue;

		// grab the name information
		ULONG resultLength = 0;
		const auto status = NtQueryObject(hTarget, ObjectNameInformation, nameBuffer.data(), nameBufferSize, &resultLength);

		// close the duplicated handle. we'll re-open it with DUPLICATE_CLOSE_SOURCE if we determine that it's one of the leaked handles
		CloseHandle(hTarget);

		if (!NT_SUCCESS(status))
			continue;

		if (resultLength < sizeof(OBJECT_NAME_INFORMATION))
		{
			printf("CleanupHandles: NtQueryObject call for ObjectNameInformation returned a buffer of size %u, which is smaller than the size of OBJECT_NAME_INFORMATION (%zu)\n", resultLength, sizeof(OBJECT_NAME_INFORMATION));
			continue;
		}

		if (resultLength > nameBuffer.size())
		{
			printf("CleanupHandles: NtQueryObject call for ObjectNameInformation returned a buffer of size %u, which is larger than the maximum buffer size (%u)\n", resultLength, nameBufferSize);
			continue;
		}

		const POBJECT_NAME_INFORMATION objName = reinterpret_cast<POBJECT_NAME_INFORMATION>(nameBuffer.data());

		// sometimes handles will point to objects without names. we only care about ones with names
		if (objName->Name.Buffer != nullptr && objName->Name.Length > 0)
		{
			// verify that the length makes sense
			const auto maxAllowedNameLength = nameBuffer.size() - sizeof(OBJECT_NAME_INFORMATION);
			if (objName->Name.Length > maxAllowedNameLength)
			{
				continue;
			}

			// is the handle pointing to the registry key of interest?
			if (_wcsnicmp(L"\\REGISTRY\\MACHINE\\SYSTEM\\ControlSet001\\Control\\Class\\{4d36e96e-e325-11ce-bfc1-08002be10318}", objName->Name.Buffer, objName->Name.Length / sizeof(wchar_t)) == 0)
			{
				//printf("Closing handle %p - %ws\n", hTarget, objName->Name.Buffer);
				if (DuplicateHandle(hProcess.get(), h, GetCurrentProcess(), &hTarget, 0, FALSE, DUPLICATE_CLOSE_SOURCE))
				{
					CloseHandle(hTarget);
					closeCount++;
				}
			}
		}
	}

	printf("Closed %u handles.\n", closeCount);
}

void EnumerateProcessesForCleanup()
{
	const HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE)
	{
		printf("Failed to create process list snapshot. Error: %u\n", GetLastError());
		return;
	}

	PROCESSENTRY32 pe = { 0 };
	pe.dwSize = sizeof(pe);

	// skip the idle process
	Process32First(hSnapshot, &pe);

	// loop through all processes and cleanup handles in the ones of interest
	while (Process32Next(hSnapshot, &pe))
	{
		if (_wcsicmp(L"chrome.exe", pe.szExeFile) == 0 ||
			_wcsicmp(L"typora.exe", pe.szExeFile) == 0 ||
			_wcsicmp(L"slack.exe", pe.szExeFile) == 0 ||
			_wcsicmp(L"whatsapp.exe", pe.szExeFile) == 0 ||
			_wcsicmp(L"mspaint.exe", pe.szExeFile) == 0 ||
			_wcsicmp(L"spotify.exe", pe.szExeFile) == 0 ||
			_wcsicmp(L"telegram.exe", pe.szExeFile) == 0 
			)
		{
			printf("Handling process %u - %ws\n", pe.th32ProcessID, pe.szExeFile);
			CleanupHandles(pe.th32ProcessID);
		}
	}

	CloseHandle(hSnapshot);
}

bool GetRegistryKeyObjectTypeIndex(PUCHAR pIndexOut)
{
	NTSTATUS status;
	ULONG returnLength;

	ULONG bufferSize = 0x1000;
	std::vector<BYTE> buffer(bufferSize);

	while ((status = NtQueryObject(
		NULL,
		OBJECT_INFORMATION_CLASS::ObjectTypesInformation,
		buffer.data(),
		bufferSize,
		&returnLength)) == STATUS_INFO_LENGTH_MISMATCH)
	{
		bufferSize *= 2;

		// Fail if we're resizing the buffer to something very large.
		if (bufferSize > 1024 * 1024 * 4)
		{
			printf("GetRegistryKeyObjectTypeIndex: NtQueryObject requested buffer exceeding 4MB. Buffer size: %u\n", bufferSize);
			return false;
		}

		buffer.resize(bufferSize);
	}

	if (!NT_SUCCESS(status))
	{
		printf("GetRegistryKeyObjectTypeIndex: NtQueryObject failed with status %u. Last error: %u\n", status, GetLastError());
		return false;
	}

	const auto objectTypes = reinterpret_cast<POBJECT_TYPES_INFORMATION>(buffer.data());
	PUCHAR ptr = reinterpret_cast<PUCHAR>(&objectTypes->Types);
	printf("%u types found.\n", objectTypes->NumberOfTypes);
	for (ULONG i = 0; i < objectTypes->NumberOfTypes; i++)
	{
		auto objectType = reinterpret_cast<POBJECT_TYPE_INFORMATION>(ptr);
		//printf("Type %u name %ws\n", objectType->TypeIndex, objectType->TypeName.Buffer);
		if (_wcsicmp(L"Key", objectType->TypeName.Buffer) == 0)
		{
			printf("Registry key type is %u\n", objectType->TypeIndex);
			*pIndexOut = objectType->TypeIndex;
			return true;
		}

		ptr += sizeof(OBJECT_TYPE_INFORMATION);
		ptr += objectType->TypeName.Length + sizeof(wchar_t);
		if ((size_t)ptr % 8 != 0)
			ptr += 8 - ((size_t)ptr % 8);
	}

	printf("Exhausted search without finding registry key type.\n");
	return false;
}

int main()
{
	if (!GetRegistryKeyObjectTypeIndex(&RegistryKeyHandleType))
	{
		printf("Couldn't find registry key object type index.\n");
	}

	while (true)
	{
		printf("Clearing handles...\n");
		EnumerateProcessesForCleanup();
		printf("Waiting...\n\n");
		Sleep(60 * 1000);
	}

	return 0;
}