/*
	This is a diskfilter for Windows NT that uses a Linux swap partition to
	provide a temporary storage area formated to the FAT file system,
	release 1.
	Copyright (C) 1999	Bo Branten
	This program is free software; you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation; either version 2 of the License, or
	(at your option) any later version.
	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.
	You should have received a copy of the GNU General Public License
	along with this program; if not, write to the Free Software
	Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

	Please send comments, corrections and contributions to bosse@acc.umu.se

	The most recent version of this program is available from:
	http://www.acc.umu.se/~bosse/

	Revision history:

	1. Initial release.
*/

#include <ntddk.h>
#include <ntdddisk.h>
#include "swap.h"
#include "mkdosfs.h"

#define PARAMETER_KEY		L"\\Parameters"
#define SWAPDEVICE_VALUE	L"SwapDevice"
#define ROOT_DIR_ENTRYS		512

typedef struct _DEVICE_EXTENSION
{
	PDEVICE_OBJECT target_device;
} DEVICE_EXTENSION, *PDEVICE_EXTENSION;

NTSTATUS
DriverEntry (
	IN PDRIVER_OBJECT	DriverObject,
	IN PUNICODE_STRING	RegistryPath
);

NTSTATUS
Dispatch (
	IN PDEVICE_OBJECT	DeviceObject,
	IN PIRP 			Irp
);

NTSTATUS
IoCompletion (
	IN PDEVICE_OBJECT	DeviceObject,
	IN PIRP				Irp,
	IN PVOID			Context
);

BOOLEAN
IsDeviceMounted (
	IN PUNICODE_STRING	DeviceName
);

NTSTATUS
IsDeviceLinuxSwap (
	IN PDEVICE_OBJECT	DeviceObject
);

NTSTATUS
FormatDeviceToFat (
	IN PDEVICE_OBJECT	DeviceObject
);

NTSTATUS 
DeviceIoControl (
	IN PDEVICE_OBJECT	DeviceObject,
	IN ULONG			IoctlCode,
	IN PVOID			InputBuffer,
	IN ULONG			InputBufferSize,
	OUT PVOID			OutputBuffer,
	OUT PULONG			OutputBufferSize
);

NTSTATUS
ReadBlockDevice (
	IN PDEVICE_OBJECT	DeviceObject,
	IN PLARGE_INTEGER	Offset,
	IN ULONG			Length,
	OUT PVOID			Buffer
);

NTSTATUS
WriteBlockDevice (
	IN PDEVICE_OBJECT	DeviceObject,
	IN PLARGE_INTEGER	Offset,
	IN ULONG			Length,
	IN PVOID			Buffer
);

#pragma code_seg("init")

NTSTATUS
DriverEntry (
	IN PDRIVER_OBJECT	DriverObject,
	IN PUNICODE_STRING	RegistryPath
	)
{
	UNICODE_STRING				parameter_path;
	UNICODE_STRING				device_name;
	RTL_QUERY_REGISTRY_TABLE	query_table[2];
	NTSTATUS					status;
	PDEVICE_OBJECT				device_object;
	PDEVICE_EXTENSION			device_extension;

	KdPrint(("SwapFs: Loading\n"));

	DriverObject->MajorFunction[IRP_MJ_CREATE]			= Dispatch;
	DriverObject->MajorFunction[IRP_MJ_CLOSE]			= Dispatch;
	DriverObject->MajorFunction[IRP_MJ_FLUSH_BUFFERS]	= Dispatch;
	DriverObject->MajorFunction[IRP_MJ_SHUTDOWN]		= Dispatch;
	DriverObject->MajorFunction[IRP_MJ_READ]			= Dispatch;
	DriverObject->MajorFunction[IRP_MJ_WRITE]			= Dispatch;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]	= Dispatch;

	parameter_path.Length = 0;

	parameter_path.MaximumLength =
		RegistryPath->Length + sizeof(PARAMETER_KEY) + sizeof(WCHAR);

	parameter_path.Buffer =
		(PWSTR) ExAllocatePool(PagedPool, parameter_path.MaximumLength);

	if (!parameter_path.Buffer)
	{
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	RtlCopyUnicodeString(&parameter_path, RegistryPath);

	RtlAppendUnicodeToString(&parameter_path, PARAMETER_KEY);

	device_name.Length = 0;
	device_name.MaximumLength = 0;
	device_name.Buffer = NULL;

	RtlZeroMemory(&query_table[0], sizeof(query_table));

	query_table[0].Flags =
		RTL_QUERY_REGISTRY_DIRECT | RTL_QUERY_REGISTRY_REQUIRED;
    query_table[0].Name = SWAPDEVICE_VALUE;
    query_table[0].EntryContext = &device_name;

    status = RtlQueryRegistryValues(
		RTL_REGISTRY_ABSOLUTE,
		parameter_path.Buffer,
		&query_table[0],
		NULL,
		NULL
		);

	ExFreePool(parameter_path.Buffer);

	if (!NT_SUCCESS(status) || !device_name.Buffer)
	{
		KdPrint(("SwapFs: DriverEntry: Query registry failed\n"));
		return STATUS_DEVICE_CONFIGURATION_ERROR;
	}

	if (IsDeviceMounted(&device_name))
	{
		ExFreePool(device_name.Buffer);
		return STATUS_DEVICE_CONFIGURATION_ERROR;
	}

	status = IoCreateDevice(
		DriverObject,
		sizeof(DEVICE_EXTENSION),
		NULL,
		FILE_DEVICE_DISK,
		0,
		FALSE,
		&device_object
		);

	if (!NT_SUCCESS(status))
	{
		ExFreePool(device_name.Buffer);
		return status;
	}

	device_object->Flags |= DO_DIRECT_IO;

	device_extension = (PDEVICE_EXTENSION) device_object->DeviceExtension;

	status = IoAttachDevice(
		device_object,
		&device_name,
		&device_extension->target_device
		);

	ExFreePool(device_name.Buffer);

	if (!NT_SUCCESS(status))
	{
		IoDeleteDevice(device_object);
		return status;
	}

	device_object->AlignmentRequirement =
		device_extension->target_device->AlignmentRequirement;

	status = IsDeviceLinuxSwap(device_extension->target_device);

	if (!NT_SUCCESS(status))
	{
		IoDetachDevice(device_extension->target_device);
		IoDeleteDevice(device_object);
		return status;
	}

	status = FormatDeviceToFat(device_extension->target_device);

	if (!NT_SUCCESS(status))
	{
		IoDetachDevice(device_extension->target_device);
		IoDeleteDevice(device_object);
		return status;
	}

	return STATUS_SUCCESS;
}

#pragma code_seg()

NTSTATUS
Dispatch (
	IN PDEVICE_OBJECT	DeviceObject,
	IN PIRP 			Irp
	)
{
	PDEVICE_EXTENSION	device_extension;
	PIO_STACK_LOCATION	io_stack;
	PIO_STACK_LOCATION	next_io_stack;

	device_extension = (PDEVICE_EXTENSION) DeviceObject->DeviceExtension;

	io_stack		= IoGetCurrentIrpStackLocation(Irp);
	next_io_stack	= IoGetNextIrpStackLocation(Irp);

	*next_io_stack = *io_stack;

	IoSetCompletionRoutine(Irp, NULL, NULL, FALSE, FALSE, FALSE);

	switch(io_stack->MajorFunction)
	{
	case IRP_MJ_READ:
		next_io_stack->Parameters.Read.ByteOffset.QuadPart +=
			sizeof(union swap_header);
		break;

	case IRP_MJ_WRITE:
		next_io_stack->Parameters.Write.ByteOffset.QuadPart +=
			sizeof(union swap_header);
		break;

	case IRP_MJ_DEVICE_CONTROL:
		if (io_stack->Parameters.DeviceIoControl.IoControlCode ==
			IOCTL_DISK_GET_PARTITION_INFO)
		{
			IoSetCompletionRoutine(
				Irp,
				IoCompletion,
				DeviceObject,
				TRUE,
				FALSE,
				FALSE
				);
		}
	}

	return IoCallDriver(device_extension->target_device, Irp);
}

NTSTATUS
IoCompletion (
	IN PDEVICE_OBJECT	DeviceObject,
	IN PIRP				Irp,
	IN PVOID			Context
	)
{
	PPARTITION_INFORMATION partition;

	ASSERT(Irp->AssociatedIrp.SystemBuffer != NULL);

	partition = (PPARTITION_INFORMATION) Irp->AssociatedIrp.SystemBuffer;

	partition->PartitionLength.QuadPart -= sizeof(union swap_header);

	if (Irp->PendingReturned)
	{
		IoMarkIrpPending(Irp);
	}

	return Irp->IoStatus.Status;
}

#pragma code_seg("init")

BOOLEAN
IsDeviceMounted (
	IN PUNICODE_STRING DeviceName
	)
{
	NTSTATUS		status;
	PFILE_OBJECT	file_object;
	PDEVICE_OBJECT	device_object;

	ASSERT(DeviceName != NULL);

	status = IoGetDeviceObjectPointer(
		DeviceName,
		FILE_READ_ATTRIBUTES,
		&file_object,
		&device_object
		);

	if (!NT_SUCCESS(status))
	{
		return FALSE;
	}

	ObDereferenceObject(file_object);

	if (device_object->Vpb && (device_object->Vpb->Flags & VPB_MOUNTED))
	{
		KdPrint(("SwapFs: IsDeviceMounted: Device is already mounted\n"));

		return TRUE;
	}
	else
	{
		return FALSE;
	}
}

NTSTATUS
IsDeviceLinuxSwap (
	IN PDEVICE_OBJECT DeviceObject
	)
{
	union swap_header*	swap_header;
	LARGE_INTEGER		offset;
	NTSTATUS			status;
	UCHAR				magic[10];

	ASSERT(DeviceObject != NULL);

	swap_header = (union swap_header*)
		ExAllocatePool(NonPagedPoolCacheAligned, sizeof(union swap_header));

	if (!swap_header)
	{
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	offset.QuadPart = SWAP_HEADER_OFFSET;

	status = ReadBlockDevice(
		DeviceObject,
		&offset,
		sizeof(union swap_header),
		swap_header
		);

	RtlCopyMemory(magic, swap_header->magic.magic, 10);

	ExFreePool(swap_header);

	if (!NT_SUCCESS(status))
	{
		KdPrint(("SwapFs: IsDeviceLinuxSwap: "
			"ReadBlockDevice failed with error 0x%x\n", status));

		return status;
	}
	else if (RtlCompareMemory(magic, SWAP_HEADER_MAGIC_V1, 10) == 10 ||
			 RtlCompareMemory(magic, SWAP_HEADER_MAGIC_V2, 10) == 10)
	{
		KdPrint(("SwapFs: IsDeviceLinuxSwap: "
			"Linux swap device found\n"));

		return STATUS_SUCCESS;
	}
	else
	{
		KdPrint(("SwapFs: IsDeviceLinuxSwap: "
			"Device is not a Linux swap device\n"));

		return STATUS_UNRECOGNIZED_VOLUME;
	}
}

NTSTATUS
FormatDeviceToFat (
	IN PDEVICE_OBJECT DeviceObject
	)
{
	ULONG						size;
	NTSTATUS					status;
	DISK_GEOMETRY				disk_geometry;
	PARTITION_INFORMATION		partition_information;