Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
#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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment