diskperf.cpp

实现USB存储设备的只读操作,可以实现所有USB设备

extern "C"
{

#define INITGUID

#include "ntddk.h"
#include "ntdddisk.h"
#include "stdarg.h"
#include "stdio.h"
#include <ntddvol.h>

#include <mountdev.h>
#include "wmistr.h"
#include "wmidata.h"
#include "wmiguid.h"
#include "wmilib.h"
#include "scsi.h"

#ifdef POOL_TAGGING
#ifdef ExAllocatePool
#undef ExAllocatePool
#endif
#define ExAllocatePool(a,b) ExAllocatePoolWithTag(a,b,'frPD')
#endif

#define DISKPERF_MAXSTR         64

//////////////////////////////////////////////////////////////////////////
// 控制码
#define DISK_READONLY	\
	CTL_CODE(FILE_DEVICE_UNKNOWN, 0x830, METHOD_BUFFERED, FILE_ANY_ACCESS)

#define DISK_NONREADONLY	\
	CTL_CODE(FILE_DEVICE_UNKNOWN, 0x831, METHOD_BUFFERED, FILE_ANY_ACCESS)


//  是否仅可读

BOOLEAN IsReadOnly = FALSE;

//////////////////////////////////////////////////////////////////////////

// Device Extension

typedef struct _DEVICE_EXTENSION {

    // Back pointer to device object

    PDEVICE_OBJECT DeviceObject;

    // Target Device Object

    PDEVICE_OBJECT TargetDeviceObject;

    // Physical device object
    PDEVICE_OBJECT PhysicalDeviceObject;

    // Disk number for reference in WMI

    ULONG       DiskNumber;

    // If device is enabled for counting always

    LONG        EnabledAlways;

    // Use to keep track of Volume info from ntddvol.h

    WCHAR StorageManagerName[8];

    // Disk performance counters
    // and locals used to compute counters

    ULONG   Processors;
    PDISK_PERFORMANCE DiskCounters;    // per processor counters
    LARGE_INTEGER LastIdleClock;
    LONG QueueDepth;
    LONG CountersEnabled;

    // must synchronize paging path notifications
    KEVENT PagingPathCountEvent;
    ULONG  PagingPathCount;

    // Physical Device name or WMI Instance Name

    UNICODE_STRING PhysicalDeviceName;
    WCHAR PhysicalDeviceNameBuffer[DISKPERF_MAXSTR];

    // Private context for using WmiLib
    WMILIB_CONTEXT WmilibContext;

	// sun add
	struct _DEVICE_EXTENSION* edx;

} DEVICE_EXTENSION, *PDEVICE_EXTENSION;

typedef struct _EXTRA_DEVICE_EXTENSION:public DEVICE_EXTENSION
{
	PDEVICE_EXTENSION pdx;

} EXTRA_DEVICE_EXTENSION,*PEXTRA_DEVICE_EXTENSION;

#define DEVICE_EXTENSION_SIZE sizeof(DEVICE_EXTENSION)
#define PROCESSOR_COUNTERS_SIZE FIELD_OFFSET(DISK_PERFORMANCE, QueryTime)

/*
Layout of Per Processor Counters is a contiguous block of memory:
    Processor 1
+-----------------------+     +-----------------------+
|PROCESSOR_COUNTERS_SIZE| ... |PROCESSOR_COUNTERS_SIZE|
+-----------------------+     +-----------------------+
where PROCESSOR_COUNTERS_SIZE is less than sizeof(DISK_PERFORMANCE) since
we only put those we actually use for counting.
*/

UNICODE_STRING DiskPerfRegistryPath;

// Function declarations

NTSTATUS
DriverEntry(
    IN PDRIVER_OBJECT DriverObject,
    IN PUNICODE_STRING RegistryPath
    );

NTSTATUS
DiskPerfForwardIrpSynchronous(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

NTSTATUS
DiskPerfAddDevice(
    IN PDRIVER_OBJECT DriverObject,
    IN PDEVICE_OBJECT PhysicalDeviceObject
    );


NTSTATUS
DiskPerfDispatchPnp(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

NTSTATUS
DiskPerfDispatchPower(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

NTSTATUS
DiskPerfStartDevice(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

NTSTATUS
DiskPerfRemoveDevice(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

NTSTATUS
DiskPerfSendToNextDriver(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );


NTSTATUS
DiskPerfCreate(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

NTSTATUS
DiskPerfReadWrite(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

NTSTATUS
DiskPerfIoCompletion(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp,
    IN PVOID Context
    );

NTSTATUS
DiskPerfDeviceControl(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

NTSTATUS
DiskPerfShutdownFlush(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

VOID
DiskPerfUnload(
    IN PDRIVER_OBJECT DriverObject
    );

NTSTATUS DiskPerfWmi(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

VOID
DiskPerfLogError(
    IN PDEVICE_OBJECT DeviceObject,
    IN ULONG UniqueId,
    IN NTSTATUS ErrorCode,
    IN NTSTATUS Status
    );

NTSTATUS
DiskPerfRegisterDevice(
    IN PDEVICE_OBJECT DeviceObject
    );

NTSTATUS
DiskPerfIrpCompletion(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp,
    IN PVOID Context
    );

NTSTATUS
DiskperfQueryWmiRegInfo(
    IN PDEVICE_OBJECT DeviceObject,
    OUT ULONG *RegFlags,
    OUT PUNICODE_STRING InstanceName,
    OUT PUNICODE_STRING *RegistryPath,
    OUT PUNICODE_STRING MofResourceName,
    OUT PDEVICE_OBJECT *Pdo
    );

NTSTATUS
DiskperfQueryWmiDataBlock(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp,
    IN ULONG GuidIndex,
    IN ULONG InstanceIndex,
    IN ULONG InstanceCount,
    IN OUT PULONG InstanceLengthArray,
    IN ULONG BufferAvail,
    OUT PUCHAR Buffer
    );

VOID
DiskPerfSyncFilterWithTarget(
    IN PDEVICE_OBJECT FilterDevice,
    IN PDEVICE_OBJECT TargetDevice
    );

NTSTATUS
DiskperfWmiFunctionControl(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp,
    IN ULONG GuidIndex,
    IN WMIENABLEDISABLECONTROL Function,
    IN BOOLEAN Enable
    );

VOID
DiskPerfAddCounters(
    IN OUT PDISK_PERFORMANCE TotalCounters,
    IN PDISK_PERFORMANCE NewCounters,
    IN LARGE_INTEGER Frequency
    );

NTSTATUS
DiskPerfScsi(
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	);

#if DBG

ULONG DiskPerfDebug = 0;

VOID
DiskPerfDBGOUT(
    ULONG DBGOUTLevel,
    PCCHAR DebugMessage,
    ...
    );

#define DBGOUT(x)   DiskPerfDBGOUT x

#else

#define DBGOUT(x)

#endif

// Define the sections that allow for discarding (i.e. paging) some of
// the code.

#ifdef ALLOC_PRAGMA
#pragma alloc_text (INIT, DriverEntry)
#pragma alloc_text (PAGE, DiskPerfCreate)
#pragma alloc_text (PAGE, DiskPerfAddDevice)
#pragma alloc_text (PAGE, DiskPerfDispatchPnp)
#pragma alloc_text (PAGE, DiskPerfStartDevice)
#pragma alloc_text (PAGE, DiskPerfRemoveDevice)
#pragma alloc_text (PAGE, DiskPerfUnload)
#pragma alloc_text (PAGE, DiskPerfWmi)
#pragma alloc_text (PAGE, DiskperfQueryWmiRegInfo)
#pragma alloc_text (PAGE, DiskperfQueryWmiDataBlock)
#pragma alloc_text (PAGE, DiskPerfRegisterDevice)
#pragma alloc_text (PAGE, DiskPerfSyncFilterWithTarget)
#pragma alloc_text (PAGE, DiskPerfScsi)
#endif

WMIGUIDREGINFO DiskperfGuidList[] =
{
    { &DiskPerfGuid,
      1,
      0
    }
};

#define DiskperfGuidCount (sizeof(DiskperfGuidList) / sizeof(WMIGUIDREGINFO))

#define USE_PERF_CTR

#ifdef USE_PERF_CTR
#define DiskPerfGetClock(a, b) (a) = KeQueryPerformanceCounter((b))
#else
#define DiskPerfGetClock(a, b) KeQuerySystemTime(&(a))
#endif


/*++

Routine Description:

Installable driver initialization entry point.
This entry point is called directly by the I/O manager to set up the disk
performance driver. The driver object is set up and then the Pnp manager
calls DiskPerfAddDevice to attach to the boot devices.

Arguments:

DriverObject - The disk performance driver object.

RegistryPath - pointer to a unicode string representing the path,
to driver-specific key in the registry.

Return Value:

STATUS_SUCCESS if successful

--*/

NTSTATUS
DriverEntry(
    IN PDRIVER_OBJECT DriverObject,
    IN PUNICODE_STRING RegistryPath
    )

{

    ULONG               ulIndex;
    PDRIVER_DISPATCH  * dispatch;
	DbgPrint(("access DriverEntry success in diskperf.c! \n"));

    // Remember registry path

    DiskPerfRegistryPath.MaximumLength = RegistryPath->Length
                                            + sizeof(UNICODE_NULL);
    DiskPerfRegistryPath.Buffer = (PWSTR)ExAllocatePool(
                                    PagedPool,
                                    DiskPerfRegistryPath.MaximumLength); // ExAllocatePool 按类型进行申请内存
    if (DiskPerfRegistryPath.Buffer != NULL)
    {
        RtlCopyUnicodeString(&DiskPerfRegistryPath, RegistryPath); // 复制UNICODE STRING 
    } else {
        DiskPerfRegistryPath.Length = 0;
        DiskPerfRegistryPath.MaximumLength = 0;
    }

    // Create dispatch points

    for (ulIndex = 0, dispatch = DriverObject->MajorFunction;
         ulIndex <= IRP_MJ_MAXIMUM_FUNCTION;
         ulIndex++, dispatch++) {

        *dispatch = DiskPerfSendToNextDriver;
    }
	
	
    // Set up the device driver entry points.

    DriverObject->MajorFunction[IRP_MJ_CREATE]          = DiskPerfCreate;
    DriverObject->MajorFunction[IRP_MJ_READ]            = DiskPerfReadWrite;
    DriverObject->MajorFunction[IRP_MJ_WRITE]           = DiskPerfReadWrite;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]  = DiskPerfDeviceControl;
    DriverObject->MajorFunction[IRP_MJ_SYSTEM_CONTROL]  = DiskPerfWmi;

    DriverObject->MajorFunction[IRP_MJ_SHUTDOWN]        = DiskPerfShutdownFlush;
    DriverObject->MajorFunction[IRP_MJ_FLUSH_BUFFERS]   = DiskPerfShutdownFlush;
    DriverObject->MajorFunction[IRP_MJ_PNP]             = DiskPerfDispatchPnp;
    DriverObject->MajorFunction[IRP_MJ_POWER]           = DiskPerfDispatchPower;

    DriverObject->DriverExtension->AddDevice            = DiskPerfAddDevice;
    DriverObject->DriverUnload                          = DiskPerfUnload;

	DriverObject->MajorFunction[IRP_MJ_SCSI]			= DiskPerfScsi;
	DbgPrint(("DriverEntry exit! in diskperf.c \n"));


    return(STATUS_SUCCESS);

} // end DriverEntry()

#define FILTER_DEVICE_PROPOGATE_FLAGS            0
#define FILTER_DEVICE_PROPOGATE_CHARACTERISTICS (FILE_REMOVABLE_MEDIA |  \
                                                 FILE_READ_ONLY_DEVICE | \
                                                 FILE_FLOPPY_DISKETTE    \
                                                 )



NTSTATUS
USBSCSICompletion( IN PDEVICE_OBJECT DeviceObject,
				  IN PIRP Irp,
				  IN PVOID Context )
{
	PDEVICE_EXTENSION        deviceExtension;
	NTSTATUS                status;
	PIO_STACK_LOCATION        irpStack;

	PSCSI_REQUEST_BLOCK        CurSrb;
	PMODE_PARAMETER_HEADER    modeData;
	PDEVICE_OBJECT            pDeviceObject;

	PCDB    cdb ;
	UCHAR    opCode;
	KIRQL  IrqLevel;


	PDEVICE_EXTENSION StorExtension = ( PDEVICE_EXTENSION )	DeviceObject->DeviceExtension;

	irpStack = IoGetCurrentIrpStackLocation( Irp );

	CurSrb = (PSCSI_REQUEST_BLOCK)ExAllocatePoolWithTag(NonPagedPool,
		sizeof(SCSI_REQUEST_BLOCK),
		'brs');
	if (CurSrb == NULL) 
	{
		//DbgPrint("观察：CurSrb==NULL\n");
	}
	else
	{
		//DbgPrint("观察：CurSrb!=NULL\n");
	}

	RtlZeroMemory(CurSrb, SCSI_REQUEST_BLOCK_SIZE); 

	if (irpStack->MajorFunction==IRP_MJ_INTERNAL_DEVICE_CONTROL)
	{