zz: 自定义StartIo

ZZ from : http://www.joenchen.com/archives/424

这个系统给的StartIo虽然好用, 但是很多时候我们还是需要使用自定义的StartIo例程的. 因为系统提供的只能使用一个队列.如果是我们自己建立的话, 可以建立多个队列. 灵活性高, 但是稍微复杂些. 其实要说复杂也不尽然, 用那个系统的时候感觉不是很好,因为封装的时候不给力, 还需要在我们自己的函数中帮系统做一些事情.我倒. 还是这个好, 想建立几个队列就几个队列. 而且代码也不是那么的复杂. 也就是几个函数的调用了!  

 

自定义的StartIo例程的话需要程序员自己维护Irp队列, 这样程序员的话可以维护多个队列, 分别对应不同的Irp请求. 首先要搞这个自定义StartIo的话, 要搞个结构KDEVICE_QUEUE, 放在设备扩展里面就可以了. 如果想维护多个, 那么多放几个结构就可以了. 在DriverEntry里面调用keInitializeDeviceQueue函数初始化队列.当然有队列的话, 需要插入, 删除. 插入是KeInsertDeviceQueue, 还有个KeRemoveDeviceQueue删除队列中的元素,.

 

当然在编写分发函数的时候, 在派遣例程中要调用IoMarkIrpPending挂起当前的Irp.然后准备将Irp进入队列, 当然在进入队列之前需要将当前IRQL提升至DISPATCH_LEVEL.这个需要记得, 不要忘记, 然后在调用KeInsertDeviceQueue函数的时候如果返回FALSE的话, 表示队列已经满了, 需要开始调用自定义的StartIo例程了. 下面的代码就是这样做的.

那么在自定义的StartIo例程中, 处理传进StartIo中的Irp, 然后枚举队列中的Irp.依次出队列, 处理!那么自定义的StartIo就那么多东西了, 看代码吧, 这边是用户态的代码:

/*	Windows 内核下StartIo例程试验 3环代码
	编译方法参见makefile. TAB = 8
*/
#include <stdio.h>
#include <windows.h>

#pragma comment( linker, "/Entry:Jmain" )
#pragma comment( linker, "/subsystem:console" )

#define DEVICE_NAME	"\\\\.\\SysLinkCustomStartIo"
//===========================================================================
//线程过程, 向设备发送15次, 写入请求
//===========================================================================
DWORD __stdcall ThreadProc( PVOID pContext ) {
	UCHAR ucBuf[10];
	ULONG i, j;
	DWORD dwByteWrite;
	BOOL bRet;
	OVERLAPPED StOverlapped[15] = {0};
	HANDLE hEvent[15] = {0};

	__try {
		for( i = 0; i < sizeof( StOverlapped ) / sizeof( StOverlapped[0] ); i++ ) {

			StOverlapped[i].hEvent = CreateEvent( NULL, FALSE, TRUE, NULL );
			if ( !StOverlapped[i].hEvent ) {
				printf( "创建同步事件失败!\n" );
				return -1;
			}
		}

		for( i = 0; i < sizeof( StOverlapped ) / sizeof( StOverlapped[0] ); i++ ) {
			hEvent[i] = StOverlapped[i].hEvent;
		}

		Sleep( 1000 );
		for( i = 0; i < sizeof( StOverlapped ) / sizeof( StOverlapped[0] ); i++ ) {

			RtlFillMemory( ucBuf, sizeof( ucBuf ), i + 'a' );

			bRet = WriteFile( *( PHANDLE )pContext, ucBuf, sizeof( ucBuf ),
					  &dwByteWrite, &StOverlapped[i]  );

			if ( !bRet && GetLastError() != ERROR_IO_PENDING ) {
				printf( "写入设备失败!\n" );
				return -1;
			} else {
				for( j = 0; j < sizeof( ucBuf ); j++ ) {
					printf( "%c\t", ucBuf[j] );
				}

				printf( "\n" );
			}
		}

		//这个复杂度加了一些中间有可能还会有取消IRP的出现
		//CancelIo(*(PHANDLE)pContext);
		WaitForMultipleObjects( sizeof( StOverlapped ) / sizeof( StOverlapped[0] ),
					&hEvent[0], TRUE, INFINITE );

		printf( "设备处理完毕!\n" );

	} __finally {
		for( i = 0; i < sizeof( hEvent ) / sizeof( hEvent[0] ); i++ ) {
			if ( StOverlapped[i].hEvent ) {
				CloseHandle( StOverlapped[i].hEvent );
			}
		}
	}

	return 0;
}
//===========================================================================
int Jmain( ) {
	HANDLE hDevice = NULL;
	HANDLE hThead[2] = {0};
	DWORD dwTemp;

	do {
		hDevice = CreateFile( DEVICE_NAME, GENERIC_READ | GENERIC_WRITE, 0, NULL,
			OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED, NULL );

		if ( hDevice == INVALID_HANDLE_VALUE ) {
			printf( "打开设备失败\n" );
			break;
		}

		hThead[0] = CreateThread( NULL, 0, &ThreadProc, &hDevice, 0, &dwTemp );

		if ( !hThead[0] ) {
			printf( "创建线程1失败!\n" );
			break;
		} else {
			printf( "创建线程1成功, 线程1已经开始运行!\n" );
		}

		hThead[1] = CreateThread( NULL, 0, &ThreadProc, &hDevice, 0, &dwTemp );

		if ( !hThead[1] ) {
			printf( "创建线程2失败!\n" );
			break;
		} else {
			printf( "创建线程2成功, 线程2已经开始运行!\n" );
		}

		printf( "主线程开始等待两个线程返回!\n" );

		WaitForMultipleObjects( 2, hThead, TRUE, INFINITE );
		printf( "两个线程都已经返回!\n" );

	} while ( FALSE );

//---------------------------------------------------------------------------
	if ( hDevice ) {
		CloseHandle( hDevice );
	}
	if ( hThead[0] ) {
		CloseHandle( hThead[0] );
	}
	if ( hThead[1] ) {
		CloseHandle( hThead[1] );
	}

	system( "pause" );

	return 0;
}

Driver code如下：

/*      Windows 内核下自定义StartIo 3环代码
        编译方法参见makefile. TAB = 8
*/
#include <ntddk.h>

#define DEVICE_NAME	L"\\Device\\CustomStartIo"
#define SYSLINK_NAME	L"\\??\\SysLinkCustomStartIo"

typedef struct tagDevice_Ext {
	KDEVICE_QUEUE DeviceQueue;		//设备队列
	PDEVICE_OBJECT pDeviceObj;
	UNICODE_STRING USzDeviceName;
	UNICODE_STRING USzSysLinkName;
} DEVICE_EXT, *PDEVICE_EXT;
//===========================================================================
//驱动卸载例程
//===========================================================================
#pragma code_seg( "PAGE" )
VOID DriverUnload( PDRIVER_OBJECT pDriverObj ) {
	PDEVICE_EXT pDeviceObj;
	PDEVICE_OBJECT pNextDeviceObj = NULL;

	pNextDeviceObj = pDriverObj->DeviceObject;

	while ( pNextDeviceObj != NULL ) {
		pDeviceObj = pNextDeviceObj->DeviceExtension;

		IoDeleteDevice( pDeviceObj->pDeviceObj );
		IoDeleteSymbolicLink( &pDeviceObj->USzSysLinkName );

		KdPrint( ( "删除%wZ设备成功!\n", &pDeviceObj->USzDeviceName ) );

		pNextDeviceObj = pNextDeviceObj->NextDevice;
	}
}
//===========================================================================
//所有不关心的Irp处理例程
//===========================================================================
#pragma code_seg( "PAGE" )
NTSTATUS DispatchRoutine( PDEVICE_OBJECT pDeviceObj, PIRP pIrp ) {

	pIrp->IoStatus.Information = 0;
	pIrp->IoStatus.Status = STATUS_SUCCESS;

	IoCompleteRequest( pIrp, IO_NO_INCREMENT );
	return STATUS_SUCCESS;
}
//===========================================================================
//自定义的StartIo例程
//===========================================================================
#pragma code_seg( )
VOID CustomStartIo( PDEVICE_OBJECT pDeviceObj, PIRP pIrp ) {
	PIRP NewIrp = NULL;
	KEVENT Event;
	ULONG ulWriteLen, i;
	LARGE_INTEGER liTimeOut;
	PDEVICE_EXT pDeviceExt = NULL;
	PIO_STACK_LOCATION Stack = NULL;
	PKDEVICE_QUEUE_ENTRY pQueue = NULL;

	PAGED_CODE_LOCKED();

	pDeviceExt = ( PDEVICE_EXT )pDeviceObj->DeviceExtension;
	NewIrp = pIrp;

	do {

//---------------------------------------------------------------------------
		//打印用户层传递的数据
//---------------------------------------------------------------------------
		KeInitializeEvent( &Event, NotificationEvent, FALSE );

		//等3秒
		liTimeOut.QuadPart = -3 * 1000 * 1000 * 10;

		//定义一个3秒的延时，主要是为了模拟该IRP操作需要大概3秒左右时间
		KeWaitForSingleObject( &Event, Executive, KernelMode, FALSE, &liTimeOut );

		Stack = IoGetCurrentIrpStackLocation( NewIrp );

		ulWriteLen = Stack->Parameters.Write.Length;

		for( i = 0; i < ulWriteLen; i++ ) {
			KdPrint( ( "%c\t",
				   *( ( UCHAR* )( NewIrp->AssociatedIrp.SystemBuffer ) ) ) );
		}
		KdPrint( ( "\n" ) );
//---------------------------------------------------------------------------
		NewIrp->IoStatus.Status = STATUS_SUCCESS;
		NewIrp->IoStatus.Information = 0;
		IoCompleteRequest( NewIrp, IO_NO_INCREMENT );
		KdPrint( ( "CustomStartIo调用一次:%x\n", pQueue ) );

		//将存放的Irp出队列
		pQueue = KeRemoveDeviceQueue( &pDeviceExt->DeviceQueue );

		//一直处理到队列为空为止
		if ( pQueue == NULL ) {
			break;
		}

		NewIrp = CONTAINING_RECORD( pQueue, IRP,
					    Tail.Overlay.DeviceQueueEntry );
	} while( TRUE );
}
//===========================================================================
//取消例程
//===========================================================================
VOID OnCancelIRP( PDEVICE_OBJECT DeviceObject, PIRP Irp  ) {

	//释放Cancel自旋锁
	IoReleaseCancelSpinLock( Irp->CancelIrql );

	//设置完成状态为STATUS_CANCELLED
	Irp->IoStatus.Status = STATUS_CANCELLED;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest( Irp, IO_NO_INCREMENT );

	KdPrint( ( "离开取消例程\n" ) );
}
//===========================================================================
//写入请求处理
//===========================================================================
#pragma code_seg( "PAGE" )
NTSTATUS DispatchWrite( PDEVICE_OBJECT pDeviceObj, PIRP pIrp ) {
	KIRQL OldIrql;
	BOOLEAN bRet;
	PDEVICE_EXT pDeviceExt = NULL;

	PAGED_CODE();

	pDeviceExt = ( PDEVICE_EXT )pDeviceObj->DeviceExtension;

	//将IRP设置为挂起
	IoMarkIrpPending( pIrp );

	//设置取消例程
	IoSetCancelRoutine( pIrp, OnCancelIRP );

	//提升IRP至DISPATCH_LEVEL
	KeRaiseIrql( DISPATCH_LEVEL, &OldIrql );

	KdPrint( ( "DispatchWrite设备扩展指针:%x\n",
		   &pIrp->Tail.Overlay.DeviceQueueEntry ) );

	//如果KeInsertDeviceQueue返回FALSE的时候,表明Irp没有插入到
	//队列, 而是需要立即被结束
	bRet = KeInsertDeviceQueue( &pDeviceExt->DeviceQueue,
				    &pIrp->Tail.Overlay.DeviceQueueEntry );
	if ( !bRet ) {
		KdPrint( ( "调用CustomStartIo例程!\n" ) );

		//调用自定义StartIo例程.
		CustomStartIo( pDeviceObj, pIrp );
	}

	//将IRP降至原来IRQL
	KeLowerIrql( OldIrql );

	//返回pending状态
	return STATUS_PENDING;
}
//===========================================================================
//驱动入口
//===========================================================================
#pragma code_seg( "INIT" )
NTSTATUS DriverEntry( PDRIVER_OBJECT pDriverObj, PUNICODE_STRING pRegPath ) {
	ULONG i;
	NTSTATUS Status;
	PDEVICE_EXT pDeviceExt = NULL;
	PDEVICE_OBJECT pDeviceObj = NULL;
	UNICODE_STRING USzDeviceName = RTL_CONSTANT_STRING( DEVICE_NAME );
	UNICODE_STRING USzSysLinkName = RTL_CONSTANT_STRING( SYSLINK_NAME );

	Status = IoCreateDevice( pDriverObj, sizeof( DEVICE_EXT ), &USzDeviceName,
				 FILE_DEVICE_UNKNOWN, 0, TRUE, &pDeviceObj );
	if ( !NT_SUCCESS( Status ) ) {
		KdPrint( ( "创建设备失败!\n" ) );
		return Status;
	}

	Status = IoCreateSymbolicLink( &USzSysLinkName, &USzDeviceName );
	if ( !NT_SUCCESS( Status ) ) {
		KdPrint( ( "创建符号链接失败!\n" ) );
		IoDetachDevice( pDeviceObj );
		return Status;
	}

	//设置设备属性, 设备扩展
	pDeviceObj->Flags |= DO_BUFFERED_IO;
	pDeviceExt = pDeviceObj->DeviceExtension;
	pDeviceExt->pDeviceObj = pDeviceObj;
	pDeviceExt->USzDeviceName = USzDeviceName;
	pDeviceExt->USzSysLinkName = USzSysLinkName;

	//初始化设备队列(自己实现StartIo)
	RtlZeroMemory( &pDeviceExt->DeviceQueue, sizeof( pDeviceExt->DeviceQueue ) );
	KeInitializeDeviceQueue( &pDeviceExt->DeviceQueue );

	for( i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; i++ ) {
		pDriverObj->MajorFunction[i] = &DispatchRoutine;
	}
	pDriverObj->MajorFunction[IRP_MJ_WRITE] = &DispatchWrite;
	pDriverObj->DriverUnload = &DriverUnload;

	return Status;
}