#include "ntddk.h"
#include "irqcos.h"


/*********************************************************************************
*
*		DRIVER ENTRY
*
**********************************************************************************/

NTSTATUS
DriverEntry(
    IN PDRIVER_OBJECT DriverObject,
    IN PUNICODE_STRING RegistryPath
    )
{
    PDEVICE_OBJECT deviceObject = NULL;
    NTSTATUS status;
    UNICODE_STRING uniNtNameString;
    UNICODE_STRING uniWin32NameString;
    KAFFINITY Affinity;
    ULONG MappedVector;
    PIRQCOS_DEVICE_EXTENSION extension;

	// Create dispatch points for create/open, close, unload, and ioctl
	DriverObject->MajorFunction[IRP_MJ_CREATE] = IRQCOSOpenClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = IRQCOSOpenClose;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = IRQCOSDeviceControl;
	DriverObject->DriverUnload  = IRQCOSUnload;

    // Create counted string version of our device name.
    RtlInitUnicodeString( &uniNtNameString, NT_DEVICE_NAME );

    // Create the device object
    status = IoCreateDevice(DriverObject,
							sizeof(IRQCOS_DEVICE_EXTENSION),
							&uniNtNameString,
							FILE_DEVICE_KRNLDRVR,
							0,
							FALSE,
							&deviceObject
							);

    if ( NT_SUCCESS(status) )
    {
		// fill in the device extension
		extension = (PIRQCOS_DEVICE_EXTENSION) deviceObject->DeviceExtension;
		extension->DeviceObject = deviceObject;
		extension->PendingIrp = NULL;
		extension->ReturnData[0] =
		extension->ReturnData[1] =
		extension->ReturnData[2] =
		extension->ReturnData[3] =
		extension->ReturnData[4] =
		extension->ReturnData[5] = 0;
		extension->InterruptRequested =
		extension->InterruptOccurred =
		extension->InterruptConnected = 
		extension->HasPendingIrp = FALSE;
		extension->BaseAddress = 0;

		// tell I/O Manager we want it to read/write data via buffers (NOT relevant to DeviceIOControl)
		deviceObject->Flags |= DO_BUFFERED_IO;

		// register DPC routine
//		IoInitializeDpcRequest(deviceObject, IRQCOSDpcForIsrRoutine);
		
        // Create counted string version of our Win32 device name.
        RtlInitUnicodeString( &uniWin32NameString, DOS_DEVICE_NAME );
    
        // Create a link from our device name to a name in the Win32 namespace.
        status = IoCreateSymbolicLink( &uniWin32NameString, &uniNtNameString );

        if (!NT_SUCCESS(status))
        {
			IoDeleteDevice( DriverObject->DeviceObject );
        }
    }
    return status;
}

/*********************************************************************************
*
*		OPEN / CLOSE (nothing happens so they're the same)
*
**********************************************************************************/

NTSTATUS
IRQCOSOpenClose(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    )
{
    Irp->IoStatus.Status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;

    IoCompleteRequest( Irp, IO_NO_INCREMENT );

    return STATUS_SUCCESS;
}

/*********************************************************************************
*
*		UNLOAD
*
**********************************************************************************/

VOID
IRQCOSUnload(
    IN PDRIVER_OBJECT DriverObject
    )
{
    UNICODE_STRING uniWin32NameString;
    PIRQCOS_DEVICE_EXTENSION extension;

    extension = DriverObject->DeviceObject->DeviceExtension;

    // Disconnect the interrupt
    if (extension->InterruptConnected)
		IoDisconnectInterrupt(extension->InterruptObject);

    // Create counted string version of our Win32 device name.
    RtlInitUnicodeString( &uniWin32NameString, DOS_DEVICE_NAME );

    // Delete the link from our device name to a name in the Win32 namespace.
    IoDeleteSymbolicLink( &uniWin32NameString );

    // Finally delete our device object
    IoDeleteDevice( DriverObject->DeviceObject );
}

/*********************************************************************************
*
*		DEVICE CONTROL
*
**********************************************************************************/

NTSTATUS
IRQCOSDeviceControl(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    )
{
    NTSTATUS status, ioConnectStatus;
    PIO_STACK_LOCATION irpSp;
    PIRQCOS_DEVICE_EXTENSION extension;
    ULONG ioControlCode;
    KIRQL irql;
    KAFFINITY Affinity;
    ULONG MappedVector;
	PULONG pIOBuffer;
	ULONG InBufferSize, OutBufferSize;
	ULONG nPort;
	USHORT BusNumber;
	short BusType;

    extension = DeviceObject->DeviceExtension;

    irpSp = IoGetCurrentIrpStackLocation( Irp );
	Irp->IoStatus.Information = 0;

    // Save the control code in the extension so we can use it later
    ioControlCode = irpSp->Parameters.DeviceIoControl.IoControlCode;

    switch (ioControlCode)
	{
		case IOCTL_IRQCOS_INIT_COS_DRIVER:
		    InBufferSize  = irpSp->Parameters.DeviceIoControl.InputBufferLength;
			if (InBufferSize < sizeof(IRQCOS_INIT_DATA))
			{
				status = STATUS_INVALID_PARAMETER;
				break;
			}

			if (extension->InterruptConnected)
			{
			// this might crash if an interrupt is waiting but it is up to the user to disable
			// interrupts before re-initializing an IRQ
				extension->InterruptRequested = FALSE;
				IoDisconnectInterrupt(extension->InterruptObject);
				if (extension->HasPendingIrp)
				{
					extension->PendingIrp->IoStatus.Status = STATUS_CANCELLED;
					extension->PendingIrp->IoStatus.Information = 0;
					IoSetCancelRoutine(extension->PendingIrp, NULL);
					IoCompleteRequest(extension->PendingIrp, IO_NO_INCREMENT);
					extension->HasPendingIrp = FALSE;
				}
			}

			pIOBuffer = (PULONG)Irp->AssociatedIrp.SystemBuffer;
			irql = extension->IRQ = ((PIRQCOS_INIT_DATA)pIOBuffer)->IRQ;
			BusType = ((PIRQCOS_INIT_DATA)pIOBuffer)->BusType;
			BusNumber = ((PIRQCOS_INIT_DATA)pIOBuffer)->BusNumber;
			extension->BaseAddress = ((PIRQCOS_INIT_DATA)pIOBuffer)->BaseAddress;
			//
			// This call will map our IRQ to a system vector. It will also fill
			// in the IRQL (the kernel-defined level at which our ISR will run),
			// and affinity mask (which processors our ISR can run on).
			//
			// We need to do this so that when we connect to the interrupt, we
			// can supply the kernel with this information.
			//
			MappedVector = HalGetInterruptVector(
				BusType,	// Interface type
				BusNumber,	// Bus number
				irql, // Bus interrupt level
				irql, // Bus interrupt vector
				&irql,	// IRQ level
				&Affinity	// Affinity mask
				);
			ioConnectStatus = IoConnectInterrupt(
				&extension->InterruptObject,
				IRQCOSInterruptServiceRoutine,
				extension->DeviceObject,
				NULL,
				MappedVector,
				irql,
				irql,
				Latched,
				FALSE,
				Affinity,
				FALSE
				);
			status = ioConnectStatus;

			if ( NT_SUCCESS (status) )
			{
				extension->InterruptConnected = TRUE;
			}
			break;
			
		case IOCTL_IRQCOS_GET_DATA:
			OutBufferSize = irpSp->Parameters.DeviceIoControl.OutputBufferLength;
			if (OutBufferSize < sizeof(extension->ReturnData))
			{
				status = STATUS_INVALID_PARAMETER;
			}
			else if (!extension->InterruptConnected || extension->InterruptRequested)
			{
				status = STATUS_INVALID_DEVICE_REQUEST;
			}
			else if (extension->InterruptOccurred)
			{
				pIOBuffer = (PULONG)Irp->AssociatedIrp.SystemBuffer;
				RtlCopyMemory(pIOBuffer, extension->ReturnData, sizeof(extension->ReturnData));
				Irp->IoStatus.Information = sizeof(extension->ReturnData);
				
				extension->InterruptOccurred = FALSE;
				status = STATUS_SUCCESS;
			}
			else
			{
				extension->InterruptRequested = TRUE;
				extension->PendingIrp = Irp;
				extension->HasPendingIrp = TRUE;
				IoMarkIrpPending(Irp);
				IoSetCancelRoutine(Irp, IRQCOSCancelIrp);
				return STATUS_PENDING;	// do not complete Irp of course
			}
			break;
		
		case IOCTL_IRQCOS_ABORT:
			if (!extension->InterruptRequested)
				status = STATUS_INVALID_DEVICE_REQUEST;
			else
			{
				extension->PendingIrp->IoStatus.Information = 0;
				extension->PendingIrp->IoStatus.Status = STATUS_SUCCESS;
				IoSetCancelRoutine(extension->PendingIrp, NULL);
				IoCompleteRequest(extension->PendingIrp, IO_SERIAL_INCREMENT);
				extension->InterruptRequested = FALSE;
				extension->HasPendingIrp = FALSE;
			}
			break;
				
		default:
			status = STATUS_INVALID_DEVICE_REQUEST;
			break;
	}

	Irp->IoStatus.Status = status;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return status;
}

/*********************************************************************************
*
*		INTERRUPT SERVICE ROUTINE
*
**********************************************************************************/

BOOLEAN
IRQCOSInterruptServiceRoutine(
    IN PKINTERRUPT Interrupt,
    IN OUT PVOID Context
    )
{
    PDEVICE_OBJECT DeviceObject;
    PIRQCOS_DEVICE_EXTENSION extension;
	ULONG nPort;
	PULONG pIOBuffer;

    DeviceObject = Context;
    extension = DeviceObject->DeviceExtension;

	if (extension->BaseAddress != 0)
	{
		nPort = extension->BaseAddress;
		extension->ReturnData[0] = READ_PORT_UCHAR((PUCHAR)nPort++);
		extension->ReturnData[1] = READ_PORT_UCHAR((PUCHAR)nPort++);
		extension->ReturnData[2] = READ_PORT_UCHAR((PUCHAR)nPort++);
		nPort++;
		extension->ReturnData[3] = READ_PORT_UCHAR((PUCHAR)nPort++);
		extension->ReturnData[4] = READ_PORT_UCHAR((PUCHAR)nPort++);
		extension->ReturnData[5] = READ_PORT_UCHAR((PUCHAR)nPort++);
	}

	if (extension->InterruptRequested && extension->HasPendingIrp)
	{
		extension->InterruptRequested = FALSE;
		extension->HasPendingIrp = FALSE;

		pIOBuffer = (PULONG)extension->PendingIrp->AssociatedIrp.SystemBuffer;
		RtlCopyMemory(pIOBuffer, extension->ReturnData, sizeof(extension->ReturnData));
		extension->PendingIrp->IoStatus.Information = sizeof(extension->ReturnData);
		extension->PendingIrp->IoStatus.Status = STATUS_SUCCESS;
		IoSetCancelRoutine(extension->PendingIrp, NULL);
		IoCompleteRequest(extension->PendingIrp, IO_SERIAL_INCREMENT);
	}
	else
		extension->InterruptOccurred = TRUE;

/*	IoRequestDpc(DeviceObject,
		NULL,
		NULL
		);*/

	nPort = extension->BaseAddress + 0x0F;
	WRITE_PORT_UCHAR((PUCHAR)nPort, 0); 

	if (extension->IRQ >= 8 && extension->IRQ <= 15)
	{
		nPort = 0xA0;
		WRITE_PORT_UCHAR((PUCHAR)nPort, 0x20);
	}
	nPort = 0x20;
	WRITE_PORT_UCHAR((PUCHAR)nPort, 0x20);

    return TRUE;
}

/*********************************************************************************
*
*		DEFERRED PROCEDURE CALL FOR ISR
*
**********************************************************************************/

VOID
IRQCOSDpcForIsrRoutine(
    IN PKDPC Dpc,
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp,
    IN PVOID Context
    )
{
    PIRQCOS_DEVICE_EXTENSION extension;
	PULONG pIOBuffer;
	ULONG nPort;

    extension = DeviceObject->DeviceExtension;

	nPort = extension->BaseAddress + 0x0F;
	WRITE_PORT_UCHAR((PUCHAR)nPort, 0); 

	if (extension->IRQ >= 8 && extension->IRQ <= 15)
	{
		nPort = 0xA0;
		WRITE_PORT_UCHAR((PUCHAR)nPort, 0x20);
	}
	nPort = 0x20;
	WRITE_PORT_UCHAR((PUCHAR)nPort, 0x20);

    return;
}

/*********************************************************************************
*
*		CANCEL IRP
*
**********************************************************************************/

VOID
IRQCOSCancelIrp(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    )
{
    PIRQCOS_DEVICE_EXTENSION extension;

    extension = DeviceObject->DeviceExtension;

	Irp->IoStatus.Information = 0;

	if (Irp == extension->PendingIrp)
	{
		extension->HasPendingIrp = FALSE;
	}
	IoReleaseCancelSpinLock( Irp->CancelIrql );

	Irp->IoStatus.Status = STATUS_CANCELLED;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
}