//-----------------------------------------------------------------------------
//     Author : hiyohiyo
//       Mail : hiyohiyo@crystalmark.info
//        Web : http://openlibsys.org/
//    License : The modified BSD license
//
//                     Copyright 2007-2008 OpenLibSys.org. All rights reserved.
//-----------------------------------------------------------------------------

#include <ntddk.h>
#include <stddef.h>
#include "OpenLibSys.h"

//-----------------------------------------------------------------------------
//
// Global
//
//-----------------------------------------------------------------------------

static ULONG refCount;

//-----------------------------------------------------------------------------
//
// Classic NT driver
// DriverEntry / OlsDispatch / Unload
//
//-----------------------------------------------------------------------------

NTSTATUS
DriverEntry(
	IN PDRIVER_OBJECT  DriverObject,
	IN PUNICODE_STRING RegistryPath
	)

/*

Return Value:

	STATUS_SUCCESS if the driver initialized correctly, otherwise an erroror
	indicating the reason for failure.
*/

{
	NTSTATUS		status;
	UNICODE_STRING  ntDeviceName;
	UNICODE_STRING  win32DeviceName;
	PDEVICE_OBJECT  deviceObject = NULL;

	RtlInitUnicodeString(&ntDeviceName, NT_DEVICE_NAME);

	status = IoCreateDevice(
		DriverObject,					// Our Driver Object
		0,								// We don't use a device extension
		&ntDeviceName,					// Device name 
		OLS_TYPE,						// Device type
		FILE_DEVICE_SECURE_OPEN,		// Device characteristics
		FALSE,							// Not an exclusive device
		&deviceObject );				// Returned ptr to Device Object

	if(!NT_SUCCESS(status))
	{
		refCount = (ULONG)-1;
		return status;
	}
	else
	{
		refCount = 0;
	}

	// Initialize the driver object with this driver's entry points.
	DriverObject->MajorFunction[IRP_MJ_CREATE] = OlsDispatch;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = OlsDispatch;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = OlsDispatch;
	DriverObject->DriverUnload = Unload;

	// Initialize a Unicode String containing the Win32 name for our device.
	RtlInitUnicodeString(&win32DeviceName, DOS_DEVICE_NAME);

	// Create a symbolic link between our device name  and the Win32 name
	status = IoCreateSymbolicLink(&win32DeviceName, &ntDeviceName);

	if (!NT_SUCCESS(status))
	{
		// Delete everything that this routine has allocated.
		IoDeleteDevice( deviceObject );
	}

	return status;
}

NTSTATUS
OlsDispatch(
	IN	PDEVICE_OBJECT pDO,
	IN	PIRP pIrp
	)

/*++

Routine Description:
	This routine is the dispatch handler for the driver.  It is responsible
	for processing the IRPs.

Arguments:
	
	pDO - Pointer to device object.

	pIrp - Pointer to the current IRP.

Return Value:

	STATUS_SUCCESS if the IRP was processed successfully, otherwise an erroror
	indicating the reason for failure.

--*/

{
	PIO_STACK_LOCATION pIrpStack;
	NTSTATUS status;
	int index;

	//  Initialize the irp info field.
	//	  This is used to return the number of bytes transfered.
	pIrp->IoStatus.Information = 0;
	pIrpStack = IoGetCurrentIrpStackLocation(pIrp);

	//  Set default return status
	status = STATUS_NOT_IMPLEMENTED;

	// Dispatch based on major fcn code.
	switch(pIrpStack->MajorFunction)
	{
		case IRP_MJ_CREATE:
			if(refCount != (ULONG)-1){refCount++;}
			status = STATUS_SUCCESS;
			break;
		case IRP_MJ_CLOSE:
			if(refCount != (ULONG)-1){refCount--;}
			status = STATUS_SUCCESS;
			break;

		case IRP_MJ_DEVICE_CONTROL:
			//  Dispatch on IOCTL
			switch(pIrpStack->Parameters.DeviceIoControl.IoControlCode)
			{
			case IOCTL_OLS_GET_DRIVER_VERSION:
				*(PULONG)pIrp->AssociatedIrp.SystemBuffer = OLS_DRIVER_VERSION;
				pIrp->IoStatus.Information = 4;
				status = STATUS_SUCCESS;
				break;

			case IOCTL_OLS_GET_REFCOUNT:
				*(PULONG)pIrp->AssociatedIrp.SystemBuffer = refCount;
				pIrp->IoStatus.Information = sizeof(refCount);
				status = STATUS_SUCCESS;
				break;

			case IOCTL_OLS_READ_MSR:
				status = ReadMsr(
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.InputBufferLength,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.OutputBufferLength,
					(ULONG*)&pIrp->IoStatus.Information
					);
				break;
			case IOCTL_OLS_WRITE_MSR:
				status = WriteMsr(
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.InputBufferLength,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.OutputBufferLength,
					(ULONG*)&pIrp->IoStatus.Information
					);
				break;
			case IOCTL_OLS_READ_PMC:
				status = ReadPmc(
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.InputBufferLength,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.OutputBufferLength,
					(ULONG*)&pIrp->IoStatus.Information
					);
				break;
			case IOCTL_OLS_HALT:
				__halt();
				status = STATUS_SUCCESS;
				break;

			case IOCTL_OLS_READ_IO_PORT:
			case IOCTL_OLS_READ_IO_PORT_BYTE:
			case IOCTL_OLS_READ_IO_PORT_WORD:
			case IOCTL_OLS_READ_IO_PORT_DWORD:
				status = ReadIoPort(
					pIrpStack->Parameters.DeviceIoControl.IoControlCode,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.InputBufferLength,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.OutputBufferLength,
					(ULONG*)&pIrp->IoStatus.Information
					);
				break;
			case IOCTL_OLS_WRITE_IO_PORT:
			case IOCTL_OLS_WRITE_IO_PORT_BYTE:
			case IOCTL_OLS_WRITE_IO_PORT_WORD:
			case IOCTL_OLS_WRITE_IO_PORT_DWORD:
				status = WriteIoPort(
					pIrpStack->Parameters.DeviceIoControl.IoControlCode,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.InputBufferLength,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.OutputBufferLength,
					(ULONG*)&pIrp->IoStatus.Information
					);
				break;

			case IOCTL_OLS_READ_PCI_CONFIG:
				status = ReadPciConfig(
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.InputBufferLength,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.OutputBufferLength,
					(ULONG*)&pIrp->IoStatus.Information
					);
				break;
			case IOCTL_OLS_WRITE_PCI_CONFIG:
				status = WritePciConfig(
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.InputBufferLength,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.OutputBufferLength,
					(ULONG*)&pIrp->IoStatus.Information
					);
				break;

			case IOCTL_OLS_READ_MEMORY:
				status = ReadMemory(
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.InputBufferLength,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.OutputBufferLength,
					(ULONG*)&pIrp->IoStatus.Information
					);
				break;
			case IOCTL_OLS_WRITE_MEMORY:
				status = WriteMemory(
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.InputBufferLength,
					pIrp->AssociatedIrp.SystemBuffer,
					pIrpStack->Parameters.DeviceIoControl.OutputBufferLength,
					(ULONG*)&pIrp->IoStatus.Information
					);
				break;


			}
			break;
	}

	// We're done with I/O request.  Record the status of the I/O action.
	pIrp->IoStatus.Status = status;

	// Don't boost priority when returning since this took little time.
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);

	return status;
}

VOID
Unload(
	PDRIVER_OBJECT DriverObject
	)
/*++

Routine Description:

	This routine is called by the I/O system to unload the driver.

	Any resources previously allocated must be freed.

Arguments:

	DriverObject - a pointer to the object that represents our driver.

Return Value:

	None
--*/

{
	PDEVICE_OBJECT deviceObject = DriverObject->DeviceObject;
	UNICODE_STRING win32NameString;

	PAGED_CODE();

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

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

	if(deviceObject != NULL)
	{
		IoDeleteDevice(deviceObject);
	}
}

//-----------------------------------------------------------------------------
//
// CPU
//
//-----------------------------------------------------------------------------

NTSTATUS
ReadMsr(	void	*lpInBuffer, 
			ULONG	nInBufferSize, 
			void	*lpOutBuffer, 
			ULONG	nOutBufferSize, 
			ULONG	*lpBytesReturned)
{
	__try
	{
		ULONGLONG data = __readmsr(*(ULONG*)lpInBuffer);
		memcpy((PULONG)lpOutBuffer, &data, 8);
		*lpBytesReturned = 8;
		return STATUS_SUCCESS;
	}
	__except(EXCEPTION_EXECUTE_HANDLER)
	{
		*lpBytesReturned = 0;
		return STATUS_UNSUCCESSFUL;
	}
}

NTSTATUS
WriteMsr(	void	*lpInBuffer, 
			ULONG	nInBufferSize, 
			void	*lpOutBuffer, 
			ULONG	nOutBufferSize, 
			ULONG	*lpBytesReturned)
{
	__try
	{
		OLS_WRITE_MSR_INPUT* param;
		param = (OLS_WRITE_MSR_INPUT*)lpInBuffer;

		__writemsr(param->Register, param->Value.QuadPart);
		*lpBytesReturned = 0;
		return STATUS_SUCCESS;
	}
	__except(EXCEPTION_EXECUTE_HANDLER)
	{
		*lpBytesReturned = 0;
		return STATUS_UNSUCCESSFUL;
	}
}

NTSTATUS
ReadPmc(	void	*lpInBuffer, 
			ULONG	nInBufferSize, 
			void	*lpOutBuffer, 
			ULONG	nOutBufferSize, 
			ULONG	*lpBytesReturned)
{
	__try
	{
		ULONGLONG data = __readpmc(*(ULONG*)lpInBuffer);
		memcpy((PULONG)lpOutBuffer, &data, 8);
		*lpBytesReturned = 8;
		return STATUS_SUCCESS;
	}
	__except(EXCEPTION_EXECUTE_HANDLER)
	{
		*lpBytesReturned = 0;
		return STATUS_UNSUCCESSFUL;
	}
}

//-----------------------------------------------------------------------------
//
// IO Port
//
//-----------------------------------------------------------------------------

NTSTATUS
ReadIoPort( ULONG	ioControlCode,
			void	*lpInBuffer, 
			ULONG	nInBufferSize, 
			void	*lpOutBuffer, 
			ULONG	nOutBufferSize, 
			ULONG	*lpBytesReturned)
{
	ULONG nPort = *(ULONG*)lpInBuffer;

	switch(ioControlCode)
	{
		case IOCTL_OLS_READ_IO_PORT_BYTE:
			*(PUCHAR)lpOutBuffer = READ_PORT_UCHAR((PUCHAR)(ULONG_PTR)nPort);
			break;
		case IOCTL_OLS_READ_IO_PORT_WORD:
			*(PUSHORT)lpOutBuffer = READ_PORT_USHORT((PUSHORT)(ULONG_PTR)nPort);
			break;
		case IOCTL_OLS_READ_IO_PORT_DWORD:
			*(PULONG)lpOutBuffer = READ_PORT_ULONG((PULONG)(ULONG_PTR)nPort);
			break;
		default:
			*lpBytesReturned = 0;
			return STATUS_INVALID_PARAMETER;
			break;
	}
	
	*lpBytesReturned = nInBufferSize;
	return STATUS_SUCCESS;
}

NTSTATUS
WriteIoPort(ULONG	ioControlCode,
			void	*lpInBuffer, 
			ULONG	nInBufferSize, 
			void	*lpOutBuffer, 
			ULONG	nOutBufferSize, 
			ULONG	*lpBytesReturned)
{
	ULONG nPort;
	OLS_WRITE_IO_PORT_INPUT* param;
	
	param = (OLS_WRITE_IO_PORT_INPUT*)lpInBuffer;
	nPort = param->PortNumber;

	switch(ioControlCode)
	{

		case IOCTL_OLS_WRITE_IO_PORT_BYTE:
			WRITE_PORT_UCHAR((PUCHAR)(ULONG_PTR)nPort, param->CharData);
			break;
		case IOCTL_OLS_WRITE_IO_PORT_WORD:
			WRITE_PORT_USHORT((PUSHORT)(ULONG_PTR)nPort, param->ShortData);
			break;
		case IOCTL_OLS_WRITE_IO_PORT_DWORD:
			WRITE_PORT_ULONG((PULONG)(ULONG_PTR)nPort, param->LongData);
			break;
		default:
			return STATUS_INVALID_PARAMETER;
			break;
	}

	return STATUS_SUCCESS;
}

//-----------------------------------------------------------------------------
//
// PCI
//
//-----------------------------------------------------------------------------

NTSTATUS
ReadPciConfig(	void	*lpInBuffer, 
				ULONG	nInBufferSize, 
				void	*lpOutBuffer, 
				ULONG	nOutBufferSize, 
				ULONG	*lpBytesReturned)
{
	OLS_READ_PCI_CONFIG_INPUT *param;
	NTSTATUS status;

	if(nInBufferSize != sizeof(OLS_READ_PCI_CONFIG_INPUT))
	{
		return STATUS_INVALID_PARAMETER;
	}
	param = (OLS_READ_PCI_CONFIG_INPUT *)lpInBuffer;

	status = pciConfigRead(param->PciAddress, param->PciOffset,
						lpOutBuffer, nOutBufferSize);

	if(status == STATUS_SUCCESS)
	{
		*lpBytesReturned = nOutBufferSize;
	}
	else
	{
		*lpBytesReturned = 0;
	}

	return status;
}

NTSTATUS
WritePciConfig(	void *lpInBuffer, 
				ULONG nInBufferSize, 
				void *lpOutBuffer, 
				ULONG nOutBufferSize, 
				ULONG *lpBytesReturned)

{
	OLS_WRITE_PCI_CONFIG_INPUT *param;
	ULONG writeSize;
	NTSTATUS status;

	if(nInBufferSize < offsetof(OLS_WRITE_PCI_CONFIG_INPUT, Data))
	{
		return STATUS_INVALID_PARAMETER;
	}

	param = (OLS_WRITE_PCI_CONFIG_INPUT *)lpInBuffer;
	writeSize = nInBufferSize - offsetof(OLS_WRITE_PCI_CONFIG_INPUT, Data);
	
	*lpBytesReturned = 0;

	return pciConfigWrite(param->PciAddress, param->PciOffset,
							&param->Data, writeSize);

}

//-----------------------------------------------------------------------------
//
// Support Function
//
//-----------------------------------------------------------------------------

NTSTATUS pciConfigRead(ULONG pciAddress, ULONG offset, void *data, int length)
{
	PCI_SLOT_NUMBER slot;
	int error;
	ULONG busNumber;

	busNumber = PciGetBus(pciAddress);
	slot.u.AsULONG = 0;
	slot.u.bits.DeviceNumber = PciGetDev(pciAddress);
	slot.u.bits.FunctionNumber = PciGetFunc(pciAddress);
	error =	HalGetBusDataByOffset(PCIConfiguration, busNumber, slot.u.AsULONG,
									data, offset, length);

	if(error == 0)
	{
		return OLS_ERROR_PCI_BUS_NOT_EXIST;
	}
	else if(length != 2 && error == 2)
	{
		return OLS_ERROR_PCI_NO_DEVICE;
	}
	else if(length != error)
	{
		return OLS_ERROR_PCI_READ_CONFIG;
	}

	return STATUS_SUCCESS;
}

NTSTATUS pciConfigWrite(ULONG pciAddress, ULONG offset, void *data, int length)
{
	PCI_SLOT_NUMBER slot;
	int error;
	ULONG busNumber;

	busNumber = PciGetBus(pciAddress);

	slot.u.AsULONG = 0;
	slot.u.bits.DeviceNumber = PciGetDev(pciAddress);
	slot.u.bits.FunctionNumber = PciGetFunc(pciAddress);
	error = HalSetBusDataByOffset(PCIConfiguration, busNumber, slot.u.AsULONG,
									data, offset, length);

	if(error != length)
	{
		return OLS_ERROR_PCI_WRITE_CONFIG;
	}

	return STATUS_SUCCESS;
}


//-----------------------------------------------------------------------------
//
// Physical Memory
//
//-----------------------------------------------------------------------------

NTSTATUS
ReadMemory(	void	*lpInBuffer, 
			ULONG	nInBufferSize, 
			void	*lpOutBuffer, 
			ULONG	nOutBufferSize, 
			ULONG	*lpBytesReturned)
{
	OLS_READ_MEMORY_INPUT *param;
	ULONG	size;
	PHYSICAL_ADDRESS address;
	PVOID	maped;
	BOOLEAN	error;

	if(nInBufferSize != sizeof(OLS_READ_MEMORY_INPUT))
	{
		return STATUS_INVALID_PARAMETER;
	}

	param = (OLS_READ_MEMORY_INPUT *)lpInBuffer;
	size = param->UnitSize * param->Count;

	if(nOutBufferSize < size)
	{
		return STATUS_INVALID_PARAMETER;
	}

	address.QuadPart = param->Address.QuadPart;

#ifndef _PHYSICAL_MEMORY_SUPPORT

	if(0x000C0000 > address.QuadPart 
	|| (address.QuadPart + size - 1) > 0x000FFFFF)
	{
		return STATUS_INVALID_PARAMETER;
	}

#endif

	maped = MmMapIoSpace(address, size, FALSE);

	error = FALSE;
	switch(param->UnitSize){
		case 1:
			READ_REGISTER_BUFFER_UCHAR(maped, lpOutBuffer, param->Count);
			break;
		case 2:
			READ_REGISTER_BUFFER_USHORT(maped, lpOutBuffer, param->Count);
			break;
		case 4:
			READ_REGISTER_BUFFER_ULONG(maped, lpOutBuffer, param->Count);
			break;
		default:
			error = TRUE;
			break;
	}

	MmUnmapIoSpace(maped, size);

	if(error)
	{
		return STATUS_INVALID_PARAMETER;
	}

	*lpBytesReturned = nOutBufferSize;

	return STATUS_SUCCESS;
}

NTSTATUS
WriteMemory(void	*lpInBuffer, 
			ULONG	nInBufferSize, 
			void	*lpOutBuffer, 
			ULONG	nOutBufferSize, 
			ULONG	*lpBytesReturned)
{
#ifdef _PHYSICAL_MEMORY_SUPPORT

	OLS_WRITE_MEMORY_INPUT *param;
	ULONG size;
	PHYSICAL_ADDRESS address;
	PVOID	maped;
	BOOLEAN	error;

	if(nInBufferSize < offsetof(OLS_WRITE_MEMORY_INPUT, Data))
	{
		return STATUS_INVALID_PARAMETER;
	}

	param = (OLS_WRITE_MEMORY_INPUT *)lpInBuffer;

	size = param->UnitSize * param->Count;
	if (nInBufferSize < size + offsetof(OLS_WRITE_MEMORY_INPUT, Data))
	{
		return STATUS_INVALID_PARAMETER;
	}

	address.QuadPart = param->Address.QuadPart;

	maped = MmMapIoSpace(address, size, FALSE);

	error = FALSE;
	switch(param->UnitSize){
		case 1:
			WRITE_REGISTER_BUFFER_UCHAR(maped, 
										(UCHAR*)&param->Data, param->Count);
			break;
		case 2:
			WRITE_REGISTER_BUFFER_USHORT(maped,
										(USHORT*)&param->Data, param->Count);
			break;
		case 4:
			WRITE_REGISTER_BUFFER_ULONG(maped,
										(ULONG*)&param->Data, param->Count);
			break;
		default:
			error = TRUE;
			break;
	}

	MmUnmapIoSpace(maped, size);

	if(error)
	{
		return STATUS_INVALID_PARAMETER;
	}

	*lpBytesReturned = 0;

	return STATUS_SUCCESS;

#else

	*lpBytesReturned = 0;
	
	return STATUS_INVALID_PARAMETER;

#endif
}