asd2fsm

#include "irp.h"
#include "ntdata.h"
#include "util.h"


NTSTATUS
    IoCompletionRoutine(
    __in PDEVICE_OBJECT DeviceObject,
    __in PIRP Irp,
    __in PVOID Context)
{
    *Irp->UserIosb = Irp->IoStatus;
    if (Irp->UserEvent != NULL) {
        KeSetEvent(Irp->UserEvent, IO_NO_INCREMENT, FALSE);
    }
    if (Irp->MdlAddress != NULL) {
        IoFreeMdl(Irp->MdlAddress);
        Irp->MdlAddress = NULL;
    }
    IoFreeIrp(Irp);
    return STATUS_MORE_PROCESSING_REQUIRED;
}


VOID
    IoCancelRoutine(
    __inout PDEVICE_OBJECT DeviceObject,
    __in PIRP Irp
    )
{
    Irp->IoStatus.Status = STATUS_CANCELLED;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, 0);
    IoReleaseCancelSpinLock(Irp->CancelIrql);
}


NTSTATUS
    GetDriveObject(
    __in PUNICODE_STRING FileName,
    __out PDEVICE_OBJECT *DeviceObject,
    __out PDEVICE_OBJECT *RealDevice)
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    HANDLE FileHandle = NULL;
    IO_STATUS_BLOCK IoStatusBlock = { 0 };
    UNICODE_STRING DeviceName = { 0 };
    OBJECT_ATTRIBUTES oa = { 0 };
    PFILE_OBJECT FileObject = NULL;
    WCHAR WcDeviceName[] = L"\\DosDevices\\*:\\";

    if (FileName->Length < 2 * sizeof(WCHAR)) {
        return STATUS_INVALID_PARAMETER;
    }

    WcDeviceName[12] = FileName->Buffer[0];
    RtlInitUnicodeString(&DeviceName, WcDeviceName);
    InitializeObjectAttributes(&oa, &DeviceName,
        OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, NULL, NULL);

    Status = IoCreateFile(
        &FileHandle,
        FILE_READ_ATTRIBUTES | SYNCHRONIZE,
        &oa,
        &IoStatusBlock,
        NULL,
        FILE_ATTRIBUTE_NORMAL,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        FILE_OPEN,
        FILE_SYNCHRONOUS_IO_NONALERT | FILE_DIRECTORY_FILE,
        NULL,
        0,
        CreateFileTypeNone,
        NULL,
        IO_NO_PARAMETER_CHECKING);
    if(!NT_SUCCESS(Status)) {
        YyPrint("Failed to open %C drive: 0x%08x\n", FileName->Buffer[0], Status);
        return Status;
    }

    Status = ObReferenceObjectByHandle(FileHandle, FILE_READ_ACCESS,
        *IoFileObjectType, KernelMode, (PVOID *)&FileObject, NULL);
    if(!NT_SUCCESS(Status)) {
        YyPrint("Failed to get PFILE_OBJECT: 0x%08x\n", Status);
        ZwClose(FileHandle);
        return Status;
    }

    // If FileObject has a mounted Vpb, use its DeviceObject.
    if (FileObject->Vpb != NULL && FileObject->Vpb->DeviceObject != NULL) {
        *DeviceObject = FileObject->Vpb->DeviceObject;
        *RealDevice = FileObject->Vpb->RealDevice;
#pragma prefast(suppress:28175, "Safe to use the 'Vpb' member of _DEVICE_OBJECT")
    } else if (!(FileObject->Flags & FO_DIRECT_DEVICE_OPEN) &&
        FileObject->DeviceObject->Vpb != NULL &&
        FileObject->DeviceObject->Vpb->DeviceObject != NULL) {
            // Otherwise, if the real device has a VPB that indicates that it
            // is mounted, then use the file system device object associated
            // with the VPB.
#pragma prefast(suppress:28175, "Safe to use the 'Vpb' member of _DEVICE_OBJECT")
            *DeviceObject = FileObject->DeviceObject->Vpb->DeviceObject;
#pragma prefast(suppress:28175, "Safe to use the 'Vpb' member of _DEVICE_OBJECT")
            *RealDevice = FileObject->DeviceObject->Vpb->RealDevice;
    } else {
        Status = STATUS_UNSUCCESSFUL;
    }

    ObDereferenceObject(FileObject);
    ZwClose(FileHandle);

    return Status;
}


NTSTATUS
    IrpCreateFile(
    __out PHANDLE Handle,
    __in PUNICODE_STRING FileName,
    __in ACCESS_MASK DesiredAccess,
    __out PIO_STATUS_BLOCK IoStatusBlock,
    __in ULONG FileAttributes,
    __in ULONG ShareAccess,
    __in ULONG CreateDisposition,
    __in ULONG CreateOptions)
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    OBJECT_ATTRIBUTES oa = { 0 };
    PDEVICE_OBJECT DeviceObject = NULL, RealDevice = NULL;
    PFILE_OBJECT FileObject = NULL;
    PIRP Irp = NULL;
    KEVENT SyncEvent = { 0 };
    PIO_STACK_LOCATION IrpSp = NULL;
    ACCESS_STATE AccessState = { 0 };
    AUX_ACCESS_DATA AuxData = { 0 };
    IO_SECURITY_CONTEXT SecurityContext = { 0 };

    Status = GetDriveObject(FileName, &DeviceObject, &RealDevice);
    if (!NT_SUCCESS(Status)) {
        YyPrint("Failed to get DeviceObject: 0x%08x\n", Status);
        return Status;
    }

    Irp = IoAllocateIrp(DeviceObject->StackSize, FALSE);
    if(Irp == NULL) {
        YyPrint("Failed to allocate IRP\n");
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    Status = SeCreateAccessState(&AccessState, &AuxData, DesiredAccess,
        IoGetFileObjectGenericMapping());
    if (!NT_SUCCESS(Status)) {
        YyPrint("SeCreateAccessState failed: 0x%08x\n", Status);
        IoFreeIrp(Irp);
        return Status;
    }

    InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, NULL, NULL);
    Status = ObCreateObject(KernelMode, *IoFileObjectType, &oa, KernelMode,
        NULL, sizeof(FILE_OBJECT), 0, 0, (PVOID *)&FileObject);
    if (!NT_SUCCESS(Status)) {
        YyPrint("Failed to create object: 0x%08x\n", Status);
        IoFreeIrp(Irp);
        return Status;
    }

    RtlZeroMemory(FileObject, sizeof(FILE_OBJECT));
    // FileName format: C:\Windows\system32
    // FileObject->FileName format: \Windows\system32
    FileName->Buffer += 2;
    FileName->Length -= 2 * sizeof(WCHAR);
    FileName->MaximumLength -= 2 * sizeof(WCHAR);
    if (FileName->Length > 0) {
        FileObject->FileName.Buffer = (PWCH)ExAllocatePoolWithTag(NonPagedPool,
            FileName->MaximumLength, 'fobj');
        if (FileObject->FileName.Buffer == NULL) {
            YyPrint("Failed to allocate space for FileObject->FileName.Buffer\n");
            IoFreeIrp(Irp);
            ObDereferenceObject(FileObject);
            return STATUS_INSUFFICIENT_RESOURCES;
        }
    } else {
        FileObject->FileName.MaximumLength = 0;
        FileObject->FileName.Buffer = NULL;
    }

    RtlCopyMemory(FileObject->FileName.Buffer, FileName->Buffer, FileName->Length);
    FileObject->FileName.Length = FileName->Length;
    FileObject->FileName.MaximumLength = FileName->MaximumLength;
    FileObject->Type = IO_TYPE_FILE;
    FileObject->Size = sizeof(FILE_OBJECT);
    // DeviceObject: C/D/E... drive
    FileObject->DeviceObject = RealDevice;
    FileObject->Flags = FO_SYNCHRONOUS_IO;
    KeInitializeEvent(&FileObject->Lock, SynchronizationEvent, FALSE);
    KeInitializeEvent(&FileObject->Event, NotificationEvent, FALSE);

    KeInitializeEvent(&SyncEvent, SynchronizationEvent, FALSE);

    SecurityContext.SecurityQos = NULL;
    SecurityContext.AccessState = &AccessState;
    SecurityContext.DesiredAccess = DesiredAccess;
    SecurityContext.FullCreateOptions = 0;

    Irp->MdlAddress = NULL;
    Irp->AssociatedIrp.SystemBuffer = NULL;
    Irp->Flags = IRP_CREATE_OPERATION | IRP_SYNCHRONOUS_API;
    Irp->RequestorMode = KernelMode;
    Irp->UserIosb = IoStatusBlock;
    Irp->UserEvent = &SyncEvent;
    Irp->PendingReturned = FALSE;
    Irp->Cancel = FALSE;
    Irp->CancelRoutine = NULL;
    Irp->Tail.Overlay.Thread = PsGetCurrentThread();
    Irp->Tail.Overlay.AuxiliaryBuffer = NULL;
    Irp->Tail.Overlay.OriginalFileObject = FileObject;

    IrpSp = IoGetNextIrpStackLocation(Irp);
    IrpSp->MajorFunction = IRP_MJ_CREATE;
    IrpSp->DeviceObject = DeviceObject;
    IrpSp->FileObject = FileObject;
    IrpSp->Parameters.Create.SecurityContext = &SecurityContext;
    IrpSp->Parameters.Create.Options = (CreateDisposition << 24) | CreateOptions;
    IrpSp->Parameters.Create.FileAttributes = (USHORT)FileAttributes;
    IrpSp->Parameters.Create.ShareAccess = (USHORT)ShareAccess;
    IrpSp->Parameters.Create.EaLength = 0;

    IoSetCompletionRoutine(Irp, IoCompletionRoutine, NULL, TRUE, TRUE, TRUE);
    IoSetCancelRoutine(Irp, IoCancelRoutine);
    Status = IoCallDriver(DeviceObject, Irp);
    if(Status == STATUS_PENDING) {
        KeWaitForSingleObject(&SyncEvent, Executive, KernelMode, FALSE, 0);
    }

    Status = IoStatusBlock->Status;

    if(!NT_SUCCESS(Status)) {
        YyPrint("Failed: 0x%08x\n", Status);
        // FileObject->DeviceObject can't be set to NULL before ObDereferenceObject
        // otherwise it won't call ExFreePoolWithTag to free FileObject->FileName.Buffer
        // which will lead to memory leak
        ObDereferenceObject(FileObject);
    } else {
#pragma prefast(suppress:28175, "Safe to use the 'ReferenceCount' member of _DEVICE_OBJECT")
        InterlockedIncrement((volatile LONG *)&FileObject->DeviceObject->ReferenceCount);
        if (FileObject->Vpb) {
            InterlockedIncrement((volatile LONG *)&FileObject->Vpb->ReferenceCount);
        }
        Status = ObInsertObject(FileObject, NULL, DesiredAccess, 0, NULL, Handle);
        if (!NT_SUCCESS(Status)) {
            YyPrintAPI("ObInsertObject");
            ObDereferenceObject(FileObject);
        }
    }

    return Status;
}


NTSTATUS
    IrpClose(
    __in PFILE_OBJECT FileObject)
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    IO_STATUS_BLOCK IoStatusBlock = { 0 };
    PIRP Irp = NULL;
    KEVENT SyncEvent = { 0 };
    PIO_STACK_LOCATION IrpSp = NULL;

    if (FileObject->Vpb == 0 || FileObject->Vpb->DeviceObject == NULL) {
        return STATUS_UNSUCCESSFUL;
    }

    Irp = IoAllocateIrp(FileObject->Vpb->DeviceObject->StackSize, FALSE);
    if(Irp == NULL) {
        YyPrint("Failed to allocate IRP\n");
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    KeInitializeEvent(&SyncEvent, SynchronizationEvent, FALSE);

    Irp->UserEvent = &SyncEvent;
    Irp->UserIosb = &IoStatusBlock;
    Irp->RequestorMode = KernelMode;
    Irp->Flags = IRP_CLOSE_OPERATION | IRP_SYNCHRONOUS_API;
    Irp->Tail.Overlay.Thread = PsGetCurrentThread();
    Irp->Tail.Overlay.OriginalFileObject = FileObject;

    IrpSp = IoGetNextIrpStackLocation(Irp);
    IrpSp->MajorFunction = IRP_MJ_CLEANUP;
    IrpSp->FileObject = FileObject;

    Status = IoCallDriver(FileObject->Vpb->DeviceObject, Irp);
    if (Status == STATUS_PENDING) {
        KeWaitForSingleObject(&SyncEvent, Executive, KernelMode, FALSE, NULL);
    }

    Status = IoStatusBlock.Status;
    if(!NT_SUCCESS(Status)) {
        YyPrint("Failed: 0x%08x\n", Status);
        IoFreeIrp(Irp);
        return Status;
    }

    KeClearEvent(&SyncEvent);
    IoReuseIrp(Irp , STATUS_SUCCESS);

    Irp->UserEvent = &SyncEvent;
    Irp->UserIosb = &IoStatusBlock;
    Irp->Tail.Overlay.OriginalFileObject = FileObject;
    Irp->Tail.Overlay.Thread = PsGetCurrentThread();
    Irp->AssociatedIrp.SystemBuffer = NULL;
    Irp->Flags = IRP_CLOSE_OPERATION | IRP_SYNCHRONOUS_API;

    IrpSp = IoGetNextIrpStackLocation(Irp);
    IrpSp->MajorFunction = IRP_MJ_CLOSE;
    IrpSp->FileObject = FileObject;

    if (FileObject->Vpb && !(FileObject->Flags & FO_DIRECT_DEVICE_OPEN)) {
        InterlockedDecrement((volatile LONG *)&FileObject->Vpb->ReferenceCount);
        FileObject->Flags |= FO_FILE_OPEN_CANCELLED;
    }

    Status = IoCallDriver(FileObject->Vpb->DeviceObject, Irp);
    if (Status == STATUS_PENDING) {
        KeWaitForSingleObject(&SyncEvent, Executive, KernelMode, FALSE, NULL);
    }

    IoFreeIrp(Irp);

    return IoStatusBlock.Status;
}


NTSTATUS
    IrpQueryDirectoryFile(
    __in PFILE_OBJECT FileObject,
    __out PIO_STATUS_BLOCK IoStatusBlock,
    __out PVOID FileInformation,
    __in ULONG Length,
    __in FILE_INFORMATION_CLASS FileInformationClass,
    __in BOOLEAN ReturnSingleEntry,
    __in_opt PUNICODE_STRING FileName,
    __in BOOLEAN RestartScan)
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    PIRP Irp = NULL;
    KEVENT SyncEvent = { 0 };
    PIO_STACK_LOCATION IrpSp = NULL;
    MODE Mode = KernelMode;
    PVOID KernelBuffer = NULL;

    if (FileObject->Vpb == NULL || FileObject->Vpb->DeviceObject == NULL) {
        return STATUS_UNSUCCESSFUL;
    }

    Irp = IoAllocateIrp(FileObject->Vpb->DeviceObject->StackSize, FALSE);
    if(Irp == NULL) {
        YyPrint("Failed to allocate IRP\n");
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    if ((ULONG)FileInformation < MmUserProbeAddress) {
        Mode = UserMode;
    }
    RtlZeroMemory(FileInformation, Length);

    if(FileObject->DeviceObject->Flags & DO_BUFFERED_IO) {
        if (Mode == UserMode) {
            KernelBuffer = ExAllocatePoolWithTag(NonPagedPool, Length, 'kmbf');
            if (KernelBuffer == NULL) {
                YyPrint("Failed to allocate memory for kernel mode buffer\n");
                IoFreeIrp(Irp);
                return STATUS_INSUFFICIENT_RESOURCES;
            }
            Irp->AssociatedIrp.SystemBuffer = KernelBuffer;
        } else {
            Irp->AssociatedIrp.SystemBuffer = FileInformation;
        }
    } else if(FileObject->DeviceObject->Flags & DO_DIRECT_IO) {
        Irp->MdlAddress = IoAllocateMdl(FileInformation, Length, FALSE, FALSE, NULL);
        if (Irp->MdlAddress == NULL) {
            YyPrint("Failed to allocate Mdl\n");
            IoFreeIrp(Irp);
            return STATUS_INSUFFICIENT_RESOURCES;
        }
        MmBuildMdlForNonPagedPool(Irp->MdlAddress);
    } else {
        Irp->UserBuffer = FileInformation;
    }

    KeInitializeEvent(&SyncEvent, SynchronizationEvent, FALSE);

    Irp->UserEvent = &SyncEvent;
    Irp->UserIosb = IoStatusBlock;
    Irp->Tail.Overlay.Thread = PsGetCurrentThread();
    Irp->Tail.Overlay.OriginalFileObject = FileObject;
    Irp->RequestorMode = Mode;

    IrpSp = IoGetNextIrpStackLocation(Irp);
    IrpSp->MajorFunction = IRP_MJ_DIRECTORY_CONTROL;
    IrpSp->MinorFunction = IRP_MN_QUERY_DIRECTORY;
    IrpSp->FileObject = FileObject;
    if (ReturnSingleEntry) {
        SetFlag(IrpSp->Flags, SL_RETURN_SINGLE_ENTRY);
    }
    if (RestartScan) {
        SetFlag(IrpSp->Flags, SL_RESTART_SCAN);
    }
    IrpSp->Parameters.QueryDirectory.Length = Length;
    IrpSp->Parameters.QueryDirectory.FileName = FileName;
    IrpSp->Parameters.QueryDirectory.FileInformationClass = FileInformationClass;

    IoSetCompletionRoutine(Irp, IoCompletionRoutine, NULL, TRUE, TRUE, TRUE);
    IoSetCancelRoutine(Irp, IoCancelRoutine);
    Status = IoCallDriver(FileObject->Vpb->DeviceObject, Irp);
    if (Status == STATUS_PENDING) {
        // if the alertable is TRUE, it will cause BSOD
        Status = KeWaitForSingleObject(&SyncEvent, Executive, KernelMode, FALSE, NULL);
        if (Status != STATUS_SUCCESS) {
            YyPrint("KeWaitForSingleObject: 0x%08x\n", Status);
        }
    }

    if (KernelBuffer != NULL) {
        if (NT_SUCCESS(IoStatusBlock->Status)) {
            RtlCopyMemory(FileInformation, KernelBuffer, Length);
        }
        ExFreePoolWithTag(KernelBuffer, 'kmbf');
    }

    return IoStatusBlock->Status;
}


NTSTATUS
    IrpQueryInformationFile(
    __in PFILE_OBJECT FileObject,
    __out PIO_STATUS_BLOCK IoStatusBlock,
    __out PVOID FileInformation,
    __in ULONG Length,
    __in FILE_INFORMATION_CLASS FileInformationClass)
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    PIRP Irp = NULL;
    KEVENT SyncEvent = { 0 };
    PIO_STACK_LOCATION IrpSp = NULL;

    if (FileObject->Vpb == 0 || FileObject->Vpb->DeviceObject == NULL) {
        return STATUS_UNSUCCESSFUL;
    }

    Irp = IoAllocateIrp(FileObject->Vpb->DeviceObject->StackSize, FALSE);
    if(Irp == NULL) {
        YyPrint("Failed to allocate IRP\n");
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    KeInitializeEvent(&SyncEvent, SynchronizationEvent, FALSE);

    RtlZeroMemory(FileInformation, Length);
    Irp->AssociatedIrp.SystemBuffer = FileInformation;
    Irp->UserEvent = &SyncEvent;
    Irp->UserIosb = IoStatusBlock;
    Irp->RequestorMode = KernelMode;
    Irp->Tail.Overlay.Thread = PsGetCurrentThread();
    Irp->Tail.Overlay.OriginalFileObject = FileObject;

    IrpSp = IoGetNextIrpStackLocation(Irp);
    IrpSp->MajorFunction = IRP_MJ_QUERY_INFORMATION;
    IrpSp->DeviceObject = FileObject->Vpb->DeviceObject;
    IrpSp->FileObject = FileObject;
    IrpSp->Parameters.QueryFile.Length = Length;
    IrpSp->Parameters.QueryFile.FileInformationClass = FileInformationClass;

    IoSetCompletionRoutine(Irp, IoCompletionRoutine, NULL, TRUE, TRUE, TRUE);
    IoSetCancelRoutine(Irp, IoCancelRoutine);
    Status = IoCallDriver(FileObject->Vpb->DeviceObject, Irp);
    if (Status == STATUS_PENDING) {
        KeWaitForSingleObject(&SyncEvent, Executive, KernelMode, FALSE, 0);
    }

    return IoStatusBlock->Status;
}


NTSTATUS
    IrpSetInformationFile(
    __in PFILE_OBJECT FileObject,
    __out PIO_STATUS_BLOCK IoStatusBlock,
    __in PVOID FileInformation,
    __in ULONG Length,
    __in FILE_INFORMATION_CLASS FileInformationClass,
    __in BOOLEAN ReplaceIfExists)
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    PIRP Irp = NULL;
    KEVENT SyncEvent = { 0 };
    PIO_STACK_LOCATION IrpSp = NULL;

    if (FileObject->Vpb == 0 || FileObject->Vpb->DeviceObject == NULL) {
        return STATUS_UNSUCCESSFUL;
    }

    Irp = IoAllocateIrp(FileObject->Vpb->DeviceObject->StackSize, FALSE);
    if(Irp == NULL) {
        YyPrint("Failed to allocate IRP\n");
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    KeInitializeEvent(&SyncEvent, SynchronizationEvent, FALSE);

    Irp->AssociatedIrp.SystemBuffer = FileInformation;
    Irp->UserEvent = &SyncEvent;
    Irp->UserIosb = IoStatusBlock;
    Irp->RequestorMode = KernelMode;
    Irp->Tail.Overlay.Thread = PsGetCurrentThread();
    Irp->Tail.Overlay.OriginalFileObject = FileObject;

    IrpSp = IoGetNextIrpStackLocation(Irp);
    IrpSp->MajorFunction = IRP_MJ_SET_INFORMATION;
    IrpSp->DeviceObject = FileObject->Vpb->DeviceObject;
    IrpSp->FileObject = FileObject;
    IrpSp->Parameters.SetFile.ReplaceIfExists = ReplaceIfExists;
    IrpSp->Parameters.SetFile.FileObject = FileObject;
    IrpSp->Parameters.SetFile.AdvanceOnly = FALSE;
    IrpSp->Parameters.SetFile.Length = Length;
    IrpSp->Parameters.SetFile.FileInformationClass = FileInformationClass;

    IoSetCompletionRoutine(Irp, IoCompletionRoutine, NULL, TRUE, TRUE, TRUE);
    IoSetCancelRoutine(Irp, IoCancelRoutine);
    Status = IoCallDriver(FileObject->Vpb->DeviceObject, Irp);
    if (Status == STATUS_PENDING) {
        KeWaitForSingleObject(&SyncEvent, Executive, KernelMode, FALSE, 0);
    }

    return IoStatusBlock->Status;
}


NTSTATUS
    IrpReadFile(
    __in PFILE_OBJECT FileObject,
    __out PIO_STATUS_BLOCK IoStatusBlock,
    __out PVOID Buffer,
    __in ULONG Length,
    __in_opt PLARGE_INTEGER ByteOffset)
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    PIRP Irp = NULL;
    KEVENT SyncEvent = { 0 };
    PIO_STACK_LOCATION IrpSp = NULL;
    MODE Mode = KernelMode;
    PVOID KernelBuffer = NULL;

    if (FileObject->Vpb == 0 || FileObject->Vpb->DeviceObject == NULL) {
        return STATUS_UNSUCCESSFUL;
    }

    if(ByteOffset == NULL) {
        if(!(FileObject->Flags & FO_SYNCHRONOUS_IO)) {
            return STATUS_INVALID_PARAMETER;
        }
        ByteOffset = &FileObject->CurrentByteOffset;
    }

    Irp = IoAllocateIrp(FileObject->Vpb->DeviceObject->StackSize, FALSE);
    if(Irp == NULL) {
        YyPrint("Failed to allocate IRP\n");
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    if ((ULONG)Buffer < MmUserProbeAddress) {
        Mode = UserMode;
    }
    RtlZeroMemory(Buffer, Length);

    if(FileObject->DeviceObject->Flags & DO_BUFFERED_IO) {
        if (Mode == UserMode) {
            KernelBuffer = ExAllocatePoolWithTag(NonPagedPool, Length, 'kmbf');
            if (KernelBuffer == NULL) {
                YyPrint("Failed to allocate memory for kernel mode buffer\n");
                IoFreeIrp(Irp);
                return STATUS_INSUFFICIENT_RESOURCES;
            }
            Irp->AssociatedIrp.SystemBuffer = KernelBuffer;
        } else {
            Irp->AssociatedIrp.SystemBuffer = Buffer;
        }
    } else if(FileObject->DeviceObject->Flags & DO_DIRECT_IO) {
        Irp->MdlAddress = IoAllocateMdl(Buffer, Length, FALSE, FALSE, NULL);
        if (Irp->MdlAddress == NULL) {
            YyPrint("Failed to allocate Mdl\n");
            IoFreeIrp(Irp);
            return STATUS_INSUFFICIENT_RESOURCES;
        }
        MmBuildMdlForNonPagedPool(Irp->MdlAddress);
    } else {
        Irp->UserBuffer = Buffer;
    }

    KeInitializeEvent(&SyncEvent, SynchronizationEvent, FALSE);

    Irp->UserEvent = &SyncEvent;
    Irp->UserIosb = IoStatusBlock;
    Irp->RequestorMode = Mode;
    Irp->Flags = IRP_READ_OPERATION;
    Irp->Tail.Overlay.Thread = PsGetCurrentThread();
    Irp->Tail.Overlay.OriginalFileObject = FileObject;

    IrpSp = IoGetNextIrpStackLocation(Irp);
    IrpSp->MajorFunction = IRP_MJ_READ;
    IrpSp->MinorFunction = IRP_MN_NORMAL;
    IrpSp->DeviceObject = FileObject->Vpb->DeviceObject;
    IrpSp->FileObject = FileObject;
    IrpSp->Parameters.Read.Length = Length;
    IrpSp->Parameters.Read.ByteOffset = *ByteOffset;

    IoSetCompletionRoutine(Irp, IoCompletionRoutine, NULL, TRUE, TRUE, TRUE);
    IoSetCancelRoutine(Irp, IoCancelRoutine);
    Status = IoCallDriver(FileObject->Vpb->DeviceObject, Irp);
    if (Status == STATUS_PENDING) {
        KeWaitForSingleObject(&SyncEvent, Executive, KernelMode, FALSE, 0);
    }

    if (KernelBuffer != NULL) {
        if (NT_SUCCESS(IoStatusBlock->Status)) {
            RtlCopyMemory(Buffer, KernelBuffer, Length);
        }
        ExFreePoolWithTag(KernelBuffer, 'kmbf');
    }

    return IoStatusBlock->Status;
}


NTSTATUS
    IrpWriteFile(
    __in PFILE_OBJECT FileObject,
    __out PIO_STATUS_BLOCK IoStatusBlock,
    __in PVOID Buffer,
    __in ULONG Length,
    __in_opt PLARGE_INTEGER ByteOffset)
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    PIRP Irp = NULL;
    KEVENT SyncEvent = { 0 };
    PIO_STACK_LOCATION IrpSp = NULL;
    MODE Mode = KernelMode;
    PVOID KernelBuffer = NULL;

    if (FileObject->Vpb == 0 || FileObject->Vpb->DeviceObject == NULL) {
        return STATUS_UNSUCCESSFUL;
    }

    if (ByteOffset == NULL) {
        if (!(FileObject->Flags & FO_SYNCHRONOUS_IO)) {
            return STATUS_INVALID_PARAMETER;
        }
        ByteOffset = &FileObject->CurrentByteOffset;
    }

    Irp = IoAllocateIrp(FileObject->Vpb->DeviceObject->StackSize, FALSE);
    if(Irp == NULL) {
        YyPrint("Failed to allocate IRP\n");
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    if ((ULONG)Buffer < MmUserProbeAddress) {
        Mode = UserMode;
    }

    if(FileObject->DeviceObject->Flags & DO_BUFFERED_IO) {
        if (Mode == UserMode) {
            KernelBuffer = ExAllocatePoolWithTag(NonPagedPool, Length, 'kmbf');
            if (KernelBuffer == NULL) {
                YyPrint("Failed to allocate memory for kernel mode buffer\n");
                IoFreeIrp(Irp);
                return STATUS_INSUFFICIENT_RESOURCES;
            }
            Irp->AssociatedIrp.SystemBuffer = KernelBuffer;
        } else {
            Irp->AssociatedIrp.SystemBuffer = Buffer;
        }
    } else if(FileObject->DeviceObject->Flags & DO_DIRECT_IO) {
        Irp->MdlAddress = IoAllocateMdl(Buffer, Length, FALSE, FALSE, NULL);
        if (Irp->MdlAddress == NULL) {
            YyPrint("Failed to allocate Mdl\n");
            IoFreeIrp(Irp);
            return STATUS_INSUFFICIENT_RESOURCES;
        }
        MmBuildMdlForNonPagedPool(Irp->MdlAddress);
    } else {
        Irp->UserBuffer = Buffer;
    }

    KeInitializeEvent(&SyncEvent, SynchronizationEvent, FALSE);

    Irp->UserEvent = &SyncEvent;
    Irp->UserIosb = IoStatusBlock;
    Irp->RequestorMode = Mode;
    Irp->Flags = IRP_WRITE_OPERATION;
    Irp->Tail.Overlay.Thread = PsGetCurrentThread();
    Irp->Tail.Overlay.OriginalFileObject = FileObject;

    IrpSp = IoGetNextIrpStackLocation(Irp);
    IrpSp->MajorFunction = IRP_MJ_WRITE;
    IrpSp->MinorFunction = IRP_MN_NORMAL;
    IrpSp->DeviceObject = FileObject->Vpb->DeviceObject;
    IrpSp->FileObject = FileObject;
    IrpSp->Parameters.Write.Length = Length;
    IrpSp->Parameters.Write.ByteOffset = *ByteOffset;

    IoSetCompletionRoutine(Irp, IoCompletionRoutine, NULL, TRUE, TRUE, TRUE);
    IoSetCancelRoutine(Irp, IoCancelRoutine);
    Status = IoCallDriver(FileObject->Vpb->DeviceObject, Irp);
    if (Status == STATUS_PENDING) {
        KeWaitForSingleObject(&SyncEvent, Executive, KernelMode, FALSE, NULL);
    }

    if (KernelBuffer != NULL) {
        if (NT_SUCCESS(IoStatusBlock->Status)) {
            RtlCopyMemory(Buffer, KernelBuffer, Length);
        }
        ExFreePoolWithTag(KernelBuffer, 'kmbf');
    }

    return IoStatusBlock->Status;
}