#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_