encryptedioqueue.c

外国开源的硬盘加密软件

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

 Governed by the TrueCrypt License 2.6 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 "DriveFilter.h"
#include "EncryptedIoQueue.h"
#include "EncryptionThreadPool.h"
#include "Volumes.h"


static void AcquireBufferPoolMutex (EncryptedIoQueue *queue)
{
	NTSTATUS status;

	status = KeWaitForMutexObject (&queue->BufferPoolMutex, Executive, KernelMode, FALSE, NULL);
	if (!NT_SUCCESS (status))
		TC_BUG_CHECK (status);
}


static void ReleaseBufferPoolMutex (EncryptedIoQueue *queue)
{
	KeReleaseMutex (&queue->BufferPoolMutex, FALSE);
}


static void *GetPoolBuffer (EncryptedIoQueue *queue, ULONG requestedSize)
{
	EncryptedIoQueueBuffer *buffer;
	void *bufferAddress = NULL;
	BOOL requestedSizePresentInPool = FALSE;

	while (TRUE)
	{
		AcquireBufferPoolMutex (queue);

		for (buffer = queue->FirstPoolBuffer; ; buffer = buffer->NextBuffer)
		{
			if (buffer && buffer->Size == requestedSize)
			{
				requestedSizePresentInPool = TRUE;

				if (!buffer->InUse)
				{
					// Reuse a free buffer
					buffer->InUse = TRUE;
					bufferAddress = buffer->Address;
					break;
				}
			}

			if (!buffer || !buffer->NextBuffer)
			{
				// Allocate a new buffer
				EncryptedIoQueueBuffer *newBuffer = TCalloc (sizeof (EncryptedIoQueueBuffer));
				if (!newBuffer)
				{
					bufferAddress = NULL;
					break;
				}

				bufferAddress = TCalloc (requestedSize);
				if (bufferAddress)
				{
					newBuffer->NextBuffer = NULL;
					newBuffer->Address = bufferAddress;
					newBuffer->Size = requestedSize;
					newBuffer->InUse = TRUE;

					if (!buffer)
						queue->FirstPoolBuffer = newBuffer;
					else
						buffer->NextBuffer = newBuffer;
				}
				else
					TCfree (newBuffer);

				break;
			}
		}

		ReleaseBufferPoolMutex (queue);

		if (bufferAddress || !requestedSizePresentInPool || queue->StartPending)
			break;

		TCSleep (TC_ENC_IO_QUEUE_MEM_ALLOC_RETRY_DELAY);
	}

	return bufferAddress;
}


static void ReleasePoolBuffer (EncryptedIoQueue *queue, void *address)
{
	EncryptedIoQueueBuffer *buffer;
	AcquireBufferPoolMutex (queue);
	
	for (buffer = queue->FirstPoolBuffer; buffer != NULL; buffer = buffer->NextBuffer)
	{
		if (buffer->Address == address)
		{
			ASSERT (buffer->InUse);

			buffer->InUse = FALSE;
			break;
		}
	}

	ReleaseBufferPoolMutex (queue);
}


static void FreePoolBuffers (EncryptedIoQueue *queue)
{
	EncryptedIoQueueBuffer *buffer;
	AcquireBufferPoolMutex (queue);

	for (buffer = queue->FirstPoolBuffer; buffer != NULL; )
	{
		EncryptedIoQueueBuffer *nextBuffer = buffer->NextBuffer;

		ASSERT (!buffer->InUse);

		TCfree (buffer->Address);
		TCfree (buffer);

		buffer = nextBuffer;
	}

	queue->FirstPoolBuffer = NULL;
	ReleaseBufferPoolMutex (queue);
}


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)
		ReleasePoolBuffer (item->Queue, item);
}


static NTSTATUS CompleteOriginalIrp (EncryptedIoQueueItem *item, NTSTATUS status, ULONG_PTR information)
{
#ifdef TC_TRACE_IO_QUEUE
	Dump ("< %I64d [%I64d] %c status=%x info=%I64d\n", item->OriginalIrpOffset, GetElapsedTime (&item->Queue->LastPerformanceCounter), item->Write ? 'W' : 'R', status, (int64) information);
#endif

	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);
			}

			ReleasePoolBuffer (queue, request);
		}
	}

	PsTerminateSystemThread (STATUS_SUCCESS);
}


static NTSTATUS TCCachedRead (EncryptedIoQueue *queue, IO_STATUS_BLOCK *ioStatus, PVOID buffer, LARGE_INTEGER offset, ULONG length)
{
	queue->LastReadOffset = offset;
	queue->LastReadLength = length;

	if (queue->ReadAheadBufferValid && queue->ReadAheadOffset.QuadPart == offset.QuadPart && queue->ReadAheadLength >= length)
	{
		memcpy (buffer, queue->ReadAheadBuffer, length);

		if (!queue->IsFilterDevice)
		{
			ioStatus->Information = length;
			ioStatus->Status = STATUS_SUCCESS;
		}

		return STATUS_SUCCESS;
	}

	if (queue->IsFilterDevice)
		return TCReadDevice (queue->LowerDeviceObject, buffer, offset, length);

	return ZwReadFile (queue->HostFileHandle, NULL, NULL, NULL, ioStatus, buffer, length, &offset, NULL);
}


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

	KeSetPriorityThread (KeGetCurrentThread(), LOW_REALTIME_PRIORITY);

	if (!queue->IsFilterDevice && queue->SecurityClientContext)
	{
#ifdef DEBUG
		NTSTATUS status =
#endif
		SeImpersonateClientEx (queue->SecurityClientContext, NULL);
		ASSERT (NT_SUCCESS (status));
	}

	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)))
		{
			InterlockedDecrement (&queue->IoThreadPendingRequestCount);
			request = CONTAINING_RECORD (listEntry, EncryptedIoRequest, ListEntry);
			
#ifdef TC_TRACE_IO_QUEUE
			Dump ("%c   %I64d [%I64d] roff=%I64d rlen=%d\n", request->Item->Write ? 'W' : 'R', request->Item->OriginalIrpOffset.QuadPart, GetElapsedTime (&queue->LastPerformanceCounter), request->Offset.QuadPart, request->Length);
#endif

			// 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;