VOID DispTdiQueryInformationExComplete(
    PVOID Context,
    ULONG Status,
    UINT ByteCount)
/*
 * FUNCTION: Completes a TDI QueryInformationEx request
 * ARGUMENTS:
 *     Context   = Pointer to the IRP for the request
 *     Status    = TDI status of the request
 *     ByteCount = Number of bytes returned in output buffer
 */
{
    PTI_QUERY_CONTEXT QueryContext;

    QueryContext = (PTI_QUERY_CONTEXT)Context;
    if (NT_SUCCESS(Status)) {
        CopyBufferToBufferChain(
            QueryContext->InputMdl,
            FIELD_OFFSET(TCP_REQUEST_QUERY_INFORMATION_EX, Context),
            (PCHAR)&QueryContext->QueryInfo.Context,
            CONTEXT_SIZE);
    }

    MmUnlockPages(QueryContext->InputMdl);
    IoFreeMdl(QueryContext->InputMdl);
    if( QueryContext->OutputMdl ) {
	MmUnlockPages(QueryContext->OutputMdl);
	IoFreeMdl(QueryContext->OutputMdl);
    }

    QueryContext->Irp->IoStatus.Information = ByteCount;
    QueryContext->Irp->IoStatus.Status      = Status;

    ExFreePoolWithTag(QueryContext, QUERY_CONTEXT_TAG);
}

VOID
NTAPI
CdRomDeAllocateMmcResources(
    IN PDEVICE_OBJECT Fdo
    )
{
    PCOMMON_DEVICE_EXTENSION commonExtension = Fdo->DeviceExtension;
    PCDROM_DATA cddata = commonExtension->DriverData;
    PCDROM_MMC_EXTENSION mmcData = &cddata->Mmc;
    //NTSTATUS status;

    if (mmcData->CapabilitiesWorkItem) {
        IoFreeWorkItem(mmcData->CapabilitiesWorkItem);
        mmcData->CapabilitiesWorkItem = NULL;
    }
    if (mmcData->CapabilitiesIrp) {
        IoFreeIrp(mmcData->CapabilitiesIrp);
        mmcData->CapabilitiesIrp = NULL;
    }
    if (mmcData->CapabilitiesMdl) {
        IoFreeMdl(mmcData->CapabilitiesMdl);
        mmcData->CapabilitiesMdl = NULL;
    }
    if (mmcData->CapabilitiesBuffer) {
        ExFreePool(mmcData->CapabilitiesBuffer);
        mmcData->CapabilitiesBuffer = NULL;
    }
    mmcData->CapabilitiesBuffer = 0;
    mmcData->IsMmc = FALSE;
    mmcData->WriteAllowed = FALSE;
    
    return;
}

/*
 *  DestroyTransferPacket
 *
 */
VOID DestroyTransferPacket(PTRANSFER_PACKET Pkt)
{
    PFUNCTIONAL_DEVICE_EXTENSION fdoExt = Pkt->Fdo->DeviceExtension;
    PCLASS_PRIVATE_FDO_DATA fdoData = fdoExt->PrivateFdoData;
    KIRQL oldIrql;

    ASSERT(!Pkt->SlistEntry.Next);
//    ASSERT(!Pkt->OriginalIrp);

    KeAcquireSpinLock(&fdoData->SpinLock, &oldIrql);

    /*
     *  Delete the packet from our all-packets queue.
     */
    ASSERT(!IsListEmpty(&Pkt->AllPktsListEntry));
    ASSERT(!IsListEmpty(&fdoData->AllTransferPacketsList));
    RemoveEntryList(&Pkt->AllPktsListEntry);
    InitializeListHead(&Pkt->AllPktsListEntry);

    KeReleaseSpinLock(&fdoData->SpinLock, oldIrql);

    IoFreeMdl(Pkt->PartialMdl);
    IoFreeIrp(Pkt->Irp);
    FREE_POOL(Pkt->RetryHistory);
    FREE_POOL(Pkt);
}

NTSTATUS
	IoCompletionRoutine(
	IN PDEVICE_OBJECT  DeviceObject,
	IN PIRP  Irp,
	IN PVOID  Context
	)
{
	UNREFERENCED_PARAMETER(DeviceObject);
	UNREFERENCED_PARAMETER(Context);

	PAGED_CODE();

	KdPrint(("SYS(%d:%d):IoCompletionRoutine!\n", PsGetCurrentProcessId(), PsGetCurrentThreadId()));
	*Irp->UserIosb = Irp->IoStatus;

	if (Irp->UserEvent)
		KeSetEvent(Irp->UserEvent, IO_NO_INCREMENT, 0);

	if (Irp->MdlAddress)
	{
		IoFreeMdl(Irp->MdlAddress);
		Irp->MdlAddress = NULL;
	}

	IoFreeIrp(Irp);

	return STATUS_MORE_PROCESSING_REQUIRED;
}

NTSTATUS
IoCompletionRoutine(
	IN PDEVICE_OBJECT  DeviceObject,
	IN PIRP  Irp,
	IN PVOID  Context
	)
{
	//DbgPrint(("IoCompletionRoutine!\n"));
	*Irp->UserIosb = Irp->IoStatus;


	if (Irp->MdlAddress)
	{
		//MmUnmapLockedPages( Irp->MdlAddress, 
		//	MmGetSystemAddressForMdlSafe( Irp->MdlAddress, NormalPagePriority ) );
		MmUnlockPages( Irp->MdlAddress );
		IoFreeMdl(Irp->MdlAddress);
		Irp->MdlAddress = NULL;
	}

	if (Irp->UserEvent)
		KeSetEvent(Irp->UserEvent, IO_NO_INCREMENT, 0);

	IoFreeIrp(Irp);

	return STATUS_MORE_PROCESSING_REQUIRED;
	//return STATUS_SUCCESS;
}

/// <summary>
/// Unmap memory region, release corresponding MDL, and remove region form list
/// </summary>
/// <param name="pPageEntry">Region data</param>
/// <param name="pFoundEntry">Process data</param>
/// <returns>Status code</returns>
NTSTATUS BBUnmapRegionEntry( IN PMAP_ENTRY pPageEntry, IN PPROCESS_MAP_ENTRY pFoundEntry )
{
    NTSTATUS status = STATUS_SUCCESS;

    UNREFERENCED_PARAMETER( pFoundEntry );

    // MDL is valid
    if (pPageEntry->pMdl)
    {
        // If MDL is mapped
        if (pPageEntry->newPtr)
        {
            DPRINT( "BlackBone: %s: Unmapping region at 0x%p from process %u\n", __FUNCTION__, pPageEntry->newPtr, pFoundEntry->target.pid );
            MmUnmapLockedPages( (PVOID)pPageEntry->newPtr, pPageEntry->pMdl );
            pPageEntry->newPtr = 0;
        }

        if (pPageEntry->locked)
            MmUnlockPages( pPageEntry->pMdl );

        IoFreeMdl( pPageEntry->pMdl );
    }

    RemoveEntryList( &pPageEntry->link );
    ExFreePoolWithTag( pPageEntry, BB_POOL_TAG );

    return status;
}

int hax_pin_user_pages(uint64 start_uva, uint64 size, hax_memdesc_user *memdesc)
{
    PMDL pmdl = NULL;

    if (!memdesc) {
        hax_error("%s: memdesc == NULL\n", __func__);
        return -EINVAL;
    }

    // TODO: Check whether [start_uva, start_uva + size) is a valid UVA range

    pmdl = IoAllocateMdl((PVOID)start_uva, size, FALSE, FALSE, NULL);
    if (!pmdl) {
        hax_error("Failed to allocate MDL for va: 0x%llx, size: 0x%llx.\n",
                  start_uva, size);
        return -EFAULT;
    }

    try {
        MmProbeAndLockPages(pmdl, UserMode, IoWriteAccess);
    } except (EXCEPTION_EXECUTE_HANDLER) {
        hax_error("Failed to probe pages for guest's memory! va: 0x%llx\n",
                  start_uva);
        IoFreeMdl(pmdl);
        return -ENOMEM;
    }

    memdesc->pmdl = pmdl;
    return 0;
}

static NTSTATUS
redirect_irp_completion(PDEVICE_OBJECT dev_obj, PIRP irp, PVOID ctxt)
{
    IO_STACK_LOCATION *const isl = IoGetCurrentIrpStackLocation(irp);
    struct scsifilt *const sf = get_scsifilt(dev_obj);
    SCSI_REQUEST_BLOCK *const new_srb = isl->Parameters.Scsi.Srb;
    SCSI_REQUEST_BLOCK *const srb = new_srb->OriginalRequest;
    MDL *const mdl = irp->MdlAddress;
    void *const buf = new_srb->DataBuffer;

    UNREFERENCED_PARAMETER(ctxt);

    srb->SrbStatus = new_srb->SrbStatus;
    srb->ScsiStatus = new_srb->ScsiStatus;
    memcpy(srb->DataBuffer, buf, srb->DataTransferLength);

    IoFreeIrp(irp);
    IoFreeMdl(mdl);
    ExFreePool(buf);
    XmFreeMemory(new_srb);

    queue_srb_for_completion(sf, srb);

    return STATUS_MORE_PROCESSING_REQUIRED;
}

VOID
LlcTerminate(
    VOID
    )

/*++

Routine Description:

    The routines terminates the LLC protocol module and frees its global
    resources. This assumes all adapter bindings to be closed.

Arguments:

    None.

Return Value:

    None.

--*/

{
    NDIS_STATUS Status;

    ASSUME_IRQL(PASSIVE_LEVEL);

    LlcTerminateTimerSystem();
    NdisDeregisterProtocol(&Status, LlcProtocolHandle);
    IoFreeMdl(pXidMdl);
}

VOID
NTAPI
FatDestroyIrpContext(PFAT_IRP_CONTEXT IrpContext)
{
    PAGED_CODE();

    /* Make sure it has no pinned stuff */
    ASSERT(IrpContext->PinCount == 0);

    /* If there is a FatIo context associated with it - free it */
    if (IrpContext->FatIoContext)
    {
        if (!(IrpContext->Flags & IRPCONTEXT_STACK_IO_CONTEXT))
        {
            /* If a zero mdl was allocated - free it */
            if (IrpContext->FatIoContext->ZeroMdl)
                IoFreeMdl(IrpContext->FatIoContext->ZeroMdl);

            /* Free memory of FatIo context */
            ExFreePool(IrpContext->FatIoContext);
        }
    }

    /* Free memory */
    ExFreeToNPagedLookasideList(&FatGlobalData.IrpContextList, IrpContext);
}

void 
UserMemoryManager::CleanupCurrentProcess()
{
	PEPROCESS	curProc	= PsGetCurrentProcess();

	mLMIList.PublicLock();
	LockedMemInfo	*pLMI = mLMIList.Head();

	// Unlock memory for this process.
	while (NULL!=(int)(pLMI))
	{
		LockedMemInfo	*pNext	= mLMIList.Next(pLMI);

		if (pLMI->proc == curProc)
		{
			MmUnlockPages(pLMI->pMdl);
			IoFreeMdl(pLMI->pMdl);

			mLMIList.Remove(pLMI);
			mLMIAllocator.Free(pLMI);
		}

		pLMI = pNext;
	}

	mLMIList.PublicUnlock();
}

void 
NTAPI 
StreamOobInjectCompletionFn(
   _Inout_ void* context,
   _Inout_ NET_BUFFER_LIST* netBufferList,
   BOOLEAN dispatchLevel
   )
/* ++

   Injection completion function for injecting an NBL created using 
   FwpsAllocateNetBufferAndNetBufferList. This function frees up 
   resources allocated during StreamOobReinjectData().

-- */
{
   MDL* mdl = (MDL*)context;

   UNREFERENCED_PARAMETER(dispatchLevel);

   FwpsFreeNetBufferList(netBufferList);

   if (mdl != NULL)
   {
      IoFreeMdl(mdl);

      //
      // The MDL mapped over a pool alloc which we need to free here.
      //

      ExFreePoolWithTag(
         mdl->MappedSystemVa, 
         STREAM_EDITOR_MDL_DATA_TAG
         );
   }
}

NTSTATUS
Ext2ReadWriteBlockSyncCompletionRoutine (
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp,
    IN PVOID Context    )
{
    PEXT2_RW_CONTEXT pContext = (PEXT2_RW_CONTEXT)Context;

    if (Irp != pContext->MasterIrp) {

        if (!NT_SUCCESS(Irp->IoStatus.Status)) {
            pContext->MasterIrp->IoStatus = Irp->IoStatus;
        }

        IoFreeMdl(Irp->MdlAddress);
        IoFreeIrp(Irp );
    }

    if (InterlockedDecrement(&pContext->Blocks) == 0) {

        pContext->MasterIrp->IoStatus.Information = 0;
        if (NT_SUCCESS(pContext->MasterIrp->IoStatus.Status)) {

            pContext->MasterIrp->IoStatus.Information =
                pContext->Length;
        }

        KeSetEvent(&pContext->Event, 0, FALSE);
    }

    return STATUS_MORE_PROCESSING_REQUIRED;
}

NTSTATUS
NTAPI
MiSimpleReadComplete(PDEVICE_OBJECT DeviceObject,
                     PIRP Irp,
                     PVOID Context)
{
    PMDL Mdl = Irp->MdlAddress;

    /* Unlock MDL Pages, page 167. */
    DPRINT("MiSimpleReadComplete %p\n", Irp);
    while (Mdl)
    {
        DPRINT("MDL Unlock %p\n", Mdl);
        MmUnlockPages(Mdl);
        Mdl = Mdl->Next;
    }

    /* Check if there's an MDL */
    while ((Mdl = Irp->MdlAddress))
    {
        /* Clear all of them */
        Irp->MdlAddress = Mdl->Next;
        IoFreeMdl(Mdl);
    }

    return STATUS_SUCCESS;
}

VOID
NTAPI
CcMdlReadComplete(IN PFILE_OBJECT FileObject,
                  IN PMDL MdlChain)
{
	IoFreeMdl(MdlChain);
}

VOID
NdisFreeBuffer(
    IN PNDIS_BUFFER Buffer
    )
{
    IoFreeMdl(Buffer);
}

/*
* TsmiHandleMemWrite
*
* Purpose:
*
* Patch vbox dll in memory.
*
* Warning: potential BSOD-generator due to nonstandard way of loading, take care with patch offsets.
*
*/
NTSTATUS TsmiHandleMemWrite(
    _In_ PVOID SrcAddress,
    _In_ PVOID DestAddress,
    _In_ ULONG Size
)
{
    PMDL        mdl;
    NTSTATUS    status = STATUS_SUCCESS;

    PAGED_CODE();

    mdl = IoAllocateMdl(DestAddress, Size, FALSE, FALSE, NULL);
    if (mdl == NULL) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    if (DestAddress >= MmSystemRangeStart)
        if (!MmIsAddressValid(DestAddress)) {
            return STATUS_ACCESS_VIOLATION;
        }
    MmProbeAndLockPages(mdl, KernelMode, IoReadAccess);
    DestAddress = MmGetSystemAddressForMdlSafe(mdl, HighPagePriority);
    if (DestAddress != NULL) {
        status = MmProtectMdlSystemAddress(mdl, PAGE_EXECUTE_READWRITE);
        __movsb((PUCHAR)DestAddress, (const UCHAR *)SrcAddress, Size);
        MmUnmapLockedPages(DestAddress, mdl);
        MmUnlockPages(mdl);
    }
    else {
        status = STATUS_ACCESS_VIOLATION;
    }

    IoFreeMdl(mdl);
    return status;
}

void
StreamEditEvtDriverUnload(
   _In_ WDFDRIVER driverObject
   )
{

   UNREFERENCED_PARAMETER(driverObject);

   if (!configEditInline)
   {
      OobEditShutdown(&gStreamEditor);
   }

   if (gStreamEditor.scratchBuffer != NULL)
   {
      ExFreePoolWithTag(
         gStreamEditor.scratchBuffer,
         STREAM_EDITOR_FLAT_BUFFER_TAG
         );

   }

   StreamEditUnregisterCallout();

   FwpsInjectionHandleDestroy(gInjectionHandle);

   NdisFreeNetBufferListPool(gNetBufferListPool);
   NdisFreeGenericObject(gNdisGenericObj);

   IoFreeMdl(gStringToReplaceMdl);
}

NTSTATUS
Ext2LockUserBuffer (IN PIRP     Irp,
                    IN ULONG            Length,
                    IN LOCK_OPERATION   Operation)
{
    NTSTATUS Status;
    ASSERT(Irp != NULL);

    if (Irp->MdlAddress != NULL) {
        return STATUS_SUCCESS;
    }

    IoAllocateMdl(Irp->UserBuffer, Length, FALSE, FALSE, Irp);
    if (Irp->MdlAddress == NULL) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    __try {

        MmProbeAndLockPages(Irp->MdlAddress, Irp->RequestorMode, Operation);
        Status = STATUS_SUCCESS;

    } __except (EXCEPTION_EXECUTE_HANDLER) {

        DbgBreak();
        IoFreeMdl(Irp->MdlAddress);
        Irp->MdlAddress = NULL;
        Status = STATUS_INVALID_USER_BUFFER;
    }

    return Status;
}

PMDL
Ext2CreateMdl (
    IN PVOID Buffer,
    IN BOOLEAN bPaged,
    IN ULONG Length,
    IN LOCK_OPERATION Operation
)
{
    NTSTATUS Status;
    PMDL Mdl = NULL;

    ASSERT (Buffer != NULL);
    Mdl = IoAllocateMdl (Buffer, Length, FALSE, FALSE, NULL);
    if (Mdl == NULL) {
        Status = STATUS_INSUFFICIENT_RESOURCES;
    } else {
        __try {
            if (bPaged) {
                MmProbeAndLockPages(Mdl, KernelMode, Operation);
            } else {
                MmBuildMdlForNonPagedPool (Mdl);
            }
            Status = STATUS_SUCCESS;
        } __except (EXCEPTION_EXECUTE_HANDLER) {
            IoFreeMdl (Mdl);
            Mdl = NULL;
            DbgBreak();
            Status = STATUS_INVALID_USER_BUFFER;
        }
    }
    return Mdl;
}