MFTPrefetchPrototype.cpp

#include <ntifs.h>
#include <ntddk.h>
#include <ntstrsafe.h>

#define MFTPF_DEV_SYMLINK_NAME L"MFTPrefetch"

#define MFTPF_DEVICE 0x8000

#define IOCTL_MFTPF_PREFETCH_MFT CTL_CODE (MFTPF_DEVICE, 0x800, METHOD_NEITHER, FILE_ANY_ACCESS)

const size_t ValidInputSizeCc = 1;
const size_t ValidInputSize = ValidInputSizeCc * sizeof (WCHAR);	// A simple drive letter

#define MFTPF_DRIVER_NAME "MFT prefetch driver"

namespace {

UNICODE_STRING SymLink = RTL_CONSTANT_STRING (LR"(\??\)" MFTPF_DEV_SYMLINK_NAME);
const ULONG PoolTag = 'FMFP';

template <typename T, size_t S>
constexpr size_t SizeOfArray (T (&)[S])
{
	return S;
}

template<typename T>
T RoundUp (T a, T b)
{
	return (((T)(a) +(T)(b) - 1) & ~((T)(b) - 1));
}

void LogWithStatus (const char* pMessage, NTSTATUS status)
{
	KdPrint ((MFTPF_DRIVER_NAME ": %s (status 0x08X)\n", pMessage, status));

	UNREFERENCED_PARAMETER (pMessage);
	UNREFERENCED_PARAMETER (status);
}

NTSTATUS OpenFile (UNICODE_STRING* pPath, HANDLE* pHandleOut)
{
	OBJECT_ATTRIBUTES objAttr;
	InitializeObjectAttributes (&objAttr, pPath, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);
	IO_STATUS_BLOCK ioBlock;

	NTSTATUS result  = ZwOpenFile (pHandleOut,
								   FILE_READ_ATTRIBUTES,
								   &objAttr,
								   &ioBlock,
								   FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
								   0);

	return result;
}

NTSTATUS PrefetchPages (PFILE_OBJECT pFileObject)
{
	LARGE_INTEGER fileSize;
	NTSTATUS result = FsRtlGetFileSize (pFileObject, &fileSize);
	if (!NT_SUCCESS (result))
		return result;

	const ULONG nPages = static_cast<ULONG> (RoundUp<UINT64> (fileSize.QuadPart, PAGE_SIZE) / PAGE_SIZE);

	READ_LIST* pReadList = static_cast<READ_LIST*> (ExAllocatePoolWithTag (PagedPool,
																		   FIELD_OFFSET (READ_LIST, List) + nPages * sizeof (FILE_SEGMENT_ELEMENT),
																		   PoolTag));

	if (pReadList == nullptr)
		return STATUS_NO_MEMORY;

	pReadList->FileObject = pFileObject;
	pReadList->IsImage = FALSE;
	pReadList->NumberOfEntries = nPages;

	for (ULONG offset = 0; offset < fileSize.QuadPart; offset += PAGE_SIZE)
		pReadList->List[offset / PAGE_SIZE].Alignment = offset;

	result = MmPrefetchPages (1, &pReadList);

	ExFreePoolWithTag (pReadList, PoolTag);

	return result;
}

extern "C" {

NTSTATUS CreateClose (PDEVICE_OBJECT /*DeviceObject*/, PIRP Irp)
{
	Irp->IoStatus.Status = STATUS_SUCCESS;
	Irp->IoStatus.Information = 0;

	IoCompleteRequest (Irp, IO_NO_INCREMENT);

	return STATUS_SUCCESS;
}

NTSTATUS DeviceControl (PDEVICE_OBJECT /*DeviceObject*/, PIRP Irp)
{
	KdPrint ((MFTPF_DRIVER_NAME " entered " __FUNCTION__ "\n"));

	PIO_STACK_LOCATION pStackLocation = IoGetCurrentIrpStackLocation (Irp);
	NTSTATUS result = STATUS_INVALID_DEVICE_REQUEST;

	switch (pStackLocation->Parameters.DeviceIoControl.IoControlCode) {
		case IOCTL_MFTPF_PREFETCH_MFT: {
			if (pStackLocation->Parameters.DeviceIoControl.InputBufferLength != ValidInputSize) {
				KdPrint ((MFTPF_DRIVER_NAME " recevied an input with invalid length\n"));

				result =  STATUS_INVALID_PARAMETER;

				break;
			}
			
			// So the input size is correct, let's copy, so it remains stable during the remainder of this function
			WCHAR driveLetter;
			static_assert (sizeof driveLetter == ValidInputSize);
			__try {
				RtlCopyMemory (&driveLetter, pStackLocation->Parameters.DeviceIoControl.Type3InputBuffer, ValidInputSize);
			} __except (EXCEPTION_EXECUTE_HANDLER) {
				KdPrint ((MFTPF_DRIVER_NAME " was unable to copy the input buffer\n"));

				result = STATUS_INVALID_PARAMETER;

				break;
			}

			// Construct MFT path
			UNICODE_STRING mftPath;
			const WCHAR MFTPathScheme[] = LR"(\??\X:\$Mft)";
			constexpr size_t MFTPathSize = SizeOfArray (MFTPathScheme);
			WCHAR mftPathBuffer[MFTPathSize];
			RtlInitUnicodeString (&mftPath, mftPathBuffer);
			mftPath.MaximumLength = sizeof mftPathBuffer;

			result = RtlUnicodeStringPrintf (&mftPath, LR"(\??\%C:\$Mft)", driveLetter);
			if (!NT_SUCCESS (result)) {
				KdPrint ((MFTPF_DRIVER_NAME " was unable to construct MFT path (drive: %C)\n", driveLetter));

				break;
			}

			HANDLE hMFT;
			if (!NT_SUCCESS (OpenFile (&mftPath, &hMFT))) {
				KdPrint ((MFTPF_DRIVER_NAME " was unable to open MFT at path: %wZ\n", mftPath));

				result = STATUS_NO_SUCH_FILE;

				break;
			}

			PFILE_OBJECT pMFTFileObject;
			result = ObReferenceObjectByHandle (hMFT,
												FILE_READ_ATTRIBUTES,
												*IoFileObjectType,
												KernelMode,
												reinterpret_cast <PVOID*> (&pMFTFileObject),
												nullptr);
			if (!NT_SUCCESS (result)) {
				KdPrint ((MFTPF_DRIVER_NAME " was unable to dereference MFT file object\n", driveLetter));

				break;
			}

			result = PrefetchPages (pMFTFileObject);

			//ObDereferenceObject (pMFTFileObject);

			//ZwClose (hMFT);

			if (NT_SUCCESS (result))
				KdPrint ((MFTPF_DRIVER_NAME " successfully completed request for %wZ\n", mftPath));
			else
				LogWithStatus ("failed to prefetch MFT", result);

			break;
		}
		default:
			KdPrint ((MFTPF_DRIVER_NAME " received an unsupported IOCTL code\n"));

			break;
	}

	Irp->IoStatus.Status = result;
	Irp->IoStatus.Information = 0;

	IoCompleteRequest (Irp, IO_NO_INCREMENT);

	return result;
}

void DriverUnload (PDRIVER_OBJECT DriverObject)
{
	IoDeleteSymbolicLink (&SymLink);

	IoDeleteDevice (DriverObject->DeviceObject);

	KdPrint ((MFTPF_DRIVER_NAME " unloaded\n"));
}

NTSTATUS DriverEntry (PDRIVER_OBJECT DriverObject, PUNICODE_STRING /*RegistryPath*/)
{
	KdPrint (("Loading " MFTPF_DRIVER_NAME "\n"));

	DriverObject->DriverUnload = DriverUnload;

	DriverObject->MajorFunction[IRP_MJ_CREATE] = CreateClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = CreateClose;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DeviceControl;

	PDEVICE_OBJECT pDevObj;
	UNICODE_STRING devName = RTL_CONSTANT_STRING (LR"(\Device\MFTPrefetch)");
	NTSTATUS status = IoCreateDevice (DriverObject, 0, &devName, FILE_DEVICE_UNKNOWN, 0, FALSE, &pDevObj);
	if (!NT_SUCCESS (status)) {
		LogWithStatus ("failed to create device object", status);

		return status;
	}

	status = IoCreateSymbolicLink (&SymLink, &devName);

	if (!NT_SUCCESS (status)) {
		LogWithStatus ("failed to create symbolic link", status);

		IoDeleteDevice (pDevObj);

		return status;
	}

	KdPrint ((MFTPF_DRIVER_NAME " loaded\n"));

	return STATUS_SUCCESS;
}

}	// extern "C"
}	// namespace