/***************************************************************
程序名稱：Hello World for WDM
文件名稱：HelloWDM.cpp
作者：羅聰
日期：2002-8-16
***************************************************************/

//一定要的頭文件，聲明了函數模塊和變量：
#include "HelloWDM.h"
#include "Ioctl.h"


KSPIN_LOCK BufferLock;
PUCHAR	Buffer = NULL;
ULONG	BufferSize = 0;

/***************************************************************
函數名稱：DriverEntry()
功能描述：WDM程序入口
***************************************************************/
//extern "C"是必須的，表示“用C鏈接”。如果你的文件名是HelloWDM.c的話，這句可以省略。
extern "C"
NTSTATUS DriverEntry(    IN PDRIVER_OBJECT DriverObject,
                        IN PUNICODE_STRING RegistryPath)
{
	NTSTATUS StatusL = STATUS_SUCCESS; 
	OBJECT_ATTRIBUTES ObjectAttributesL; 
	UNICODE_STRING UniFileNameL; 
	IO_STATUS_BLOCK IoStatusBlockL;
	HANDLE hFile ; 

	PCWSTR FileName = L"\\\\??\\\\d:\\\\readme.dat";
    DbgPrint("Starting DriverEntry()\n");
    
    //指定“添加設備”消息由函數“HelloWDMAddDevice()”來處理：
    DriverObject->DriverExtension->AddDevice = HelloWDMAddDevice;
    DriverObject->DriverUnload = Wdm1Unload;
    //指定“即插即用”消息由函數“HelloWDMPnp()”來處理：
    DriverObject->MajorFunction[IRP_MJ_PNP] = HelloWDMPnp;
    
    DriverObject->MajorFunction[IRP_MJ_CREATE] = Wdm1Create;
    DriverObject->MajorFunction[IRP_MJ_CLOSE] = Wdm1Close;

	DriverObject->MajorFunction[IRP_MJ_READ] = Wdm1Read;
	DriverObject->MajorFunction[IRP_MJ_WRITE] = Wdm1Write;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = Wdm1DeviceControl;
	
	KeInitializeSpinLock(&BufferLock);
	DbgPrint("DriverEntry completed");    
	
	
	//RtlInitUnicodeString(&UniFileNameL,L"\\\\??\\\\C:\\\\TestDat.dat\");
	RtlInitUnicodeString(&UniFileNameL,FileName);
	InitializeObjectAttributes(&ObjectAttributesL,&UniFileNameL, OBJ_CASE_INSENSITIVE,NULL,NULL); 

	StatusL = ZwCreateFile(&hFile, 
			  GENERIC_WRITE|SYNCHRONIZE|FILE_APPEND_DATA, 
			  &ObjectAttributesL, 
			  &IoStatusBlockL, 
			  0, 
			  FILE_ATTRIBUTE_NORMAL, 
			  FILE_SHARE_DELETE, 
			  FILE_OPEN_IF, 
			  FILE_SYNCHRONOUS_IO_NONALERT, 
			  NULL, 
			  0); 
	if(!NT_SUCCESS(StatusL)) 
	{ 
	    //DbgPrint("Cann\'t ZwCreateFile %s : 0x%x\\n\"); 
	}
	//返回一個NTSTATUS值STATUS_SUCCESS。幾乎所有的驅動程序例程都必須返回一個NTSTATUS值，這些值在NTSTATUS.H DDK頭文件中有詳細的定義。
    return STATUS_SUCCESS;
}

NTSTATUS Wdm1Create(IN PDEVICE_OBJECT fdo,
					IN PIRP Irp)
{
	PIO_STACK_LOCATION IrpStack = IoGetCurrentIrpStackLocation(Irp);
	DbgPrint( "Create File is %T", &(IrpStack->FileObject->FileName));

	// Complete successfully
	return CompleteIrp(Irp,STATUS_SUCCESS,0);
}

NTSTATUS Wdm1Close(	IN PDEVICE_OBJECT fdo,
					IN PIRP Irp)
{
	DbgPrint("Wdm1Close");

	// Complete successfully
	return CompleteIrp(Irp,STATUS_SUCCESS,0);
}

NTSTATUS Wdm1DeviceControl(	IN PDEVICE_OBJECT fdo,
							IN PIRP Irp)
{
	PIO_STACK_LOCATION IrpStack = IoGetCurrentIrpStackLocation(Irp);
	NTSTATUS status = STATUS_SUCCESS;
	ULONG BytesTxd = 0;

	ULONG ControlCode = IrpStack->Parameters.DeviceIoControl.IoControlCode;
	ULONG InputLength = IrpStack->Parameters.DeviceIoControl.InputBufferLength;
	ULONG OutputLength = IrpStack->Parameters.DeviceIoControl.OutputBufferLength;

	DbgPrint("DeviceIoControl: Control code %x InputLength %d OutputLength %d",
				ControlCode, InputLength, OutputLength);

	// Get access to the shared buffer
	KIRQL irql;
	KeAcquireSpinLock(&BufferLock,&irql);
	switch( ControlCode)
	{
	///////	Zero Buffer
	case IOCTL_WDM1_ZERO_BUFFER:
		// Zero the buffer
		if( Buffer!=NULL && BufferSize>0)
			RtlZeroMemory(Buffer,BufferSize);
		break;

	///////	Remove Buffer
	case IOCTL_WDM1_REMOVE_BUFFER:
		if( Buffer!=NULL)
		{
			ExFreePool(Buffer);
			Buffer = NULL;
			BufferSize = 0;
		}
		break;

	///////	Get Buffer Size as ULONG
	case IOCTL_WDM1_GET_BUFFER_SIZE:
		if( OutputLength<sizeof(ULONG))
			status = STATUS_INVALID_PARAMETER;
		else
		{
			BytesTxd = sizeof(ULONG);
			RtlCopyMemory(Irp->AssociatedIrp.SystemBuffer,&BufferSize,sizeof(ULONG));
		}
		break;

	///////	Get Buffer
	case IOCTL_WDM1_GET_BUFFER:
		if( OutputLength>BufferSize)
			status = STATUS_INVALID_PARAMETER;
		else
		{
			BytesTxd = OutputLength;
			RtlCopyMemory(Irp->AssociatedIrp.SystemBuffer,Buffer,BytesTxd);
		}
		break;

	///////	Invalid request
	default:
		status = STATUS_INVALID_DEVICE_REQUEST;
	}
	// Release shared buffer
	KeReleaseSpinLock(&BufferLock,irql);

	DbgPrint("DeviceIoControl: %d bytes written",(int)BytesTxd);

	// Complete IRP
	return CompleteIrp(Irp,status,BytesTxd);
}

NTSTATUS Wdm1Read(IN PDEVICE_OBJECT fdo,
				  IN PIRP Irp)
{
	PIO_STACK_LOCATION IrpStack = IoGetCurrentIrpStackLocation(Irp);
	NTSTATUS status = STATUS_SUCCESS;
	LONG BytesTxd = 0;

	// Get call parameters
	LONGLONG FilePointer = IrpStack->Parameters.Read.ByteOffset.QuadPart;
	ULONG ReadLen = IrpStack->Parameters.Read.Length;
	DbgPrint("Read %d bytes from file pointer %d",(int)ReadLen,(int)FilePointer);

	// Get access to the shared buffer
	KIRQL irql;
	KeAcquireSpinLock(&BufferLock,&irql);

	// Check file pointer
	if( FilePointer<0)
		status = STATUS_INVALID_PARAMETER;
	if( FilePointer>=(LONGLONG)BufferSize)
		status = STATUS_END_OF_FILE;

	if( status==STATUS_SUCCESS)
	{
		// Get transfer count
		if( ((ULONG)FilePointer)+ReadLen>BufferSize)
		{
			BytesTxd = BufferSize - (ULONG)FilePointer;
			if( BytesTxd<0) BytesTxd = 0;
		}
		else
			BytesTxd = ReadLen;

		// Read from shared buffer
		if( BytesTxd>0 && Buffer!=NULL)
			RtlCopyMemory( Irp->AssociatedIrp.SystemBuffer, Buffer+FilePointer, BytesTxd);
	}

	// Release shared buffer
	KeReleaseSpinLock(&BufferLock,irql);

	DbgPrint("Read: %d bytes returned",(int)BytesTxd);

	// Complete IRP
	return CompleteIrp(Irp,status,BytesTxd);
}

NTSTATUS Wdm1Write(	IN PDEVICE_OBJECT fdo,
					IN PIRP Irp)
{
	PIO_STACK_LOCATION IrpStack = IoGetCurrentIrpStackLocation(Irp);
	NTSTATUS status = STATUS_SUCCESS;
	LONG BytesTxd = 0;

	// Get call parameters
	LONGLONG FilePointer = IrpStack->Parameters.Write.ByteOffset.QuadPart;
	ULONG WriteLen = IrpStack->Parameters.Write.Length;
	DbgPrint("Write %d bytes from file pointer %d",(int)WriteLen,(int)FilePointer);

	if( FilePointer<0)
		status = STATUS_INVALID_PARAMETER;
	else
	{
		// Get access to the shared buffer
		KIRQL irql;
		KeAcquireSpinLock(&BufferLock,&irql);

		BytesTxd = WriteLen;

		// (Re)allocate buffer if necessary
		if( ((ULONG)FilePointer)+WriteLen>BufferSize)
		{
			ULONG NewBufferSize = ((ULONG)FilePointer)+WriteLen;
			PVOID NewBuffer = ExAllocatePool(NonPagedPool,NewBufferSize);
			if( NewBuffer==NULL)
			{
				BytesTxd = BufferSize - (ULONG)FilePointer;
				if( BytesTxd<0) BytesTxd = 0;
			}
			else
			{
				RtlZeroMemory(NewBuffer,NewBufferSize);
				if( Buffer!=NULL)
				{
					RtlCopyMemory(NewBuffer,Buffer,BufferSize);
					ExFreePool(Buffer);
				}
				Buffer = (PUCHAR)NewBuffer;
				BufferSize = NewBufferSize;
			}
		}

		// Write to shared memory
		if( BytesTxd>0 && Buffer!=NULL)
			RtlCopyMemory( Buffer+FilePointer, Irp->AssociatedIrp.SystemBuffer, BytesTxd);

		// Release shared buffer
		KeReleaseSpinLock(&BufferLock,irql);
	}

	DbgPrint("Write: %d bytes written",(int)BytesTxd);

	// Complete IRP
	return CompleteIrp(Irp,status,BytesTxd);
}




NTSTATUS CompleteIrp( PIRP Irp, NTSTATUS status, ULONG info)
{
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = info;
	IoCompleteRequest(Irp,IO_NO_INCREMENT);
	return status;
}



/***************************************************************
函數名稱：HelloWDMAddDevice()
功能描述：處理“添加設備”消息
***************************************************************/
NTSTATUS HelloWDMAddDevice(IN PDRIVER_OBJECT DriverObject,
                           IN PDEVICE_OBJECT PhysicalDeviceObject)
{
    //定義一個NTSTATUS類型的返回值：
    NTSTATUS status;
    //定義一個功能設備對象（Functional Device Object）：
    PDEVICE_OBJECT fdo;

	UNICODE_STRING ntDeviceName; 
	UNICODE_STRING win32DeviceName; 
	RtlInitUnicodeString(&ntDeviceName, L"\\Device\\Twdm1");
	RtlInitUnicodeString(&win32DeviceName, L"\\DosDevices\\Twdm1");    
    
    DbgPrint("Starting AddDevice()\n");

    //創建我們的功能設備對象，并儲存到fdo中：
    status = IoCreateDevice(
        DriverObject,                //驅動程序對象
        sizeof(DEVICE_EXTENSION),    //要求的設備擴展的大小
        &ntDeviceName,                        //設備名稱，這里為NULL
        FILE_DEVICE_UNKNOWN,        //設備的類型，在標準頭文件WDM.H或NTDDK.H中列出的FILE_DEVICE_xxx值之一
        0,                            //各種常量用OR組合在一起，指示可刪除介質、只讀等。
        FALSE,                        //如果一次只有一個線程可以訪問該設備，為TRUE，否則為FALSE
        &fdo);                        //返回的設備對象

    //NT_SUCCESS宏用于測試IoCreateDevice內核是否成功完成。不要忘記檢查對內核的所有調用是否成功。
    //NT_ERROR宏不等同于!NT_SUCCESS，最好使用!NT_SUCCESS，因為除了錯誤外，它還截獲警告信息。
    if( !NT_SUCCESS(status))
    {
    	DbgPrint("AddDevice(): IoCreateDevice failed!");
        return status;
    }
    else
    {
    	DbgPrint("AddDevice(): IoCreateDevice OK!");
    }

	status = IoCreateSymbolicLink( &win32DeviceName, &ntDeviceName ); 
	if (!NT_SUCCESS(status)) 
	{ 
		DbgPrint("TWDM: IoCreateSymbolicLink() faild ! \n"); 
	} 
	else 
	{ 
		DbgPrint("TWDM: IoCreateSymbolicLink() ok ! \n"); 
	} 

    //創建一個設備擴展對象dx，用于存儲指向fdo的指針：
    PDEVICE_EXTENSION dx = (PDEVICE_EXTENSION)fdo->DeviceExtension;
    dx->fdo = fdo;

    //用IoAttachDeviceToDeviceStack函數把HelloWDM設備掛接到設備棧：
    dx->NextStackDevice = IoAttachDeviceToDeviceStack(fdo, PhysicalDeviceObject);

    //設置fdo的flags。有兩個“位”是必須改變的，一個是必須清除DO_DEVICE_INITIALIZING標志，如果在DriverEntry例程中調用IoCreateDevice()，就不需要清除這個標志位。還有一個是必須設置DO_BUFFER_IO標志位：
    fdo->Flags |= DO_BUFFERED_IO | DO_POWER_PAGABLE;
    fdo->Flags &= ~DO_DEVICE_INITIALIZING;
    
    DbgPrint("AddDevice() OK!\n");

    //返回值：
    return STATUS_SUCCESS;
}


/***************************************************************
函數名稱：HelloWDMPnp()
功能描述：處理“即插即用”消息
***************************************************************/
NTSTATUS HelloWDMPnp(IN PDEVICE_OBJECT fdo,
                        IN PIRP Irp)
{
    //創建一個設備擴展對象dx，用于存儲指向fdo的指針：
    PDEVICE_EXTENSION dx=(PDEVICE_EXTENSION)fdo->DeviceExtension;
	UNICODE_STRING win32DeviceName; 
    //首先要通過函數IoGetCurrentIrpStackLocation()得到當前的IRP，并由此得到Minor Function：
    PIO_STACK_LOCATION IrpStack = IoGetCurrentIrpStackLocation(Irp);
    ULONG MinorFunction = IrpStack->MinorFunction;

    //然后把這個Minor Function傳遞給下一個設備棧：
    IoSkipCurrentIrpStackLocation(Irp);
    NTSTATUS status = IoCallDriver( dx->NextStackDevice, Irp);

    //處理“即插即用”次功能代碼：
    //當Minor Function等于IRP_MN_REMOVE_DEVICE時，說明有設備被拔出或卸下，這時要取消資源分配并刪除設備：
    if( MinorFunction==IRP_MN_REMOVE_DEVICE)
    {
        //取消設備接口：
      RtlInitUnicodeString(&win32DeviceName, L"\\DosDevices\\Twdm1"); 
	  IoDeleteSymbolicLink(&win32DeviceName); 


        //調用IoDetachDevice()把fdo從設備棧中脫開：
        if (dx->NextStackDevice)
            IoDetachDevice(dx->NextStackDevice);
        //刪除fdo：
        IoDeleteDevice(fdo);
    }

    //返回值：
    return status;
}

VOID Wdm1Unload(IN PDRIVER_OBJECT DriverObject)
{
	// Free buffer (do not need to acquire spin lock)
	if( Buffer!=NULL)
		ExFreePool(Buffer);

	DbgPrint("Wdm1Unload");
}