windows系统下usb过滤驱动实现

#include "stddcls.h" #include "driver.h" #include <ntddk.h> #include <srb.h> #include <scsi.h> #include <usbdi.h> #include <usbdlib.h> NTSTATUS AddDevice(IN PDRIVER_OBJECT DriverObject, IN PDEVICE_OBJECT pdo); VOID DriverUnload(IN PDRIVER_OBJECT fido); NTSTATUS DispatchAny(IN PDEVICE_OBJECT fido, IN PIRP Irp); NTSTATUS DispatchPower(IN PDEVICE_OBJECT fido, IN PIRP Irp); NTSTATUS DispatchPnp(IN PDEVICE_OBJECT fido, IN PIRP Irp); NTSTATUS DispatchWmi(IN PDEVICE_OBJECT fido, IN PIRP Irp); ULONG GetDeviceTypeToUse(PDEVICE_OBJECT pdo); NTSTATUS PnpStartDevice(IN PDEVICE_OBJECT fido, IN PIRP Irp); NTSTATUS ReadDeviceDescriptor(IN PDEVICE_OBJECT fido); NTSTATUS CallRemoveUSBDevice(IN PDEVICE_OBJECT fido); NTSTATUS StartDeviceCompletionRoutine(PDEVICE_OBJECT fido, PIRP Irp, PDEVICE_EXTENSION pdx); NTSTATUS StartPnPQueryDeviceCompletionRoutine(PDEVICE_OBJECT fido, PIRP Irp, PDEVICE_EXTENSION pdx); NTSTATUS QueryDeviceRelationCompletionRoutine(PDEVICE_OBJECT fido, PIRP Irp, PDEVICE_EXTENSION pdx); NTSTATUS UsageNotificationCompletionRoutine(PDEVICE_OBJECT fido, PIRP Irp, PDEVICE_EXTENSION pdx); LONG BulkUsb_IoDecrement(IN OUT PDEVICE_EXTENSION DeviceExtension); LONG BulkUsb_IoIncrement(IN OUT PDEVICE_EXTENSION DeviceExtension); int QueryRegDWORD(UNICODE_STRING *pValueName, ULONG *pReturn); int SetRegDWORD(UNICODE_STRING *pValueName, ULONG nSetValue); bool USBDeviceIsRegist(ULONG vid_pid); /////////////////////////////////////////////////////////////////////////////// #pragma INITCODE extern "C" NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath) { // DriverEntry KdPrint((DRIVERNAME " - Entering DriverEntry: DriverObject %8.8lX\n", DriverObject)); // Initialize function pointers UNICODE_STRING ValueName; RtlInitUnicodeString( &ValueName, L"FilterIsRun"); SetRegDWORD(&ValueName, 0); DriverObject->DriverUnload = DriverUnload; DriverObject->DriverExtension->AddDevice = AddDevice; for (int i = 0; i < arraysize(DriverObject->MajorFunction); ++i) DriverObject->MajorFunction[i] = DispatchAny; DriverObject->MajorFunction[IRP_MJ_POWER] = DispatchPower; DriverObject->MajorFunction[IRP_MJ_PNP] = DispatchPnp; DriverObject->MajorFunction[IRP_MJ_SCSI] = DispatchForSCSI; return STATUS_SUCCESS; } // DriverEntry /////////////////////////////////////////////////////////////////////////////// #pragma PAGEDCODE VOID DriverUnload(IN PDRIVER_OBJECT DriverObject) { // DriverUnload UNICODE_STRING ValueName; PAGED_CODE(); RtlInitUnicodeString( &ValueName, L"FilterIsRun"); SetRegDWORD(&ValueName, 0); KdPrint((DRIVERNAME " - Entering DriverUnload: DriverObject %8.8lX\n", DriverObject)); } // DriverUnload /////////////////////////////////////////////////////////////////////////////// NTSTATUS AddDevice(IN PDRIVER_OBJECT DriverObject, IN PDEVICE_OBJECT pdo) { // AddDevice PAGED_CODE(); NTSTATUS status; UNICODE_STRING ValueName; ULONG dwReturn; PDEVICE_OBJECT fido; status = IoCreateDevice(DriverObject, sizeof(DEVICE_EXTENSION), NULL, FILE_DEVICE_UNKNOWN, //GetDeviceTypeToUse(pdo), 0, FALSE, &fido); if (!NT_SUCCESS(status)) { // can't create device object KdPrint((DRIVERNAME " - IoCreateDevice failed - %X\n", status)); return status; } // can't create device object PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fido->DeviceExtension; do { // finish initialization IoInitializeRemoveLock(&pdx->RemoveLock, 0, 0, 0); pdx->DeviceObject = fido; pdx->Pdo = pdo; //将过滤驱动附加在底层驱动之上 PDEVICE_OBJECT fdo = IoAttachDeviceToDeviceStack(fido, pdo); if (!fdo) { // can't attach KdPrint((DRIVERNAME " - IoAttachDeviceToDeviceStack failed\n")); status = STATUS_DEVICE_REMOVED; break; } // can't attach //记录底层驱动 pdx->LowerDeviceObject = fdo; //由于不知道底层驱动是直接IO还是BufferIO，因此将标志都置上 fido->Flags |= fdo->Flags & (DO_DIRECT_IO | DO_BUFFERED_IO | DO_POWER_PAGABLE); // Clear the "initializing" flag so that we can get IRPs fido->Flags &= ~DO_DEVICE_INITIALIZING; } while (FALSE); // finish initialization if (!NT_SUCCESS(status)) { // need to cleanup if (pdx->LowerDeviceObject) IoDetachDevice(pdx->LowerDeviceObject); IoDeleteDevice(fido); } // need to cleanup KdPrint((DRIVERNAME " - ********************** AddDevice status is - %X **********************\n", status)); pdx->OutStandingIO = 1; KeInitializeSpinLock(&pdx->IOCountLock); KeInitializeEvent(&pdx->StopEvent, SynchronizationEvent, TRUE); KeInitializeEvent(&pdx->RemoveEvent, SynchronizationEvent, FALSE); pdx->bUSBRemove = false; /// 判断是否需要过滤设备 //初始化ValueName RtlInitUnicodeString( &ValueName, L"FilterIsRun"); if (QueryRegDWORD(&ValueName, &dwReturn) == 0) { KdPrint((DRIVERNAME " - AddDevice QueryRegDWORD return error! \n")); return status; } /// 是否查询,0不过滤非法设备，1开始过滤非法设备 if (0 == dwReturn) { KdPrint((DRIVERNAME " - AddDevice QueryRegDWORD return is 0x%d! \n", dwReturn)); return status; } ReadDeviceDescriptor(fido); return status; } // AddDevice /////////////////////////////////////////////////////////////////////////////// #pragma LOCKEDCODE NTSTATUS CompleteRequest(IN PIRP Irp, IN NTSTATUS status, IN ULONG_PTR info) { // CompleteRequest PIO_STACK_LOCATION pStackIrp = IoGetCurrentIrpStackLocation(Irp); Irp->IoStatus.Status = status; Irp->IoStatus.Information = info; IoCompleteRequest(Irp, IO_NO_INCREMENT); return status; } // CompleteRequest NTSTATUS USBSCSICompletion( IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context ) { PDEVICE_EXTENSION pdx = ( PDEVICE_EXTENSION ) DeviceObject->DeviceExtension; IoAcquireRemoveLock(&pdx->RemoveLock,Irp); PIO_STACK_LOCATION irpStack = IoGetCurrentIrpStackLocation( Irp ); PSCSI_REQUEST_BLOCK CurSrb=irpStack->Parameters.Scsi.Srb; PCDB cdb = (PCDB)CurSrb->Cdb; UCHAR opCode=cdb->CDB6GENERIC.OperationCode; if(opCode==SCSIOP_MODE_SENSE && CurSrb->DataBuffer && CurSrb->DataTransferLength >= sizeof(MODE_PARAMETER_HEADER)) { KdPrint(("SCSIOP_MODE_SENSE comming!\n")); PMODE_PARAMETER_HEADER modeData = (PMODE_PARAMETER_HEADER)CurSrb->DataBuffer; modeData->DeviceSpecificParameter |= MODE_DSP_WRITE_PROTECT; } if ( Irp->PendingReturned ) { IoMarkIrpPending( Irp ); } IoReleaseRemoveLock(&pdx->RemoveLock,Irp); return Irp->IoStatus.Status ; } #pragma LOCKEDCODE NTSTATUS DispatchForSCSI(IN PDEVICE_OBJECT fido, IN PIRP Irp) { // KdPrint((DRIVERNAME " - Enter DispatchForSCSI \n")); PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fido->DeviceExtension; PIO_STACK_LOCATION irpStack = IoGetCurrentIrpStackLocation(Irp); // Pass request down without additional processing NTSTATUS status; status = IoAcquireRemoveLock(&pdx->RemoveLock, Irp); if (!NT_SUCCESS(status)) return CompleteRequest(Irp, status, 0); IoCopyCurrentIrpStackLocationToNext(Irp); IoSetCompletionRoutine( Irp, USBSCSICompletion, NULL, TRUE, TRUE, TRUE ); status = IoCallDriver(pdx->LowerDeviceObject, Irp); IoReleaseRemoveLock(&pdx->RemoveLock, Irp); return status; } /////////////////////////////////////////////////////////////////////////////// #pragma LOCKEDCODE // make no assumptions about pageability of dispatch fcns NTSTATUS DispatchAny(IN PDEVICE_OBJECT fido, IN PIRP Irp) { // DispatchAny PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fido->DeviceExtension; PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp); #if DBG 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_