/*
 Derived from source code of TrueCrypt 7.1a, which is
 Copyright (c) 2008-2012 TrueCrypt Developers Association and which is governed
 by the TrueCrypt License 3.0.

 Modifications and additions to the original source code (contained in this file)
 and all other portions of this file are Copyright (c) 2013-2017 IDRIX
 and are governed by the Apache License 2.0 the full text of which is
 contained in the file License.txt included in VeraCrypt binary and source
 code distribution packages.
*/

#include "TCdefs.h"
#include "Ntdriver.h"
#include "Ntvol.h"
#include "DriveFilter.h"
#include "VolumeFilter.h"

typedef DriveFilterExtension VolumeFilterExtension;

// Number of times the filter driver answered that an unencrypted volume
// is read-only (or mounted an outer/normal TrueCrypt volume as read only)
uint32 HiddenSysLeakProtectionCount = 0;


NTSTATUS VolumeFilterAddDevice (PDRIVER_OBJECT driverObject, PDEVICE_OBJECT pdo)
{
	VolumeFilterExtension *Extension = NULL;
	NTSTATUS status;
	PDEVICE_OBJECT filterDeviceObject = NULL;
	PDEVICE_OBJECT attachedDeviceObject;

	Dump ("VolumeFilterAddDevice pdo=%p\n", pdo);

	attachedDeviceObject = IoGetAttachedDeviceReference (pdo);
	status = IoCreateDevice (driverObject, sizeof (VolumeFilterExtension), NULL, attachedDeviceObject->DeviceType, 0, FALSE, &filterDeviceObject);

	ObDereferenceObject (attachedDeviceObject);

	if (!NT_SUCCESS (status))
	{
		filterDeviceObject = NULL;
		goto err;
	}

	Extension = (VolumeFilterExtension *) filterDeviceObject->DeviceExtension;
	memset (Extension, 0, sizeof (VolumeFilterExtension));

	status = IoAttachDeviceToDeviceStackSafe (filterDeviceObject, pdo, &(Extension->LowerDeviceObject));
	if (status != STATUS_SUCCESS)
	{
		goto err;
	}

	if (!Extension->LowerDeviceObject)
	{
		status = STATUS_DEVICE_REMOVED;
		goto err;
	}

	Extension->IsVolumeFilterDevice = TRUE;
	Extension->DeviceObject = filterDeviceObject;
	Extension->Pdo = pdo;

	IoInitializeRemoveLock (&Extension->Queue.RemoveLock, 'LRCV', 0, 0);

	filterDeviceObject->Flags |= Extension->LowerDeviceObject->Flags & (DO_DIRECT_IO | DO_BUFFERED_IO | DO_POWER_PAGABLE);
	filterDeviceObject->Flags &= ~DO_DEVICE_INITIALIZING;

	return status;

err:
	if (filterDeviceObject)
	{
		if (Extension && Extension->LowerDeviceObject)
			IoDetachDevice (Extension->LowerDeviceObject);

		IoDeleteDevice (filterDeviceObject);
	}

	return status;
}


static NTSTATUS PassIrp (PDEVICE_OBJECT deviceObject, PIRP irp)
{
	IoSkipCurrentIrpStackLocation (irp);
	return IoCallDriver (deviceObject, irp);
}


static NTSTATUS PassFilteredIrp (PDEVICE_OBJECT deviceObject, PIRP irp, PIO_COMPLETION_ROUTINE completionRoutine, PVOID completionRoutineArg)
{
	IoCopyCurrentIrpStackLocationToNext (irp);

	if (completionRoutine)
		IoSetCompletionRoutine (irp, completionRoutine, completionRoutineArg, TRUE, TRUE, TRUE);

	return IoCallDriver (deviceObject, irp);
}


static NTSTATUS OnDeviceUsageNotificationCompleted (PDEVICE_OBJECT filterDeviceObject, PIRP Irp, VolumeFilterExtension *Extension)
{
	if (Irp->PendingReturned)
		IoMarkIrpPending (Irp);

	if (!(Extension->LowerDeviceObject->Flags & DO_POWER_PAGABLE))
		filterDeviceObject->Flags &= ~DO_POWER_PAGABLE;

	IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
	return STATUS_CONTINUE_COMPLETION;
}


static NTSTATUS OnStartDeviceCompleted (PDEVICE_OBJECT filterDeviceObject, PIRP Irp, VolumeFilterExtension *Extension)
{
	if (Irp->PendingReturned)
		IoMarkIrpPending (Irp);

	if (Extension->LowerDeviceObject->Characteristics & FILE_REMOVABLE_MEDIA)
		filterDeviceObject->Characteristics |= FILE_REMOVABLE_MEDIA;

	IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
	return STATUS_CONTINUE_COMPLETION;
}

static BOOL IsSystemVolumePartition (VolumeFilterExtension *Extension)
{
	NTSTATUS status;
	BOOL bRet = FALSE;
	DriveFilterExtension *bootDriveExtension = GetBootDriveFilterExtension();
	STORAGE_DEVICE_NUMBER storageDeviceNumber;

	if (!bootDriveExtension->SystemStorageDeviceNumberValid)
		TC_BUG_CHECK (STATUS_INVALID_PARAMETER);

	status = SendDeviceIoControlRequest (Extension->LowerDeviceObject, IOCTL_STORAGE_GET_DEVICE_NUMBER, NULL, 0, &storageDeviceNumber, sizeof (storageDeviceNumber));

	if (NT_SUCCESS (status) && bootDriveExtension->SystemStorageDeviceNumber == storageDeviceNumber.DeviceNumber)
	{
		PARTITION_INFORMATION_EX partition;
		status = SendDeviceIoControlRequest (Extension->LowerDeviceObject, IOCTL_DISK_GET_PARTITION_INFO_EX, NULL, 0, &partition, sizeof (partition));

		if (NT_SUCCESS (status) && partition.StartingOffset.QuadPart == bootDriveExtension->ConfiguredEncryptedAreaStart)
		{
			bRet = TRUE;
		}
	}

	return bRet;
}


static NTSTATUS DispatchControl (PDEVICE_OBJECT DeviceObject, PIRP Irp, VolumeFilterExtension *Extension, PIO_STACK_LOCATION irpSp)
{
	NTSTATUS status = IoAcquireRemoveLock (&Extension->Queue.RemoveLock, Irp);
	UNREFERENCED_PARAMETER(DeviceObject);
	if (!NT_SUCCESS (status))
		return TCCompleteIrp (Irp, status, 0);

	if (IsHiddenSystemRunning())
	{
		switch (irpSp->Parameters.DeviceIoControl.IoControlCode)
		{
		case IOCTL_DISK_IS_WRITABLE:
			{
				// All volumes except the system volume must be read-only
				if (IsSystemVolumePartition(Extension))
				{
					IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
					return TCCompleteDiskIrp (Irp, STATUS_SUCCESS, 0);
				}

				IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);

				++HiddenSysLeakProtectionCount;
				return TCCompleteDiskIrp (Irp, STATUS_MEDIA_WRITE_PROTECTED, 0);
			}

		case TC_IOCTL_DISK_IS_WRITABLE:
			Dump ("TC_IOCTL_DISK_IS_WRITABLE pdo=%p\n", Extension->Pdo);

			if (!ProbingHostDeviceForWrite)
				break;

			// Probe the real state of the device as the user is mounting a TrueCrypt volume.

			// Volume filter may be attached to a merged drive+volume PDO. First test if TC_IOCTL_DISK_IS_WRITABLE works for the underlying device.
			status = SendDeviceIoControlRequest (Extension->LowerDeviceObject, TC_IOCTL_DISK_IS_WRITABLE, NULL, 0, NULL, 0);

			if (NT_SUCCESS (status) || status == STATUS_MEDIA_WRITE_PROTECTED)
			{
				IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
				return TCCompleteDiskIrp (Irp, status, 0);
			}

			status = SendDeviceIoControlRequest (Extension->LowerDeviceObject, IOCTL_DISK_IS_WRITABLE, NULL, 0, NULL, 0);

			IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
			return TCCompleteDiskIrp (Irp, status, 0);

		case IOCTL_STORAGE_MANAGE_DATA_SET_ATTRIBUTES:

			// Filter IOCTL_STORAGE_MANAGE_DATA_SET_ATTRIBUTES to enable potential future use of hidden systems on drives that use the trim operation but not wear-leveling (if any appear in future). The documentation forbids users to create hidden volumes/systems on drives that use wear-leveling and consequently also on drives that use trim (as trim is used only by drives that use wear-leveling, as of 2010).

			IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
			return TCCompleteDiskIrp (Irp, STATUS_SUCCESS, 0);

		case IOCTL_DISK_GROW_PARTITION:
			if (IsSystemVolumePartition(Extension))
			{
				Dump ("VolumeFilter-DispatchControl: IOCTL_DISK_GROW_PARTITION blocked\n");
				IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
				return TCCompleteDiskIrp (Irp, STATUS_UNSUCCESSFUL, 0);
			}
			break;
		}
	}

	status = PassIrp (Extension->LowerDeviceObject, Irp);
	IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
	return status;
}


static NTSTATUS DispatchPnp (PDEVICE_OBJECT DeviceObject, PIRP Irp, VolumeFilterExtension *Extension, PIO_STACK_LOCATION irpSp)
{
	NTSTATUS status;

	status = IoAcquireRemoveLock (&Extension->Queue.RemoveLock, Irp);
	if (!NT_SUCCESS (status))
		return TCCompleteIrp (Irp, status, 0);

	switch (irpSp->MinorFunction)
	{
	case IRP_MN_START_DEVICE:
		Dump ("IRP_MN_START_DEVICE volume pdo=%p\n", Extension->Pdo);
		return PassFilteredIrp (Extension->LowerDeviceObject, Irp, OnStartDeviceCompleted, Extension);

	case IRP_MN_DEVICE_USAGE_NOTIFICATION:
		{
			PDEVICE_OBJECT attachedDevice = IoGetAttachedDeviceReference (DeviceObject);

			if (attachedDevice == DeviceObject || (attachedDevice->Flags & DO_POWER_PAGABLE))
				DeviceObject->Flags |= DO_POWER_PAGABLE;

			ObDereferenceObject (attachedDevice);
		}

		return PassFilteredIrp (Extension->LowerDeviceObject, Irp, OnDeviceUsageNotificationCompleted, Extension);


	case IRP_MN_REMOVE_DEVICE:
		Dump ("IRP_MN_REMOVE_DEVICE volume pdo=%p\n", Extension->Pdo);

		IoReleaseRemoveLockAndWait (&Extension->Queue.RemoveLock, Irp);
		status = PassIrp (Extension->LowerDeviceObject, Irp);

		IoDetachDevice (Extension->LowerDeviceObject);

		IoDeleteDevice (DeviceObject);
		return status;

	default:
		status = PassIrp (Extension->LowerDeviceObject, Irp);
		IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
	}

	return status;
}


static NTSTATUS DispatchPower (PDEVICE_OBJECT DeviceObject, PIRP Irp, VolumeFilterExtension *Extension)
{
	NTSTATUS status;
	UNREFERENCED_PARAMETER(DeviceObject);
	PoStartNextPowerIrp (Irp);

	status = IoAcquireRemoveLock (&Extension->Queue.RemoveLock, Irp);
	if (!NT_SUCCESS (status))
		return TCCompleteIrp (Irp, status, 0);

	IoSkipCurrentIrpStackLocation (Irp);
	status = PoCallDriver (Extension->LowerDeviceObject, Irp);

	IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
	return status;
}


NTSTATUS VolumeFilterDispatchIrp (PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	VolumeFilterExtension *Extension = (VolumeFilterExtension *) DeviceObject->DeviceExtension;
	PIO_STACK_LOCATION irpSp = IoGetCurrentIrpStackLocation (Irp);
	NTSTATUS status;

	ASSERT (!Extension->bRootDevice && Extension->IsVolumeFilterDevice);

	switch (irpSp->MajorFunction)
	{
	case IRP_MJ_DEVICE_CONTROL:
		return DispatchControl (DeviceObject, Irp, Extension, irpSp);

	case IRP_MJ_PNP:
		return DispatchPnp (DeviceObject, Irp, Extension, irpSp);

	case IRP_MJ_POWER:
		return DispatchPower (DeviceObject, Irp, Extension);

	default:
		status = IoAcquireRemoveLock (&Extension->Queue.RemoveLock, Irp);
		if (!NT_SUCCESS (status))
			return TCCompleteIrp (Irp, status, 0);

		status = PassIrp (Extension->LowerDeviceObject, Irp);

		IoReleaseRemoveLock (&Extension->Queue.RemoveLock, Irp);
		return status;
	}
}