/*
 Copyright (c) 2008 TrueCrypt Foundation. All rights reserved.

 Governed by the TrueCrypt License 2.5 the full text of which is contained
 in the file License.txt included in TrueCrypt binary and source code
 distribution packages.
*/

#include "TCdefs.h"
#include "Apidrvr.h"
#include "Ntdriver.h"
#include "EncryptedIoQueue.h"
#include "EncryptionThreadPool.h"
#include "Volumes.h"


static void DecrementOutstandingIoCount (EncryptedIoQueue *queue)
{
	if (InterlockedDecrement (&queue->OutstandingIoCount) == 0 && (queue->SuspendPending || queue->StopPending))
		KeSetEvent (&queue->NoOutstandingIoEvent, IO_DISK_INCREMENT, FALSE);
}


static void OnItemCompleted (EncryptedIoQueueItem *item, BOOL freeItem)
{
	DecrementOutstandingIoCount (item->Queue);
	
	if (item->Queue->IsFilterDevice)
		IoReleaseRemoveLock (&item->Queue->RemoveLock, item->OriginalIrp);

	if (NT_SUCCESS (item->Status))
	{
		if (item->Write)
			item->Queue->TotalBytesWritten += item->OriginalLength;
		else
			item->Queue->TotalBytesRead += item->OriginalLength;
	}

	if (freeItem)
		TCfree (item);
}


static NTSTATUS CompleteOriginalIrp (EncryptedIoQueueItem *item, NTSTATUS status, ULONG_PTR information)
{
	//Dump ("Queue comp  offset=%I64d  status=%x  info=%p  out=%d\n", item->OriginalOffset, status, information, item->Queue->OutstandingIoCount - 1);
	
	TCCompleteDiskIrp (item->OriginalIrp, status, information);

	item->Status = status;
	OnItemCompleted (item, TRUE);

	return status;
}


static void AcquireFragmentBuffer (EncryptedIoQueue *queue, byte *buffer)
{
	NTSTATUS status = STATUS_INVALID_PARAMETER;

	if (buffer == queue->FragmentBufferA)
	{
		status = KeWaitForSingleObject (&queue->FragmentBufferAFreeEvent, Executive, KernelMode, FALSE, NULL);
	}
	else if (buffer == queue->FragmentBufferB)
	{
		status = KeWaitForSingleObject (&queue->FragmentBufferBFreeEvent, Executive, KernelMode, FALSE, NULL);
	}

	if (!NT_SUCCESS (status))
		TC_BUG_CHECK (status);
}


static void ReleaseFragmentBuffer (EncryptedIoQueue *queue, byte *buffer)
{
	if (buffer == queue->FragmentBufferA)
	{
		KeSetEvent (&queue->FragmentBufferAFreeEvent, IO_DISK_INCREMENT, FALSE);
	}
	else if (buffer == queue->FragmentBufferB)
	{
		KeSetEvent (&queue->FragmentBufferBFreeEvent, IO_DISK_INCREMENT, FALSE);
	}
	else
	{
		TC_BUG_CHECK (STATUS_INVALID_PARAMETER);
	}
}


static VOID CompletionThreadProc (PVOID threadArg)
{
	EncryptedIoQueue *queue = (EncryptedIoQueue *) threadArg;
	PLIST_ENTRY listEntry;
	EncryptedIoRequest *request;
	UINT64_STRUCT dataUnit;

	if (IsEncryptionThreadPoolRunning())
		KeSetPriorityThread (KeGetCurrentThread(), LOW_REALTIME_PRIORITY);

	while (!queue->ThreadExitRequested)
	{
		if (!NT_SUCCESS (KeWaitForSingleObject (&queue->CompletionThreadQueueNotEmptyEvent, Executive, KernelMode, FALSE, NULL)))
			continue;

		if (queue->ThreadExitRequested)
			break;

		while ((listEntry = ExInterlockedRemoveHeadList (&queue->CompletionThreadQueue, &queue->CompletionThreadQueueLock)))
		{
			request = CONTAINING_RECORD (listEntry, EncryptedIoRequest, CompletionListEntry);

			if (request->EncryptedLength > 0 && NT_SUCCESS (request->Item->Status))
			{
				ASSERT (request->EncryptedOffset + request->EncryptedLength <= request->Offset.QuadPart + request->Length);
				dataUnit.Value = (request->Offset.QuadPart + request->EncryptedOffset) / ENCRYPTION_DATA_UNIT_SIZE;

				if (queue->CryptoInfo->bPartitionInInactiveSysEncScope)
					dataUnit.Value += queue->CryptoInfo->FirstDataUnitNo.Value;
				else if (queue->RemapEncryptedArea)
					dataUnit.Value += queue->RemappedAreaDataUnitOffset;

				DecryptDataUnits (request->Data + request->EncryptedOffset, &dataUnit, request->EncryptedLength / ENCRYPTION_DATA_UNIT_SIZE, queue->CryptoInfo);
			}

			if (request->CompleteOriginalIrp)
			{
				CompleteOriginalIrp (request->Item, request->Item->Status,
					NT_SUCCESS (request->Item->Status) ? request->Item->OriginalLength : 0);
			}
			else
			{
				InterlockedDecrement (&request->Item->OutstandingRequestCount);
				KeSetEvent (&queue->RequestCompletedEvent, IO_DISK_INCREMENT, FALSE);
			}

			TCfree (request);
		}
	}

	PsTerminateSystemThread (STATUS_SUCCESS);
}


static VOID IoThreadProc (PVOID threadArg)
{
	EncryptedIoQueue *queue = (EncryptedIoQueue *) threadArg;
	PLIST_ENTRY listEntry;
	EncryptedIoRequest *request;

	KeSetPriorityThread (KeGetCurrentThread(), LOW_REALTIME_PRIORITY);

	while (!queue->ThreadExitRequested)
	{
		if (!NT_SUCCESS (KeWaitForSingleObject (&queue->IoThreadQueueNotEmptyEvent, Executive, KernelMode, FALSE, NULL)))
			continue;

		if (queue->ThreadExitRequested)
			break;

		while ((listEntry = ExInterlockedRemoveHeadList (&queue->IoThreadQueue, &queue->IoThreadQueueLock)))
		{
			request = CONTAINING_RECORD (listEntry, EncryptedIoRequest, ListEntry);
			
			// Perform IO request if no preceding request of the item failed
			if (NT_SUCCESS (request->Item->Status))
			{
				if (queue->IsFilterDevice)
				{
					if (queue->RemapEncryptedArea && request->EncryptedLength > 0)
					{
						if (request->EncryptedLength != request->Length)
						{
							// Up to three subfragments may be required to handle a partially remapped fragment
							int subFragment;
							byte *subFragmentData = request->Data;

							for (subFragment = 0 ; subFragment < 3; ++subFragment)
							{
								LARGE_INTEGER subFragmentOffset;
								ULONG subFragmentLength;
								subFragmentOffset.QuadPart = request->Offset.QuadPart;

								switch (subFragment)
								{
								case 0:
									subFragmentLength = (ULONG) request->EncryptedOffset;
									break;

								case 1:
									subFragmentOffset.QuadPart += request->EncryptedOffset + queue->RemappedAreaOffset;
									subFragmentLength = request->EncryptedLength;
									break;

								case 2:
									subFragmentOffset.QuadPart += request->EncryptedOffset + request->EncryptedLength;
									subFragmentLength = (ULONG) (request->Length - (request->EncryptedOffset + request->EncryptedLength));
									break;
								}

								if (subFragmentLength > 0)
								{
									if (request->Item->Write)
										request->Item->Status = TCWriteDevice (queue->LowerDeviceObject, subFragmentData, subFragmentOffset, subFragmentLength);
									else
										request->Item->Status = TCReadDevice (queue->LowerDeviceObject, subFragmentData, subFragmentOffset, subFragmentLength);

									subFragmentData += subFragmentLength;
								}
							}
						}
						else
						{
							// Remap the fragment
							LARGE_INTEGER remappedOffset;
							remappedOffset.QuadPart = request->Offset.QuadPart + queue->RemappedAreaOffset;

							if (request->Item->Write)
								request->Item->Status = TCWriteDevice (queue->LowerDeviceObject, request->Data, remappedOffset, request->Length);
							else
								request->Item->Status = TCReadDevice (queue->LowerDeviceObject, request->Data, remappedOffset, request->Length);
						}
					}
					else
					{
						if (request->Item->Write)
							request->Item->Status = TCWriteDevice (queue->LowerDeviceObject, request->Data, request->Offset, request->Length);
						else
							request->Item->Status = TCReadDevice (queue->LowerDeviceObject, request->Data, request->Offset, request->Length);
					}
				}
				else
				{
					IO_STATUS_BLOCK ioStatus;

					if (request->Item->Write)
						request->Item->Status = ZwWriteFile (queue->HostFileHandle, NULL, NULL, NULL, &ioStatus, request->Data, request->Length, &request->Offset, NULL);
					else
						request->Item->Status = ZwReadFile (queue->HostFileHandle, NULL, NULL, NULL, &ioStatus, request->Data, request->Length, &request->Offset, NULL);

					if (NT_SUCCESS (request->Item->Status) && ioStatus.Information != request->Length)
						request->Item->Status = STATUS_END_OF_FILE;
				}
			}

			if (request->Item->Write)
			{
				ReleaseFragmentBuffer (queue, request->Data);

				if (request->CompleteOriginalIrp)
				{
					CompleteOriginalIrp (request->Item, request->Item->Status,
						NT_SUCCESS (request->Item->Status) ? request->Item->OriginalLength : 0);
				}
				else
				{
					InterlockedDecrement (&request->Item->OutstandingRequestCount);
					KeSetEvent (&queue->RequestCompletedEvent, IO_DISK_INCREMENT, FALSE);
				}

				TCfree (request);
			}
			else
			{
				if (NT_SUCCESS (request->Item->Status))
					memcpy (request->OrigDataBufferFragment, request->Data, request->Length);

				ReleaseFragmentBuffer (queue, request->Data);
				request->Data = request->OrigDataBufferFragment;

				ExInterlockedInsertTailList (&queue->CompletionThreadQueue, &request->CompletionListEntry, &queue->CompletionThreadQueueLock);
				KeSetEvent (&queue->CompletionThreadQueueNotEmptyEvent, IO_DISK_INCREMENT, FALSE);
			}
		}
	}

	PsTerminateSystemThread (STATUS_SUCCESS);
}


static NTSTATUS OnPassedIrpCompleted (PDEVICE_OBJECT filterDeviceObject, PIRP irp, EncryptedIoQueueItem *item)
{
	if (irp->PendingReturned)
		IoMarkIrpPending (irp);

	OnItemCompleted (item, TRUE);
	return STATUS_CONTINUE_COMPLETION;
}


static VOID MainThreadProc (PVOID threadArg)
{
	EncryptedIoQueue *queue = (EncryptedIoQueue *) threadArg;
	PLIST_ENTRY listEntry;
	EncryptedIoQueueItem *item;

	LARGE_INTEGER fragmentOffset;
	ULONG dataRemaining;
	PUCHAR activeFragmentBuffer = queue->FragmentBufferA;
	PUCHAR dataBuffer;
	EncryptedIoRequest *request;
	uint64 intersectStart;
	uint32 intersectLength;

	int mdlWaitTime;
	LARGE_INTEGER mdlWaitInterval;
	mdlWaitInterval.QuadPart = TC_ENC_IO_QUEUE_MEM_ALLOC_RETRY_DELAY * -10000;

	if (IsEncryptionThreadPoolRunning())
		KeSetPriorityThread (KeGetCurrentThread(), LOW_REALTIME_PRIORITY);

	while (!queue->ThreadExitRequested)
	{
		if (!NT_SUCCESS (KeWaitForSingleObject (&queue->MainThreadQueueNotEmptyEvent, Executive, KernelMode, FALSE, NULL)))
			continue;

		while ((listEntry = ExInterlockedRemoveHeadList (&queue->MainThreadQueue, &queue->MainThreadQueueLock)))
		{
			PIRP irp = CONTAINING_RECORD (listEntry, IRP, Tail.Overlay.ListEntry);
			PIO_STACK_LOCATION irpSp = IoGetCurrentIrpStackLocation (irp);
			
			if (queue->Suspended)
				KeWaitForSingleObject (&queue->QueueResumedEvent, Executive, KernelMode, FALSE, NULL);

			item = AllocateMemoryWithTimeout (sizeof (EncryptedIoQueueItem), TC_ENC_IO_QUEUE_MEM_ALLOC_RETRY_DELAY, TC_ENC_IO_QUEUE_MEM_ALLOC_TIMEOUT);
			if (!item)
			{
				EncryptedIoQueueItem stackItem;
				stackItem.Queue = queue;
				stackItem.OriginalIrp = irp;
				stackItem.Status = STATUS_INSUFFICIENT_RESOURCES;

				TCCompleteDiskIrp (irp, STATUS_INSUFFICIENT_RESOURCES, 0);
				OnItemCompleted (&stackItem, FALSE);
				continue;
			}

			item->Queue = queue;
			item->OriginalIrp = irp;
			item->OutstandingRequestCount = 0;
			item->Status = STATUS_SUCCESS;

			IoSetCancelRoutine (irp, NULL);
			if (irp->Cancel)
			{
				CompleteOriginalIrp (item, STATUS_CANCELLED, 0);
				continue;
			}

			switch (irpSp->MajorFunction)
			{
			case IRP_MJ_READ:
				item->Write = FALSE;
				item->OriginalOffset = irpSp->Parameters.Read.ByteOffset;
				item->OriginalLength = irpSp->Parameters.Read.Length;
				break;

			case IRP_MJ_WRITE:
				item->Write = TRUE;
				item->OriginalOffset = irpSp->Parameters.Write.ByteOffset;
				item->OriginalLength = irpSp->Parameters.Write.Length;
				break;

			default:
				CompleteOriginalIrp (item, STATUS_INVALID_PARAMETER, 0);