wmi.cpp

Programming the Microsoft Windows driver model.2nd 随书光盘。内有很多作者送的实用工具和随书源码。WDM编程

// WMI.CPP -- Windows Management Instrumentation handlers for wmiextra driver
// Copyright (C) 1999 by Walter Oney
// All rights reserved

#include "stddcls.h"
#include "driver.h"

#include <initguid.h>
#include <wdmguid.h>
#include "wmiextra.h"			// autogenerated by wmimofck

NTSTATUS QueryRegInfo(PDEVICE_OBJECT fdo, PULONG flags, PUNICODE_STRING instname, PUNICODE_STRING* regpath, PUNICODE_STRING resname, PDEVICE_OBJECT* pdo);
NTSTATUS QueryDataBlock(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG instcount, PULONG instlength, ULONG bufsize, PUCHAR buffer);
NTSTATUS SetDataBlock(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG bufsize, PUCHAR buffer);
NTSTATUS SetDataItem(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG id, ULONG bufsize, PUCHAR buffer);
NTSTATUS ExecuteMethod(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG id, ULONG cbInbuf, ULONG cbOutbuf, PUCHAR buffer);
NTSTATUS FunctionControl(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, WMIENABLEDISABLECONTROL fcn, BOOLEAN enable);

WMIGUIDREGINFO guidlist[] = {
	{&wmiextra_event_GUID, 1, WMIREG_FLAG_EVENT_ONLY_GUID},
	{&wmiextra_expensive_GUID, 1, WMIREG_FLAG_EXPENSIVE},
	{&wmiextra_method_GUID, 1, 0},
	{&wmiextra_event_test_GUID, 1, 0},
	};

#define INDEX_WMIEXTRA_EVENT		0
#define INDEX_WMIEXTRA_EXPENSIVE	1
#define INDEX_WMIEXTRA_METHOD		2
#define INDEX_WMIEXTRA_EVENT_TEST	3

WMILIB_CONTEXT libinfo = {
	arraysize(guidlist),		// GuidCount
	guidlist,					// GuidList
	QueryRegInfo,
	QueryDataBlock,
	SetDataBlock,
	SetDataItem,
	ExecuteMethod,
	FunctionControl,
	};
									  
///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS DispatchWmi(IN PDEVICE_OBJECT fdo, IN PIRP Irp)
	{							// DispatchWmi
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	NTSTATUS status = IoAcquireRemoveLock(&pdx->RemoveLock, Irp);
	if (!NT_SUCCESS(status))
		return CompleteRequest(Irp, status, 0);

#if DBG
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
	ULONG fcn = stack->MinorFunction;

	static char* fcnname[] = {
		"IRP_MN_QUERY_ALL_DATA",
		"IRP_MN_QUERY_SINGLE_INSTANCE",
		"IRP_MN_CHANGE_SINGLE_INSTANCE",
		"IRP_MN_CHANGE_SINGLE_ITEM",
		"IRP_MN_ENABLE_EVENTS",
		"IRP_MN_DISABLE_EVENTS",
		"IRP_MN_ENABLE_COLLECTION",
		"IRP_MN_DISABLE_COLLECTION",
		"IRP_MN_REGINFO",
		"IRP_MN_EXECUTE_METHOD",
		"<reserved>",
		"IRP_MN_REGINFO_EX",
		};

	KdPrint((DRIVERNAME " - WMI Request (%s)\n", fcnname[fcn]));
#endif // DBG

	// Delegate processing to the WMILIB helper library

	SYSCTL_IRP_DISPOSITION disposition;
	status = WmiSystemControl(&libinfo, fdo, Irp, &disposition);

	switch (disposition)
		{						// finish handling IRP

	case IrpProcessed:
		break;

	case IrpNotCompleted:
		IoCompleteRequest(Irp, IO_NO_INCREMENT); // WMILIB already filled in IoStatus fields
		break;

	default:
		IoSkipCurrentIrpStackLocation(Irp);
		status = IoCallDriver(pdx->LowerDeviceObject, Irp);
		break;
		}						// finish handling IRP

	IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
	return status;
	}							// DispatchWmi

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

#pragma PAGEDCODE

void WmiInitialize(PDEVICE_OBJECT fdo)
	{							// WmiInitialize
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;

	NTSTATUS status = IoWMIRegistrationControl(fdo, WMIREG_ACTION_REGISTER);
	if (!NT_SUCCESS(status))
		KdPrint((DRIVERNAME " - IoWMIRegistrationControl failed - %X\n", status));
	}							// WmiInitialize

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

#pragma PAGEDCODE

void WmiTerminate(PDEVICE_OBJECT fdo)
	{							// WmiTerminate
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	IoWMIRegistrationControl(fdo, WMIREG_ACTION_DEREGISTER);
	}							// WmiTerminate

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

#pragma PAGEDCODE

NTSTATUS QueryRegInfo(PDEVICE_OBJECT fdo, PULONG flags, PUNICODE_STRING instname, 
	PUNICODE_STRING* regpath, PUNICODE_STRING resname, PDEVICE_OBJECT* pdo)
	{							// QueryRegInfo
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	*flags = WMIREG_FLAG_INSTANCE_BASENAME;
	*regpath = &servkey;
	RtlInitUnicodeString(resname, L"MofResource");
	static WCHAR basename[] = L"WMIEXTRA";
	instname->Buffer = (PWCHAR) ExAllocatePool(PagedPool, sizeof(basename));
	if (!instname->Buffer)
		return STATUS_INSUFFICIENT_RESOURCES;
	instname->MaximumLength = sizeof(basename);
	instname->Length = sizeof(basename) - 2;
	RtlCopyMemory(instname->Buffer, basename, sizeof(basename));
	return STATUS_SUCCESS;
	}							// QueryRegInfo

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

#pragma PAGEDCODE

NTSTATUS QueryDataBlock(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, 
	ULONG instindex, ULONG instcount, PULONG instlength, ULONG bufavail, PUCHAR buffer)
	{							// QueryDataBlock
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	ULONG bufsize = 0;
	NTSTATUS status = STATUS_SUCCESS;

	switch (guidindex)
		{						// provide indicated data

	case INDEX_WMIEXTRA_EXPENSIVE:
		{						// GUID_WMIEXTRA_EXPENSIVE
		KdPrint((DRIVERNAME " - QueryDataBlock for wmiextra_expensive\n"));
		if (instindex != 0)
			{
			status = STATUS_WMI_INSTANCE_NOT_FOUND;
			break;
			}

		bufsize = wmiextra_expensive_SIZE;
		if (!instlength || bufavail < bufsize)
			status = STATUS_BUFFER_TOO_SMALL;
		else
			((Pwmiextra_expensive) buffer)->ExpensiveData = 42;
		break;
		}						// GUID_WMIEXTRA_EXPENSIVE

	case INDEX_WMIEXTRA_EVENT:
		KdPrint((DRIVERNAME " - QueryDataBlock for wmiextra_event even though defined as EVENT_ONLY!\n"));
		bufsize = 0;
		break;

	case INDEX_WMIEXTRA_METHOD:
		KdPrint((DRIVERNAME " - QueryDataBlock for wmiextra_method\n"));
		bufsize = 0;
		break;

	case INDEX_WMIEXTRA_EVENT_TEST:
		KdPrint((DRIVERNAME " - QueryDataBlock for wmiextra_event_test\n"));
		bufsize = wmiextra_event_test_SIZE;
		if (!instlength || bufavail < bufsize)
			status = STATUS_BUFFER_TOO_SMALL;
		else
			((Pwmiextra_event_test) buffer)->FireTheEvent = 0;
		break;

	default:
		status = STATUS_WMI_GUID_NOT_FOUND;
		}						// provide indicated data

	return WmiCompleteRequest(fdo, Irp, status, bufsize, IO_NO_INCREMENT);
	}							// QueryDataBlock

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

#pragma PAGEDCODE

NTSTATUS SetDataBlock(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG bufavail, PUCHAR buffer)
	{							// SetDataBlock
	// Note: the only settable data item in our schema is wmiextra_event_test, whose only
	// purpose is to provide a way for the test script to fire wmiextra_event. THIS IS
	// A HACK! In real life, the driver would decide on its own to trigger the event.

	if (guidindex != INDEX_WMIEXTRA_EVENT_TEST)
		return WmiCompleteRequest(fdo, Irp, STATUS_WMI_GUID_NOT_FOUND, 0, IO_NO_INCREMENT);
	if (instindex != 0)
		return WmiCompleteRequest(fdo, Irp, STATUS_WMI_INSTANCE_NOT_FOUND, 0, IO_NO_INCREMENT);

	// Construct an event object in nonpaged memory. The system will free this
	// when it's no longer needed

	Pwmiextra_event junk = (Pwmiextra_event) ExAllocatePool(NonPagedPool, wmiextra_event_SIZE);
	if (junk)
		{						// fire event
		junk->EventInfo = 42;
		NTSTATUS status = WmiFireEvent(fdo, (LPGUID) &wmiextra_event_GUID, 0, wmiextra_event_SIZE, junk);
		if (!NT_SUCCESS(status))
			KdPrint((DRIVERNAME " - WmiFireEvent failed - %X\n", status));
		else
			KdPrint((DRIVERNAME " - Firing wmiextra_event\n"));
		}

	return WmiCompleteRequest(fdo, Irp, STATUS_SUCCESS, 0, IO_NO_INCREMENT);
	}							// SetDataBlock

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

#pragma PAGEDCODE

NTSTATUS SetDataItem(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG id, ULONG bufavail, PUCHAR buffer)
	{							// SetDataItem

	return WmiCompleteRequest(fdo, Irp, STATUS_INVALID_DEVICE_REQUEST, 0, IO_NO_INCREMENT);
	}							// SetDataItem

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

#pragma PAGEDCODE

NTSTATUS ExecuteMethod(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG id, ULONG cbInbuf, ULONG cbOutbuf, PUCHAR buffer)
	{							// ExecuteMethod
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	NTSTATUS status = STATUS_SUCCESS;
	ULONG bufused = 0;

	// This driver defines methods only for the wmiextra_method class. More generally, you would
	// probably want to switch on "guidindex" and then inspect "instindex" and "id".

	if (guidindex != INDEX_WMIEXTRA_METHOD)
		return WmiCompleteRequest(fdo, Irp, STATUS_WMI_GUID_NOT_FOUND, 0, IO_NO_INCREMENT);
	if (instindex != 0)
		return WmiCompleteRequest(fdo, Irp, STATUS_WMI_INSTANCE_NOT_FOUND, 0, IO_NO_INCREMENT);

	// The only method defined on this class is "AnswerMethod"...

	if (id != AnswerMethod)
		return WmiCompleteRequest(fdo, Irp, STATUS_WMI_ITEMID_NOT_FOUND, 0, IO_NO_INCREMENT);

	KdPrint((DRIVERNAME " - ExecuteMethod\n"));

	if (cbInbuf < AnswerMethod_IN_SIZE)
		status = STATUS_INVALID_PARAMETER;
	else if (cbOutbuf < AnswerMethod_OUT_SIZE)
		status = STATUS_BUFFER_TOO_SMALL;
	else
		{						// perform method

		// On input, "buffer" contains the input parameters laid out to match an AnswerMethod_IN structure.
		// On output, we fill "buffer" with output values laid out to match an AnswerMethod_OUT structure.
		// These structure declarations are in "wmiextra.h" because we used the "-m" command
		// line option to wmimofck.

		PAnswerMethod_IN in = (PAnswerMethod_IN) buffer;
		PAnswerMethod_OUT out = (PAnswerMethod_OUT) buffer;

		out->TheAnswer = in->TheAnswer + 1;
		bufused = AnswerMethod_OUT_SIZE;
		}						// perform method

	return WmiCompleteRequest(fdo, Irp, status, bufused, IO_NO_INCREMENT);
	}							// ExecuteMethod

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

#pragma PAGEDCODE

NTSTATUS FunctionControl(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, WMIENABLEDISABLECONTROL fcn, BOOLEAN enable)
	{							// FunctionControl
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	NTSTATUS status = STATUS_SUCCESS;

	switch (guidindex)
		{						// process control operation
	case INDEX_WMIEXTRA_EVENT:
		if (fcn == WmiEventControl)
			{
			KdPrint((DRIVERNAME " - FunctionControl(WmiEventControl) for wmiextra_event: %s\n",
				enable ? "enable" : "disable"));
			pdx->evenable = enable;
			}
		else
			{
			KdPrint((DRIVERNAME " - FunctionControl(WmiDataBlockControl) for wmiextra_event!\n"));
			status = STATUS_WMI_NOT_SUPPORTED;
			}
		break;

	case INDEX_WMIEXTRA_EXPENSIVE:
		if (fcn == WmiDataBlockControl)
			{
			KdPrint((DRIVERNAME " - FunctionControl(WmiDataBlockControl) for wmiextra_expensive: %s\n",
				enable ? "enable" : "disable"));
			pdx->expenable = enable;
			}
		else
			{
			KdPrint((DRIVERNAME " - FunctionControl(WmiEventControl) for wmiextra_expensive!\n"));
			status = STATUS_WMI_NOT_SUPPORTED;
			}
		break;

	default:
		KdPrint((DRIVERNAME " - FunctionControl for invalid GUID # %d\n", guidindex));
		status = STATUS_WMI_GUID_NOT_FOUND;
		break;
		}						// process control operation

	return WmiCompleteRequest(fdo, Irp, status, 0, IO_NO_INCREMENT);
	}							// FunctionControl