hook.cpp

#pragma warning(disable : 4201 4311 4302)
#include "hook.hpp"
#include "utils.hpp"

namespace KHook
{
	constexpr unsigned long kPoolTag  = 'VMON';
	constexpr unsigned long kPerfTag  = 'freP';

	constexpr unsigned long kMagicPre23606  = 0x501802u;
	constexpr unsigned long kMagicPost23606 = 0x601802u;
	constexpr unsigned short kMagicTag      = 0xF33;

	constexpr unsigned long kCounterTypeOffset    = 0xE4;
	constexpr unsigned long kCounterBaseRateOfs   = 0xC0;
	constexpr unsigned long kCounterTypePhysical  = 0x5;
	constexpr long long kCounterBaseRate          = 10000000LL;

	constexpr unsigned long long kQpcBiasVa = 0xFFFFF780000003B8ull;

	const GUID kCkclGuid = { 0x54dea73a, 0xed1f, 0x42a4, { 0xaf, 0x71, 0x3e, 0x63, 0xd0, 0x56, 0xf1, 0x74 } };

	InfinityCallbackPtr m_InfinityCallback = nullptr;
	unsigned long m_BuildNumber = 0;
	void *m_SystemCallTable = nullptr;
	bool m_DetectThreadStatus = true;

	void *m_EtwpDebuggerData = nullptr;
	void *m_CkclWmiLoggerContext = nullptr;
	void **m_EtwpDebuggerDataSilo = nullptr;
	void **m_GetCpuClock = nullptr;
	unsigned long long m_OriginalGetCpuClock = 0;

	PETHREAD m_DetectThreadObject = NULL;
	CLIENT_ID m_ClientId = {};

	unsigned long long m_HvlpReferenceTscPage = 0;
	unsigned long long m_HvlGetQpcBias = 0;
	unsigned long long m_HvlpGetReferenceTimeUsingTscPage = 0;
	unsigned long long m_HalpPerformanceCounter = 0;
	unsigned long long m_HalpOriginalPerformanceCounter = 0;
	unsigned long long m_HalpOriginalPerformanceCounterCopy = 0;
	unsigned long *m_HalpPerformanceCounterType = 0;
	unsigned char m_VmHalpPerformanceCounterType = 0;
	unsigned long m_OriginalHalpPerformanceCounterType = 0;
	unsigned long long m_OriginalHvlpGetReferenceTimeUsingTscPage = 0;
	HvlGetQpcBiasPtr m_OriginalHvlGetQpcBias = nullptr;

	PLONGLONG m_QpcPointer = NULL;
	PMDL m_QpcMdl = NULL;

	NTSTATUS EventTraceControl(ETWP_TRACE_TYPE type)
	{
		auto *property = (CKCL_TRACE_PROPERTIES *)ExAllocatePoolWithTag(NonPagedPool, PAGE_SIZE, kPoolTag);
		if (!property)
			return STATUS_MEMORY_NOT_ALLOCATED;

		auto *provider_name = (wchar_t *)ExAllocatePoolWithTag(NonPagedPool, 256 * sizeof(wchar_t), kPoolTag);
		if (!provider_name)
		{
			ExFreePoolWithTag(property, kPoolTag);
			return STATUS_MEMORY_NOT_ALLOCATED;
		}

		RtlZeroMemory(property, PAGE_SIZE);
		RtlZeroMemory(provider_name, 256 * sizeof(wchar_t));

		RtlCopyMemory(provider_name, L"Circular Kernel Context Logger", sizeof(L"Circular Kernel Context Logger"));
		RtlInitUnicodeString(&property->ProviderName, provider_name);

		property->Wnode.BufferSize = PAGE_SIZE;
		property->Wnode.Flags = WNODE_FLAG_TRACED_GUID;
		property->Wnode.Guid = kCkclGuid;
		property->Wnode.ClientContext = 3;
		property->BufferSize = sizeof(unsigned long);
		property->MinimumBuffers = 2;
		property->MaximumBuffers = 2;
		property->LogFileMode = EVENT_TRACE_BUFFERING_MODE;

		unsigned long length = 0;
		if (type == EtwpUpdateTrace)
			property->EnableFlags = EVENT_TRACE_FLAG_SYSTEMCALL;

		NTSTATUS status = NtTraceControl(type, property, PAGE_SIZE, property, PAGE_SIZE, &length);

		ExFreePoolWithTag(provider_name, kPoolTag);
		ExFreePoolWithTag(property, kPoolTag);

		return status;
	}

	unsigned long long SelfGetCpuClock()
	{
		if (ExGetPreviousMode() == KernelMode)
			return __rdtsc();

		auto thread = (PKTHREAD)__readgsqword(0x188);

		unsigned int call_index = 0;
		if (m_BuildNumber <= 7601)
			call_index = *(unsigned int *)((unsigned long long)thread + 0x1f8);
		else
			call_index = *(unsigned int *)((unsigned long long)thread + 0x80);

		auto **stack_top = (void **)__readgsqword(0x1a8);
		auto **stack_frame = (void **)_AddressOfReturnAddress();

		for (auto **current = stack_top; current > stack_frame; --current)
		{
			auto *marker = (unsigned long *)current;
			if (*marker != kMagicPre23606 && *marker != kMagicPost23606)
				continue;

			--current;

			auto *tag = (unsigned short *)current;
			if (*tag != kMagicTag)
				continue;

			for (; current < stack_top; ++current)
			{
				auto *value = (unsigned long long *)current;
				if (!(PAGE_ALIGN(*value) >= m_SystemCallTable &&
					PAGE_ALIGN(*value) < (void *)((unsigned long long)m_SystemCallTable + (PAGE_SIZE * 2))))
					continue;

				void **syscall_fn = &current[9];
				if (m_InfinityCallback)
					m_InfinityCallback(call_index, syscall_fn);

				break;
			}

			break;
		}

		return __rdtsc();
	}

	EXTERN_C __int64 FakeHvlGetQpcBias()
	{
		SelfGetCpuClock();

		if (*(unsigned long long *)m_HvlpReferenceTscPage != 0)
			return *((unsigned long long *)(*(unsigned long long *)m_HvlpReferenceTscPage) + 3);

		return 0;
	}

	ULONG64 FakeGetReferenceTimeUsingTscPage()
	{
		return __rdtsc();
	}

	void DetectThreadRoutine(void *)
	{
		while (m_DetectThreadStatus)
		{
			KUtils::Sleep(1000);

			if (m_BuildNumber <= 18363)
			{
				if (MmIsAddressValid(m_GetCpuClock) && MmIsAddressValid(*m_GetCpuClock))
				{
					if (SelfGetCpuClock != *m_GetCpuClock)
					{
						if (Initialize(m_InfinityCallback))
							Start();
					}
				}
				else
				{
					Initialize(m_InfinityCallback);
				}
			}

			KeQueryPerformanceCounter(NULL);
		}

		PsTerminateSystemThread(STATUS_SUCCESS);
	}

	bool Initialize(InfinityCallbackPtr callback)
	{
		if (!m_DetectThreadStatus) return false;
		if (!MmIsAddressValid(callback)) return false;
		m_InfinityCallback = callback;

		if (!NT_SUCCESS(EventTraceControl(EtwpUpdateTrace)))
		{
			if (!NT_SUCCESS(EventTraceControl(EtwpStartTrace)))
				return false;
			if (!NT_SUCCESS(EventTraceControl(EtwpUpdateTrace)))
				return false;
		}

		m_BuildNumber = KUtils::GetSystemBuildNumber();
		if (!m_BuildNumber) return false;

		unsigned long long ntoskrnl = KUtils::GetModuleAddress("ntoskrnl.exe", nullptr);
		if (!ntoskrnl) return false;

		unsigned long long etw_data = KUtils::FindPatternImage(ntoskrnl, "\x00\x00\x2c\x08\x04\x38\x0c", "??xxxxx", ".text");
		if (!etw_data) etw_data = KUtils::FindPatternImage(ntoskrnl, "\x00\x00\x2c\x08\x04\x38\x0c", "??xxxxx", ".data");
		if (!etw_data) etw_data = KUtils::FindPatternImage(ntoskrnl, "\x00\x00\x2c\x08\x04\x38\x0c", "??xxxxx", ".rdata");
		if (!etw_data) return false;
		m_EtwpDebuggerData = (void *)etw_data;

		m_EtwpDebuggerDataSilo = *(void ***)((unsigned long long)m_EtwpDebuggerData + 0x10);
		if (!m_EtwpDebuggerDataSilo) return false;

		m_CkclWmiLoggerContext = m_EtwpDebuggerDataSilo[0x2];
		if (!m_CkclWmiLoggerContext) return false;

		if (m_BuildNumber <= 7601 || m_BuildNumber >= 22000)
			m_GetCpuClock = (void **)((unsigned long long)m_CkclWmiLoggerContext + 0x18);
		else
			m_GetCpuClock = (void **)((unsigned long long)m_CkclWmiLoggerContext + 0x28);

		if (!MmIsAddressValid(m_GetCpuClock)) return false;

		m_SystemCallTable = PAGE_ALIGN(KUtils::GetSyscallEntry(ntoskrnl));
		if (!m_SystemCallTable) return false;

		if (m_BuildNumber > 18363)
		{
			unsigned long long addr = 0;

			//
			// HvlpReferenceTscPage
			//
			addr = KUtils::FindPatternImage(ntoskrnl,
				"\x48\x8b\x05\x00\x00\x00\x00\x48\x8b\x40\x00\x48\x8b\x0d\x00\x00\x00\x00\x48\xf7\xe2",
				"xxx????xxx?xxx????xxx");
			if (!addr) return false;
			m_HvlpReferenceTscPage = KUtils::ResolveRipRelative(addr);
			if (!m_HvlpReferenceTscPage) return false;

			//
			// HvlGetQpcBias
			//
			addr = KUtils::FindPatternImage(ntoskrnl,
				"\x48\x8b\x05\x00\x00\x00\x00\x48\x85\xc0\x74\x00\x48\x83\x3d\x00\x00\x00\x00\x00\x74",
				"xxx????xxxx?xxx?????x");
			if (!addr)
			{
				addr = KUtils::FindPatternImage(ntoskrnl,
					"\x48\x8b\x05\x00\x00\x00\x00\xe8\x00\x00\x00\x00\x48\x03\xd8\x48\x89\x1f",
					"xxx????x????xxxxxx");
			}
			if (!addr) return false;
			m_HvlGetQpcBias = KUtils::ResolveRipRelative(addr);
			if (!m_HvlGetQpcBias) return false;

			//
			// HvlpGetReferenceTimeUsingTscPage
			//
			addr = KUtils::FindPatternImage(ntoskrnl,
				"\x48\x8b\x05\x00\x00\x00\x00\x48\x85\xc0\x74\x00\x33\xc9\xe8\x00\x00\x00\x00\x48\x8b\xd8",
				"xxx????xxxx?xxx????xxx");
			if (!addr)
			{
				addr = KUtils::FindPatternImage(ntoskrnl,
					"\x48\x8b\x05\x00\x00\x00\x00\xE8\x00\x00\x00\x00\x48\x03\xd8",
					"xxx????x????xxx");
			}
			if (!addr) return false;
			m_HvlpGetReferenceTimeUsingTscPage = KUtils::ResolveRipRelative(addr);
			if (!m_HvlpGetReferenceTimeUsingTscPage) return false;

			//
			// HalpPerformanceCounter
			//
			unsigned long long halp_counter_scan = KUtils::FindPatternImage(ntoskrnl,
				"\x48\x8b\x05\x00\x00\x00\x00\x48\x8b\xf9\x48\x85\xc0\x74\x00\x83\xb8",
				"xxx????xxxxxxx?xx");
			if (!halp_counter_scan) return false;
			m_HalpPerformanceCounter = KUtils::ResolveRipRelative(halp_counter_scan);
			if (!m_HalpPerformanceCounter) return false;

			//
			// HalpOriginalPerformanceCounter
			//
			addr = KUtils::FindPatternImage(ntoskrnl,
				"\x48\x8b\x05\x00\x00\x00\x00\x48\x3b\x00\x0f\x85\x00\x00\x00\x00\xA0",
				"xxx????xx?xx????x");
			if (!addr)
			{
				addr = KUtils::FindPatternImage(ntoskrnl,
					"\x48\x8b\x0d\x00\x00\x00\x00\x4c\x00\x00\x00\x00\x48\x3b\xf1",
					"xxx????x????xxx");
			}
			if (!addr) return false;
			m_HalpOriginalPerformanceCounter = KUtils::ResolveRipRelative(addr);
			if (!m_HalpOriginalPerformanceCounter) return false;

			m_HalpPerformanceCounterType = (ULONG *)((ULONG_PTR)(*(PVOID *)m_HalpPerformanceCounter) + kCounterTypeOffset);
			if (!m_HalpPerformanceCounterType) return false;

			if (*m_HalpPerformanceCounterType == kCounterTypePhysical)
			{
				/*++

				Physical machine path. The cmp operand at byte 21 of the
				HalpPerformanceCounter scan tells us what counter type
				the kernel expects for the virtualized path (7 or 8).

				--*/
				m_VmHalpPerformanceCounterType = *((char *)halp_counter_scan + 21);

				m_HalpOriginalPerformanceCounterCopy = (ULONGLONG)ExAllocatePoolWithTag(NonPagedPool, 0xFF, kPerfTag);
				if (!m_HalpOriginalPerformanceCounterCopy) return false;
				RtlZeroMemory((PVOID)m_HalpOriginalPerformanceCounterCopy, 0xFF);
				*(PULONGLONG)(m_HalpOriginalPerformanceCounterCopy + kCounterBaseRateOfs) = kCounterBaseRate;
				*(PULONG)(m_HalpOriginalPerformanceCounterCopy + kCounterTypeOffset) = kCounterTypePhysical;

				m_QpcMdl = IoAllocateMdl((PLONGLONG)kQpcBiasVa, 8, false, false, NULL);
				if (!m_QpcMdl) return false;

				MmBuildMdlForNonPagedPool(m_QpcMdl);
				m_QpcPointer = (PLONGLONG)MmMapLockedPagesSpecifyCache(
					m_QpcMdl, KernelMode, MmWriteCombined, NULL, false, NormalPagePriority);
				if (!m_QpcPointer) return false;
			}
		}

		return true;
	}

	bool Start()
	{
		if (!m_InfinityCallback) return false;
		if (!MmIsAddressValid(m_GetCpuClock)) return false;

		m_OriginalGetCpuClock = (unsigned long long)(*m_GetCpuClock);

		if (m_BuildNumber <= 18363)
		{
			*m_GetCpuClock = SelfGetCpuClock;
		}
		else
		{
			/*++

			Setting GetCpuClock to 2 forces the kernel down the
			HalpTimerQueryHostPerformanceCounter path which calls
			HvlGetQpcBias, our actual interception point.

			--*/
			*m_GetCpuClock = (void *)2;

			m_OriginalHvlGetQpcBias = (HvlGetQpcBiasPtr)(*(unsigned long long *)m_HvlGetQpcBias);

			if (m_HvlpGetReferenceTimeUsingTscPage)
			{
				m_OriginalHvlpGetReferenceTimeUsingTscPage = *(unsigned long long *)m_HvlpGetReferenceTimeUsingTscPage;
				if (m_OriginalHvlpGetReferenceTimeUsingTscPage == 0)
					*(unsigned long long *)m_HvlpGetReferenceTimeUsingTscPage = (ULONGLONG)FakeGetReferenceTimeUsingTscPage;
			}

			m_OriginalHalpPerformanceCounterType = *m_HalpPerformanceCounterType;
			if (*m_HalpPerformanceCounterType == kCounterTypePhysical)
			{
				*(unsigned long long *)m_HalpOriginalPerformanceCounter = m_HalpOriginalPerformanceCounterCopy;
				KeQueryPerformanceCounter(NULL);
				*m_HalpPerformanceCounterType = m_VmHalpPerformanceCounterType;
				KeQueryPerformanceCounter(NULL);
			}

			*(unsigned long long *)m_HvlGetQpcBias = (unsigned long long)FakeHvlGetQpcBias;
		}

		static bool thread_created = false;
		if (!thread_created)
		{
			thread_created = true;

			OBJECT_ATTRIBUTES oa{};
			HANDLE thread = NULL;
			InitializeObjectAttributes(&oa, 0, OBJ_KERNEL_HANDLE, 0, 0);

			NTSTATUS status = PsCreateSystemThread(&thread, THREAD_ALL_ACCESS, &oa, 0, &m_ClientId, DetectThreadRoutine, 0);
			if (NT_SUCCESS(status))
				ObReferenceObjectByHandle(thread, THREAD_ALL_ACCESS, NULL, KernelMode, (PVOID *)&m_DetectThreadObject, NULL);
		}

		return true;
	}

	bool Stop()
	{
		m_DetectThreadStatus = false;

		bool result = NT_SUCCESS(EventTraceControl(EtwpStopTrace)) &&
		              NT_SUCCESS(EventTraceControl(EtwpStartTrace));

		if (m_DetectThreadObject)
		{
			KeWaitForSingleObject(m_DetectThreadObject, Executive, KernelMode, false, NULL);

			UNICODE_STRING name = RTL_CONSTANT_STRING(L"ObfDereferenceObject");
			auto fn = (ObfDereferenceObjectPtr)MmGetSystemRoutineAddress(&name);
			if (fn)
			{
				fn(m_DetectThreadObject);
			}
			else
			{
				name = RTL_CONSTANT_STRING(L"ObDereferenceObject");
				auto fn_fallback = (ObDereferenceObjectPtr)MmGetSystemRoutineAddress(&name);
				if (fn_fallback) fn_fallback(m_DetectThreadObject);
			}
		}

		*m_GetCpuClock = (void *)m_OriginalGetCpuClock;

		if (m_BuildNumber > 18363)
		{
			if (m_HvlpGetReferenceTimeUsingTscPage && m_OriginalHvlpGetReferenceTimeUsingTscPage == 0)
				*(unsigned long long *)m_HvlpGetReferenceTimeUsingTscPage = m_OriginalHvlpGetReferenceTimeUsingTscPage;

			if (m_OriginalHalpPerformanceCounterType == kCounterTypePhysical)
			{
				LARGE_INTEGER before = KeQueryPerformanceCounter(NULL);
				*m_HalpPerformanceCounterType = m_OriginalHalpPerformanceCounterType;


				LARGE_INTEGER after = KeQueryPerformanceCounter(NULL);
				if (after.QuadPart - before.QuadPart > kCounterBaseRate)
					*m_QpcPointer -= (after.QuadPart - before.QuadPart);

				if (m_QpcMdl)
				{
					IoFreeMdl(m_QpcMdl);
					m_QpcMdl = NULL;
				}
			}

			*(unsigned long long *)m_HvlGetQpcBias = (unsigned long long)m_OriginalHvlGetQpcBias;
		}

		return result;
	}
}