filter.c

USB filter driver sample

#include <WDM.H>
#include <usbiodef.h>
#include <usbdi.h>
#include <usb100.h>
#include <usbdlib.h>

#include "filter.h"

#ifdef ALLOC_PRAGMA
        #pragma alloc_text(INIT, DriverEntry)
        #pragma alloc_text(PAGE, USB_AddDevice)
        #pragma alloc_text(PAGE, USB_DriverUnload)
#endif

static char* irpname[] = {
		"IRP_MJ_CREATE",
		"IRP_MJ_CREATE_NAMED_PIPE",
		"IRP_MJ_CLOSE",
		"IRP_MJ_READ",
		"IRP_MJ_WRITE",
		"IRP_MJ_QUERY_INFORMATION",
		"IRP_MJ_SET_INFORMATION",
		"IRP_MJ_QUERY_EA",
		"IRP_MJ_SET_EA",
		"IRP_MJ_FLUSH_BUFFERS",
		"IRP_MJ_QUERY_VOLUME_INFORMATION",
		"IRP_MJ_SET_VOLUME_INFORMATION",
		"IRP_MJ_DIRECTORY_CONTROL",
		"IRP_MJ_FILE_SYSTEM_CONTROL",
		"IRP_MJ_DEVICE_CONTROL",
		"IRP_MJ_INTERNAL_DEVICE_CONTROL",
		"IRP_MJ_SHUTDOWN",
		"IRP_MJ_LOCK_CONTROL",
		"IRP_MJ_CLEANUP",
		"IRP_MJ_CREATE_MAILSLOT",
		"IRP_MJ_QUERY_SECURITY",
		"IRP_MJ_SET_SECURITY",
		"IRP_MJ_POWER",
		"IRP_MJ_SYSTEM_CONTROL",
		"IRP_MJ_DEVICE_CHANGE",
		"IRP_MJ_QUERY_QUOTA",
		"IRP_MJ_SET_QUOTA",
		"IRP_MJ_PNP",
		};

static char* urb_name[] = {
		"URB_FUNCTION_SELECT_CONFIGURATION",
		"URB_FUNCTION_SELECT_INTERFACE",
		"URB_FUNCTION_ABORT_PIPE",
		"URB_FUNCTION_TAKE_FRAME_LENGTH_CONTROL",
		"URB_FUNCTION_RELEASE_FRAME_LENGTH_CONTROL",
		"URB_FUNCTION_GET_FRAME_LENGTH",
		"URB_FUNCTION_SET_FRAME_LENGTH",
		"URB_FUNCTION_GET_CURRENT_FRAME_NUMBER",
		"URB_FUNCTION_CONTROL_TRANSFER",
		"URB_FUNCTION_BULK_OR_INTERRUPT_TRANSFER",
		"URB_FUNCTION_ISOCH_TRANSFER",
		"URB_FUNCTION_GET_DESCRIPTOR_FROM_DEVICE",
		"URB_FUNCTION_SET_DESCRIPTOR_TO_DEVICE",
		"URB_FUNCTION_SET_FEATURE_TO_DEVICE",
		"URB_FUNCTION_SET_FEATURE_TO_INTERFACE",
		"URB_FUNCTION_SET_FEATURE_TO_ENDPOINT",
		"URB_FUNCTION_CLEAR_FEATURE_TO_DEVICE",
		"URB_FUNCTION_CLEAR_FEATURE_TO_INTERFACE",
		"URB_FUNCTION_CLEAR_FEATURE_TO_ENDPOINT",
		"URB_FUNCTION_GET_STATUS_FROM_DEVICE",
		"URB_FUNCTION_GET_STATUS_FROM_INTERFACE",
		"URB_FUNCTION_GET_STATUS_FROM_ENDPOINT",
		"URB_FUNCTION_RESERVED0",
		"URB_FUNCTION_VENDOR_DEVICE",
		"URB_FUNCTION_VENDOR_INTERFACE",
		"URB_FUNCTION_VENDOR_ENDPOINT",
		"URB_FUNCTION_CLASS_DEVICE",
		"URB_FUNCTION_CLASS_INTERFACE",
		"URB_FUNCTION_CLASS_ENDPOINT",
		"URB_FUNCTION_RESERVED",
		"URB_FUNCTION_RESET_PIPE",
		"URB_FUNCTION_CLASS_OTHER",
		"URB_FUNCTION_VENDOR_OTHER",
		"URB_FUNCTION_GET_STATUS_FROM_OTHER",
		"URB_FUNCTION_CLEAR_FEATURE_TO_OTHER",
		"URB_FUNCTION_SET_FEATURE_TO_OTHER",
		"URB_FUNCTION_GET_DESCRIPTOR_FROM_ENDPOINT",
		"URB_FUNCTION_SET_DESCRIPTOR_TO_ENDPOINT",
		"URB_FUNCTION_GET_CONFIGURATION",
		"URB_FUNCTION_GET_INTERFACE",
		"URB_FUNCTION_GET_DESCRIPTOR_FROM_INTERFACE",
		"URB_FUNCTION_SET_DESCRIPTOR_TO_INTERFACE",
	};

NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath)
	{
	int i;

    PAGED_CODE();

    UNREFERENCED_PARAMETER(RegistryPath);

    DBGOUT(("DriverEntry")); 

    for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++)
		{
        DriverObject->MajorFunction[i] = USB_Dispatch; 
		}

    DriverObject->DriverExtension->AddDevice = USB_AddDevice;
    DriverObject->DriverUnload = USB_DriverUnload;

    return STATUS_SUCCESS;
	}

NTSTATUS USB_AddDevice(IN PDRIVER_OBJECT driverObj, IN PDEVICE_OBJECT physicalDevObj)
	{
    NTSTATUS status;
    PDEVICE_OBJECT filterDevObj = NULL;
	PDEVICE_OBJECT controlDevObj = NULL;
    
    PAGED_CODE();

    DBGOUT(("USB_AddDevice: drvObj=%ph, pdo=%ph", driverObj, physicalDevObj)); 

	/// for filter device...
    status = IoCreateDevice(driverObj, sizeof(struct DEVICE_EXTENSION), NULL, FILE_DEVICE_UNKNOWN, 
                                FILE_AUTOGENERATED_DEVICE_NAME, FALSE, &filterDevObj);
    
    if (NT_SUCCESS(status))
		{
        struct DEVICE_EXTENSION *devExt;	/// for filter
		struct DEVICE_EXTENSION *devExt2;	/// for control
		UNICODE_STRING DevName;

		RtlInitUnicodeString(&DevName, L"\\Device\\USBfilterDrv");
		RtlInitUnicodeString(&DevlinkName, BEACON_USB_FILTER);

		/// for control device...
		status = IoCreateDevice(driverObj, sizeof(struct DEVICE_EXTENSION), &DevName, FILE_DEVICE_UNKNOWN,
									0, FALSE, &controlDevObj);

		if (NT_SUCCESS(status))
			{
			controlDevObj->Flags |= DO_DEVICE_INITIALIZING;
			devExt2 = (struct DEVICE_EXTENSION *)controlDevObj->DeviceExtension;
			RtlZeroMemory(devExt2, sizeof(struct DEVICE_EXTENSION));
			devExt2->signature = DEVICE_EXTENSION_SIGNATURE;
			devExt2->state = STATE_INITIALIZED;
			devExt2->filterDevObj = filterDevObj;
			devExt2->physicalDevObj = physicalDevObj;
			devExt2->controlDevObj = controlDevObj;
			devExt2->WhichObj = CONTROL_DEVICE_OBJ;
			
			devExt2->pendingActionCount = 0;

			controlDevObj->Flags |= DO_BUFFERED_IO | DO_DIRECT_IO;
			controlDevObj->Flags &= ~DO_DEVICE_INITIALIZING;

			DBGOUT(("Created controlDevObj : %ph", controlDevObj));
			
			status = IoCreateSymbolicLink(&DevlinkName, &DevName);
			if (NT_SUCCESS(status))
				{
				DBGOUT(("Created SymbolicLink"));
				}
			else
				{
				DBGOUT(("CreateSymbolicLink failed!!! : %04x", status));
				IoDeleteDevice(filterDevObj);
				IoDeleteDevice(controlDevObj);

				return status;
				}
			}
		else
			{
			DBGOUT(("Control Device Object Creation failed!!! : %04x", status));
			IoDeleteDevice(filterDevObj);

			return status;
			}

        ASSERT(filterDevObj);

		devExt = (struct DEVICE_EXTENSION *)filterDevObj->DeviceExtension;
		RtlZeroMemory(devExt, sizeof(struct DEVICE_EXTENSION));
        devExt->signature = DEVICE_EXTENSION_SIGNATURE;
        devExt->state = STATE_INITIALIZED;
        devExt->filterDevObj = filterDevObj;
        devExt->physicalDevObj = physicalDevObj;
		devExt->controlDevObj = controlDevObj;
		devExt->WhichObj = FILTER_DEVICE_OBJ;
		devExt->readflag = CAPTURE_STOP;
		devExt->writeflag = CAPTURE_STOP;
		devExt->bThreadStop = true;
		devExt->bInit = false;
		devExt->packet_Event = devExt->removeEvent_app = NULL;

		/// for Data capture...
		KeInitializeSemaphore(&devExt->pSemaphore, 0, MAX_PACKET);
		ExInitializeFastMutex(&devExt->FMutex);
		/// create thread...
	/*	status = PsCreateSystemThread(&hThread, 0, NULL, NULL, NULL, (PKSTART_ROUTINE)CaptureThread, devExt);
		if (NT_SUCCESS(status))
			{
			ObReferenceObjectByHandle(hThread, THREAD_ALL_ACCESS, NULL, KernelMode, (PVOID *)&devExt->pThread, NULL);
			ZwClose(hThread);
			}
		else
			{
			DBGOUT(("CreateThread failed!!!"));
			IoDeleteSymbolicLink(&DevlinkName);
			IoDeleteDevice(filterDevObj);
			IoDeleteDevice(controlDevObj);

			return status;
			}
		//*/
        
        devExt->pendingActionCount = 0;
        KeInitializeEvent(&devExt->removeEvent, NotificationEvent, FALSE);

		/// for packet delay
		KeInitializeEvent(&devExt->kTimerEvent, SynchronizationEvent, FALSE);
		KeInitializeTimer(&devExt->kTimer);
		KeInitializeDpc(&devExt->kTimerDpc, TimerDpc, devExt);
		devExt->lnTimer = RtlConvertLongToLargeInteger(-1*10000000);

        devExt->topDevObj = IoAttachDeviceToDeviceStack(filterDevObj, physicalDevObj);

        ASSERT(devExt->topDevObj);
		
		devExt2->topDevObj = devExt->topDevObj;
        DBGOUT(("Created filterDevObj %ph attached to %ph", filterDevObj, devExt->topDevObj));

		TIME_SYNCHRONIZE(&devExt->start_time);

        filterDevObj->Flags |= (devExt->topDevObj->Flags & (DO_BUFFERED_IO | DO_DIRECT_IO));
        filterDevObj->Flags |= (devExt->topDevObj->Flags & (DO_POWER_INRUSH | DO_POWER_PAGABLE /*| DO_POWER_NOOP*/)); 
        filterDevObj->Flags &= ~DO_DEVICE_INITIALIZING;
	
        RegistryAccessSample(devExt, devExt->physicalDevObj);
		}
		
    ASSERT(NT_SUCCESS(status));
    return status;
	}

VOID USB_DriverUnload(IN PDRIVER_OBJECT DriverObject)
	{
	PAGED_CODE();

	DBGOUT(("USB_DriverUnload"));
	}

NTSTATUS USB_Dispatch(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
	{
    struct DEVICE_EXTENSION *devExt;
    PIO_STACK_LOCATION irpSp;
    BOOLEAN passIrpDown = TRUE;
    UCHAR majorFunc, minorFunc;
    NTSTATUS status = STATUS_SUCCESS;

    devExt = DeviceObject->DeviceExtension;
    ASSERT(devExt->signature == DEVICE_EXTENSION_SIGNATURE);

    irpSp = IoGetCurrentIrpStackLocation(Irp);

    majorFunc = irpSp->MajorFunction;
    minorFunc = irpSp->MinorFunction;

	if (devExt->WhichObj == FILTER_DEVICE_OBJ)
		{
		/*
		 *  For all IRPs except REMOVE, we increment the PendingActionCount
		 *  across the dispatch routine in order to prevent a race condition with
		 *  the REMOVE_DEVICE IRP (without this increment, if REMOVE_DEVICE
		 *  preempted another IRP, device object and extension might get
		 *  freed while the second thread was still using it).
		 */
		if (!((majorFunc == IRP_MJ_PNP) && (minorFunc == IRP_MN_REMOVE_DEVICE)))
			{
			IncrementPendingActionCount(devExt);
			}

		if ((majorFunc != IRP_MJ_PNP) && (majorFunc != IRP_MJ_CLOSE) && 
			((devExt->state == STATE_REMOVING) || (devExt->state == STATE_REMOVED)))
			{
			/*
			 *  While the device is being removed, 
			 *  we only pass down the PNP and CLOSE IRPs.
			 *  We fail all other IRPs.
			 */
			status = Irp->IoStatus.Status = STATUS_DELETE_PENDING;
			IoCompleteRequest(Irp, IO_NO_INCREMENT);
			passIrpDown = FALSE;
			}
		else 
			{
			switch (majorFunc)
				{
				case IRP_MJ_PNP:
					status = USB_PnP(devExt, Irp);
					passIrpDown = FALSE;
					break;

				case IRP_MJ_POWER:
					status = USB_Power(devExt, Irp);
					passIrpDown = FALSE;
					break;

				case IRP_MJ_INTERNAL_DEVICE_CONTROL:
				//	status = USB_IoInternal(devExt, Irp);
				//	passIrpDown = FALSE;
					break;

				case IRP_MJ_READ:
				case IRP_MJ_WRITE:
					status = USB_Read_Write(devExt, Irp);
					passIrpDown = FALSE;
					break;

				case IRP_MJ_DEVICE_CONTROL:
				case IRP_MJ_CREATE:
				case IRP_MJ_CLOSE:
				case IRP_MJ_SYSTEM_CONTROL:
				default:
					break;
				}
			}

		if (passIrpDown)
			{
			IoCopyCurrentIrpStackLocationToNext(Irp);
			status = IoCallDriver(devExt->topDevObj, Irp);
			}

		if (!((majorFunc == IRP_MJ_PNP) && (minorFunc == IRP_MN_REMOVE_DEVICE)))
			{
			DecrementPendingActionCount(devExt);
			}

		}
	else	// Control Device Object...
		{
		Irp->IoStatus.Status = STATUS_SUCCESS;
		Irp->IoStatus.Information = 0;
		switch(majorFunc)
			{
			case IRP_MJ_CREATE:
			case IRP_MJ_CLOSE:
				break;

			case IRP_MJ_DEVICE_CONTROL:
				status = USB_IoCtrl(devExt, Irp);
				break;

			default:
				break;
			}
		
		IoCompleteRequest(Irp, IO_NO_INCREMENT);
		}
	
    return status;
	}

NTSTATUS USB_Read_Write(struct DEVICE_EXTENSION *devExt, PIRP irp)
	{
	NTSTATUS status = STATUS_SUCCESS;
	PIO_STACK_LOCATION irpSp = IoGetCurrentIrpStackLocation(irp);
	DWORD wLen;
	bool bWhich = devExt->bWhichBuffer_write;
	KIRQL oldIrql;

	if (IRP_MJ_READ == irpSp->MajorFunction)
		{
		IoCopyCurrentIrpStackLocationToNext(irp);
		IoSetCompletionRoutine(irp, USB_IoReadCompletion, devExt, TRUE, TRUE, TRUE);
		status = IoCallDriver(devExt->topDevObj, irp);
		}
	else
		{
		/// for packet delay
	//	KeSetTimer(&devExt->kTimer, devExt->lnTimer, &devExt->kTimerDpc);
	//	KeWaitForSingleObject(&devExt->kTimerEvent, Executive, KernelMode, FALSE, &devExt->lnTimer);

		wLen = irpSp->Parameters.Write.Length;
		if (wLen > 0 && devExt->writeflag == CAPTURE_START)
			{
			/// write capture
		//	KeRaiseIrql(DISPATCH_LEVEL, &oldIrql);
			oldIrql = KeGetCurrentIrql();
			if (oldIrql < DISPATCH_LEVEL)
				{
				ExAcquireFastMutex(&devExt->FMutex);
			//	KeWaitForSingleObject(&devExt->kTimerEvent, Executive, KernelMode, FALSE, &devExt->lnTimer);
				}

			if (wLen + devExt->wWriteOffset[bWhich] > MAX_BUFFER)
				{
				bWhich = bWhich ^ 0x01;
				devExt->wWriteOffset[bWhich] = 0;
				devExt->bWhichBuffer_write = bWhich;
				}

		//	DBGOUT(("write len : %u(%u)", wLen, bWhich));
			devExt->PkHdr[devExt->nCount].code = wWriteFlag;
			devExt->PkHdr[devExt->nCount].length = (WORD)wLen;
			devExt->PkHdr[devExt->nCount].offset = (WORD)devExt->wWriteOffset[bWhich];
			devExt->PkHdr[devExt->nCount].bFlag = bWhich;
			GET_TIME(&devExt->PkHdr[devExt->nCount].tv, &devExt->start_time);
			if (devExt->nCount >= MAX_PACKET-1)
				devExt->nCount = 0;
			else
				devExt->nCount++;

			RtlCopyMemory(devExt->pWriteBuffer[bWhich] + devExt->wWriteOffset[bWhich], irp->AssociatedIrp.SystemBuffer, wLen);
			InterlockedExchangeAdd((PLONG)&devExt->wWriteOffset[bWhich], wLen);
		//	devExt->wWriteOffset[bWhich] += wLen;
			if (oldIrql < DISPATCH_LEVEL)
				ExReleaseFastMutex(&devExt->FMutex);

		//	KeLowerIrql(oldIrql);
			KeReleaseSemaphore(&devExt->pSemaphore, 0, 1, FALSE);
			}
		IoCopyCurrentIrpStackLocationToNext(irp);
        status = IoCallDriver(devExt->topDevObj, irp);
		}

	return status;
	}

NTSTATUS USB_IoReadCompletion(IN PDEVICE_OBJECT DeviceObject, IN PIRP irp, IN PVOID Context)
	{
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(irp);
	struct DEVICE_EXTENSION *devExt = (struct DEVICE_EXTENSION *)Context;
	DWORD wLen;
	bool bWhich = devExt->bWhichBuffer_read;
	KIRQL oldIrql;
	
	wLen = irp->IoStatus.Information;
	if (wLen > 0 && devExt->readflag == CAPTURE_START)
		{
		/// read capture
	//	KeRaiseIrql(DISPATCH_LEVEL, &oldIrql);