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;
}