// Main program for AIOWDM driver

#include "stddcls.h"
#include "driver.h"
#include <initguid.h>
#include "guids.h"
#ifndef INITGUID
	#include <initguid.h>
#endif

NTSTATUS AddDevice(IN PDRIVER_OBJECT DriverObject, IN PDEVICE_OBJECT pdo);
VOID DriverUnload(IN PDRIVER_OBJECT fdo);
NTSTATUS OnRequestComplete(IN PDEVICE_OBJECT fdo, IN PIRP Irp, IN PKEVENT pev);

BOOLEAN IsWin98();
BOOLEAN win98 = FALSE;
PDEVICE_OBJECT pCentralDevice = NULL;

//externed by ReadWrite
UNICODE_STRING servkey;
long SlotsOpen;

///////////////////////////////////////////////////////////////////////////////

#pragma INITCODE

/*
void NoteMyExistence(void)
{
	PKEY_VALUE_PARTIAL_INFORMATION pKVPI;
	HANDLE Key;
	OBJECT_ATTRIBUTES Attr;
	PWSTR Buffer = L"\\Registry\\Machine\\";
	UNICODE_STRING EKey;
	EKey.Buffer = Buffer;
	InitializeObjectAttributes(&Attr, &EKey, 0, NULL, NULL);
	NTSTATUS status = ZwOpenKey(&Key, KEY_QUERY_VALUE, &Attr);
	if (!NT_SUCCESS(status))
		return ;

	PWSTR Value = L"I Am";
	ZwSetValueKey(Key, NULL, 0, REG_SZ, Value, 10);

	ZwClose(Key);
}
*/

/*
void DumpRegDWord(unsigned long Value)
{
	PKEY_VALUE_PARTIAL_INFORMATION pKVPI;
	HANDLE Key;
	OBJECT_ATTRIBUTES Attr;
	PWSTR Buffer = L"\\Registry\\Machine\\Software\\PCIFind\\";
	UNICODE_STRING EKey;
	EKey.Buffer = Buffer;
	InitializeObjectAttributes(&Attr, &EKey, 0, NULL, NULL);
	NTSTATUS status = ZwOpenKey(&Key, KEY_QUERY_VALUE, &Attr);
	if (!NT_SUCCESS(status))
		return ;

	ZwSetValueKey(Key, NULL, 0, REG_DWORD, &Value, 4);

	ZwClose(Key);
}
*/

///////////////////////////////////////////////////////////////////////////////

#pragma INITCODE

PDEVICE_OBJECT AIOWDMCreateDevice(IN PDRIVER_OBJECT DriverObject)
{
	NTSTATUS Status;
	UNICODE_STRING NtDeviceName;
	UNICODE_STRING Win32DeviceName;
	PDEVICE_OBJECT pDevObj;

	RtlInitUnicodeString(&NtDeviceName, DEVICE_DEVICE_NAME);

	Status = IoCreateDevice(
		DriverObject,
		sizeof(DEVICE_EXTENSION),
		&NtDeviceName,
		FILE_DEVICE_UNKNOWN,
		0,
		FALSE, // Not exclusive
		&pDevObj
	);

	if (!NT_SUCCESS(Status)) return NULL; // Give up if create failed.

	RtlInitUnicodeString(&Win32DeviceName, AIOWDM_DEVICE_NAME);
	Status = IoCreateSymbolicLink( &Win32DeviceName, &NtDeviceName );
	if ( ! NT_SUCCESS(Status) ) // If we we couldn't create the link then abort installation.
	{
		IoDeleteDevice(pDevObj);
		return NULL;
	}

	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) pDevObj->DeviceExtension;
	pdx->portbase = 0;
	pdx->bRealDevice = FALSE;

	return pDevObj;
}

///////////////////////////////////////////////////////////////////////////////

#pragma INITCODE

extern "C" NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject,
	IN PUNICODE_STRING RegistryPath)
	{							// DriverEntry

	//NoteMyExistence();
	servkey.Buffer = NULL;

	// Insist that OS support at least the WDM level of the DDK we use

	if (!IoIsWdmVersionAvailable(1, 0))
		{
		KdPrint((DRIVERNAME " - Expected version of WDM (%d.%2.2d) not available\n", 1, 0));
		return STATUS_UNSUCCESSFUL;
		}

	// See if we're running under Win98 or NT:

	win98 = IsWin98();



	// Save the name of the service key

	servkey.Buffer = (PWSTR) ExAllocatePool(NonPagedPool, RegistryPath->Length + sizeof(WCHAR));
	if (!servkey.Buffer)
	{
		KdPrint((DRIVERNAME " - Unable to allocate %d bytes for copy of service key name\n", RegistryPath->Length + sizeof(WCHAR)));
		return STATUS_INSUFFICIENT_RESOURCES;
	}
	servkey.MaximumLength = RegistryPath->Length + sizeof(WCHAR);
	RtlCopyUnicodeString(&servkey, RegistryPath);

	// Initialize function pointers

	DriverObject->DriverUnload = DriverUnload;
	DriverObject->DriverExtension->AddDevice = AddDevice;

	DriverObject->MajorFunction[IRP_MJ_CREATE] = DispatchCreate;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = DispatchClose;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DispatchControl;
	DriverObject->MajorFunction[IRP_MJ_POWER] = DispatchPower;
	DriverObject->MajorFunction[IRP_MJ_PNP] = DispatchPnp;

	SlotsOpen = MAX_SLOTS_OPEN;
	for (int I = 0; I < MAX_SLOTS_OPEN; ++I)
		GlobalExtension[I] = NULL;

	//Create central device for global commands.
	pCentralDevice = AIOWDMCreateDevice(DriverObject);
	//DumpRegDWord(sizeof(TPORTOP));

	return STATUS_SUCCESS;
	}							// DriverEntry

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

VOID DriverUnload(IN PDRIVER_OBJECT DriverObject)
{							// DriverUnload
	PAGED_CODE();

	if ( pCentralDevice != NULL ) IoDeleteDevice(pCentralDevice);
	RtlFreeUnicodeString(&servkey);
}							// DriverUnload

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS AddDevice(IN PDRIVER_OBJECT DriverObject, IN PDEVICE_OBJECT pdo)
{							// AddDevice
	PAGED_CODE();

	NTSTATUS status;
	long IRQIndex;

	//This logic isn't truly multiprocessor safe, but it's fine for up to about
	//2G instances :)
	IRQIndex = InterlockedDecrement(&SlotsOpen);
	if (IRQIndex < 0)
	{
		InterlockedIncrement(&SlotsOpen);
		return STATUS_DEVICE_ALREADY_ATTACHED;
	}

	// Create a functional device object to represent the hardware we're managing.

	PDEVICE_OBJECT fdo;
	status = IoCreateDevice(DriverObject, sizeof(DEVICE_EXTENSION), NULL,
		FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN, FALSE, &fdo);
	if (!NT_SUCCESS(status))
	{						// can't create device object
		KdPrint((DRIVERNAME " - IoCreateDevice failed - %X\n", status));
		return status;
	}						// can't create device object
	
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;

	// From this point forward, any error will have side effects that need to
	// be cleaned up. Using a try-finally block allows us to modify the program
	// easily without losing track of the side effects.

	__try
	{						// finish initialization
		pdx->DeviceObject = fdo;
		pdx->Pdo = pdo;
		IoInitializeRemoveLock(&pdx->RemoveLock, 0, 0, 0);
		pdx->state = STOPPED;		// device starts in the stopped state

		pdx->DeviceID = 0;
		pdx->IAm.Operation = 'N';
		pdx->Disable.Operation = 'N';
		pdx->Clear.Operation = 'N';
		pdx->Enable.Operation = 'N';
		pdx->HasPendingIrp = FALSE;
		pdx->InterruptOccurred = FALSE;
		pdx->InterruptRequested = FALSE;
		pdx->IRQIndex = IRQIndex;
		pdx->Bench = 0;
		pdx->bRealDevice = TRUE;

		// Initialize DPC object

		IoInitializeDpcRequest(fdo, DpcForIsr);

		// Link our device object into the stack leading to the PDO
		
		pdx->LowerDeviceObject = IoAttachDeviceToDeviceStack(fdo, pdo);
		if (!pdx->LowerDeviceObject)
		{						// can't attach device
			KdPrint((DRIVERNAME " - IoAttachDeviceToDeviceStack failed\n"));
			status = STATUS_DEVICE_REMOVED;
			__leave;
		}						// can't attach device

		// Set power management flags in the device object

		fdo->Flags |= DO_POWER_PAGABLE;

		// Register a device interface

		status = IoRegisterDeviceInterface(pdo, &GUID_INTERFACE_AIOWDM, NULL, &pdx->ifname);
		if (!NT_SUCCESS(status))
		{						// unable to register interface
			KdPrint((DRIVERNAME " - IoRegisterDeviceInterface failed - %8.8lX\n", status));
			__leave;
		}						// unable to register interface

		// Indicate that our initial power state is D0 (fully on). Also indicate that
		// we have a pagable power handler (otherwise, we'll never get idle shutdown
		// messages!)

		pdx->syspower = PowerSystemWorking;
		pdx->devpower = PowerDeviceD0;
		POWER_STATE state;
		state.DeviceState = PowerDeviceD0;
		PoSetPowerState(fdo, DevicePowerState, state);

		// Clear the "initializing" flag so that we can get IRPs

		fdo->Flags &= ~DO_DEVICE_INITIALIZING;
	}						// finish initialization
	__finally
	{						// cleanup side effects
		if ( ! NT_SUCCESS(status) )
		{					// need to cleanup
			if (pdx->ifname.Buffer)
				RtlFreeUnicodeString(&pdx->ifname);
			if (pdx->LowerDeviceObject)
				IoDetachDevice(pdx->LowerDeviceObject);
			IoDeleteDevice(fdo);
		}					// need to cleanup
		else
			GlobalExtension[IRQIndex] = pdx;
	}						// cleanup side effects

	return status;
}							// AddDevice

///////////////////////////////////////////////////////////////////////////////

#pragma LOCKEDCODE

NTSTATUS CompleteRequest(IN PIRP Irp, IN NTSTATUS status, IN ULONG_PTR info)
	{							// CompleteRequest
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = info;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return status;
	}							// CompleteRequest

NTSTATUS CompleteRequest(IN PIRP Irp, IN NTSTATUS status)
	{							// CompleteRequest
	Irp->IoStatus.Status = status;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return status;
	}							// CompleteRequest

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS ForwardAndWait(IN PDEVICE_OBJECT fdo, IN PIRP Irp)
	{							// ForwardAndWait
	ASSERT(KeGetCurrentIrql() == PASSIVE_LEVEL);
	PAGED_CODE();
	
	KEVENT event;
	KeInitializeEvent(&event, NotificationEvent, FALSE);

	IoCopyCurrentIrpStackLocationToNext(Irp);
	IoSetCompletionRoutine(Irp, (PIO_COMPLETION_ROUTINE) OnRequestComplete,
		(PVOID) &event, TRUE, TRUE, TRUE);

	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	IoCallDriver(pdx->LowerDeviceObject, Irp);
	KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
	return Irp->IoStatus.Status;
	}							// ForwardAndWait

///////////////////////////////////////////////////////////////////////////////

#pragma LOCKEDCODE

NTSTATUS OnRequestComplete(IN PDEVICE_OBJECT fdo, IN PIRP Irp, IN PKEVENT pev)
	{							// OnRequestComplete
	KeSetEvent(pev, 0, FALSE);
	return STATUS_MORE_PROCESSING_REQUIRED;
	}							// OnRequestComplete

///////////////////////////////////////////////////////////////////////////////

VOID EnableAllInterfaces(PDEVICE_EXTENSION pdx, BOOLEAN enable)
	{							// EnableAllInterfaces
	IoSetDeviceInterfaceState(&pdx->ifname, enable);
	}							// EnableAllInterfaces

///////////////////////////////////////////////////////////////////////////////

VOID DeregisterAllInterfaces(PDEVICE_EXTENSION pdx)
	{							// DeregisterAllInterfaces
	IoSetDeviceInterfaceState(&pdx->ifname, FALSE);
	RtlFreeUnicodeString(&pdx->ifname);
	}							// DeregisterAllInterfaces

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

VOID RemoveDevice(IN PDEVICE_OBJECT fdo)
	{							// RemoveDevice
	PAGED_CODE();
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	NTSTATUS status;


	if (pdx->LowerDeviceObject)
		IoDetachDevice(pdx->LowerDeviceObject);

	IoDeleteDevice(fdo);
	}							// RemoveDevice

///////////////////////////////////////////////////////////////////////////////

#pragma INITCODE

BOOLEAN IsWin98()
	{							// IsWin98
#ifdef _X86_

	// Windows 98 (including 2d ed) supports WDM version 1.0, whereas Win2K
	// supports 1.10.

	return !IoIsWdmVersionAvailable(1, 0x10);
#else // not _X86_
	return FALSE;
#endif // not _X86_
	}							// IsWin98

///////////////////////////////////////////////////////////////////////////////

#if DBG && defined(_X86_)
#pragma LOCKEDCODE

extern "C" void __declspec(naked) __cdecl _chkesp()
	{
	_asm je okay
	ASSERT(!DRIVERNAME " - Stack pointer mismatch!");
okay:
	_asm ret
	}

#endif // DBG