// Main program for powtrace driver
// Copyright (C) 2002 by Walter Oney
// All rights reserved

// POWTRACE.SYS is intended as a lower filter that will show all the power-related irps flowing
// downward from a function driver. The resulting debug trace can be useful in diagnosing power
// management problems with drivers.

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

NTSTATUS AddDevice(IN PDRIVER_OBJECT DriverObject, IN PDEVICE_OBJECT pdo);
NTSTATUS CapabilitiesQueryCompletionRoutine(PDEVICE_OBJECT fido, PIRP Irp, PDEVICE_EXTENSION pdx);
VOID DriverUnload(IN PDRIVER_OBJECT DriverObject);
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 PowerCompletionRoutine(PDEVICE_OBJECT fido, PIRP Irp, PDEVICE_EXTENSION pdx);
VOID PowerTrace(PDEVICE_EXTENSION pdx, char* ctl, ...);
VOID ShowCapabilities(PDEVICE_EXTENSION pdx, PDEVICE_CAPABILITIES pdc);
NTSTATUS StartDeviceCompletionRoutine(PDEVICE_OBJECT fido, PIRP Irp, PDEVICE_EXTENSION pdx);
NTSTATUS UsageNotificationCompletionRoutine(PDEVICE_OBJECT fido, PIRP Irp, PDEVICE_EXTENSION pdx);

char* sysnames[] = {
	"PowerSystemUnspecified",
	"PowerSystemWorking",
	"PowerSystemSleeping1",
	"PowerSystemSleeping2",
	"PowerSystemSleeping3",
	"PowerSystemHibernate",
	"PowerSystemShutdown",
	};

char* devnames[] = {
	"PowerDeviceUnspecified",
	"PowerDeviceD0",
	"PowerDeviceD1",
	"PowerDeviceD2",
	"PowerDeviceD3",
	};

char* powminor[] = {
	"IRP_MN_WAIT_WAKE",
	"IRP_MN_POWER_SEQUENCE",
	"IRP_MN_SET_POWER",
	"IRP_MN_QUERY_POWER",
	};

FAST_MUTEX filelock;			// lock for file I/O

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

#pragma PAGEDCODE

extern "C" NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject,
	IN PUNICODE_STRING RegistryPath)
	{							// DriverEntry
	DbgPrint(DRIVERNAME " - Entering DriverEntry: DriverObject %8.8lX\n", DriverObject);

	// Initialize function pointers

	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;

	ExInitializeFastMutex(&filelock);
	
	return STATUS_SUCCESS;
	}							// DriverEntry

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

#pragma PAGEDCODE

VOID DriverUnload(IN PDRIVER_OBJECT DriverObject)
	{							// DriverUnload
	PAGED_CODE();
	DbgPrint(DRIVERNAME " - Entering DriverUnload: DriverObject %8.8lX\n", DriverObject);
	}							// DriverUnload

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

#pragma PAGEDCODE

NTSTATUS AddDevice(IN PDRIVER_OBJECT DriverObject, IN PDEVICE_OBJECT pdo)
	{							// AddDevice
	PAGED_CODE();
	DbgPrint(DRIVERNAME " - Entering AddDevice: DriverObject %8.8lX, pdo %8.8lX\n", DriverObject, pdo);

	NTSTATUS status;

	// Create a filter device object

	PDEVICE_OBJECT fdo;

	status = IoCreateDevice(DriverObject, sizeof(DEVICE_EXTENSION), NULL,
		GetDeviceTypeToUse(pdo), 0, FALSE, &fdo);
	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) fdo->DeviceExtension;

	// From this point forward, any error will have side effects that need to
	// be cleaned up. Using a do-once block allows us to modify the program
	// easily without losing track of the side effects.

	do
		{						// finish initialization
		pdx->DeviceObject = fdo;
		pdx->Pdo = pdo;
		IoInitializeRemoveLock(&pdx->RemoveLock, 0, 0, 0);

		// Add our device object to the stack and propagate critical settings
		// from the immediately lower device object

		PDEVICE_OBJECT ldo = IoAttachDeviceToDeviceStack(fdo, pdo);
		if (!ldo)
			{						// can't attach device
			KdPrint((DRIVERNAME " - IoAttachDeviceToDeviceStack failed\n"));
			status = STATUS_DEVICE_REMOVED;
			break;
			}						// can't attach device

		pdx->LowerDeviceObject = ldo;

		// Copy the flags related to I/O buffering from the lower device object so the I/O manager
		// will create the expected data structures for reads and writes.

		fdo->Flags |= ldo->Flags & (DO_DIRECT_IO | DO_BUFFERED_IO | DO_POWER_PAGABLE);

		// Get registry flag FileLogging to decide whether to log to a disk file (\SystemRoot\powtrace.log)
		// or not.

		HANDLE hkey;
		status = IoOpenDeviceRegistryKey(pdo, PLUGPLAY_REGKEY_DEVICE, KEY_READ, &hkey);
		if (NT_SUCCESS(status))
			{					// check FileLogging flag
			UNICODE_STRING valname;
			RtlInitUnicodeString(&valname, L"FileLogging");

			ULONG junk;
			KEY_VALUE_PARTIAL_INFORMATION value;	// has room for a REG_DWORD value

			status = ZwQueryValueKey(hkey, &valname, KeyValuePartialInformation,
				&value, sizeof(value), &junk);

			if (NT_SUCCESS(status) && value.DataLength == sizeof(ULONG))
				pdx->filetrace = *(PULONG) value.Data != 0;

			ZwClose(hkey);
			}					// check FileLogging flag

		status = STATUS_SUCCESS;	// don't actually care if registry query worked or not

		PowerTrace(pdx, DRIVERNAME " - Initiating trace\n");

		// Clear the "initializing" flag so that we can get IRPs

		fdo->Flags &= ~DO_DEVICE_INITIALIZING;
		}						// finish initialization
	while (FALSE);

	if (!NT_SUCCESS(status))
		{					// need to cleanup
		if (pdx->LowerDeviceObject)
			IoDetachDevice(pdx->LowerDeviceObject);
		IoDeleteDevice(fdo);
		}					// need to cleanup

	return status;
	}							// AddDevice

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

#pragma LOCKEDCODE

NTSTATUS CapabilitiesQueryCompletionRoutine(PDEVICE_OBJECT fido, PIRP Irp, PDEVICE_EXTENSION pdx)
	{							// CapabilitiesQueryCompletionRoutine
	if (Irp->PendingReturned)
		IoMarkIrpPending(Irp);

	if (NT_SUCCESS(Irp->IoStatus.Status))
		ShowCapabilities(pdx, IoGetCurrentIrpStackLocation(Irp)->Parameters.DeviceCapabilities.Capabilities);

	IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
	return STATUS_SUCCESS;
	}							// CapabilitiesQueryCompletionRoutine

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

#pragma LOCKEDCODE

NTSTATUS CompleteRequest(IN PIRP Irp, IN NTSTATUS status, IN ULONG_PTR info)
	{							// CompleteRequest
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = info;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return status;
	}							// CompleteRequest

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

#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);

	// Pass request down without additional processing

	NTSTATUS status;
	status = IoAcquireRemoveLock(&pdx->RemoveLock, Irp);
	if (!NT_SUCCESS(status))
		return CompleteRequest(Irp, status, 0);
	IoSkipCurrentIrpStackLocation(Irp);
	status = IoCallDriver(pdx->LowerDeviceObject, Irp);
	IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
	return status;
	}							// DispatchAny

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

NTSTATUS DispatchPnp(IN PDEVICE_OBJECT fido, IN PIRP Irp)
	{							// DispatchPnp
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
	ULONG fcn = stack->MinorFunction;

	NTSTATUS status;
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fido->DeviceExtension;
	status = IoAcquireRemoveLock(&pdx->RemoveLock, Irp);
	if (!NT_SUCCESS(status))
		return CompleteRequest(Irp, status, 0);

	// Handle usage notification specially in order to track power pageable
	// flag correctly

	if (fcn == IRP_MN_DEVICE_USAGE_NOTIFICATION)
		{						// usage notification
		if (!fido->AttachedDevice || (fido->AttachedDevice->Flags & DO_POWER_PAGABLE))
			fido->Flags |= DO_POWER_PAGABLE;
		IoCopyCurrentIrpStackLocationToNext(Irp);
		IoSetCompletionRoutine(Irp, (PIO_COMPLETION_ROUTINE) UsageNotificationCompletionRoutine,
			(PVOID) pdx, TRUE, TRUE, TRUE);
		return IoCallDriver(pdx->LowerDeviceObject, Irp);
		}						// usage notification

	// Handle start device specially in order to correctly inherit
	// FILE_REMOVABLE_MEDIA

	if (fcn == IRP_MN_START_DEVICE)
		{						// device start
		IoCopyCurrentIrpStackLocationToNext(Irp);
		IoSetCompletionRoutine(Irp, (PIO_COMPLETION_ROUTINE) StartDeviceCompletionRoutine,
			(PVOID) pdx, TRUE, TRUE, TRUE);
		return IoCallDriver(pdx->LowerDeviceObject, Irp);
		}						// device start

	// Handle remove device specially in order to cleanup device stack

	if (fcn == IRP_MN_REMOVE_DEVICE)
		{						// remove device
		IoSkipCurrentIrpStackLocation(Irp);
		status = IoCallDriver(pdx->LowerDeviceObject, Irp);
		IoReleaseRemoveLockAndWait(&pdx->RemoveLock, Irp);
		RemoveDevice(fido);
		return status;
		}						// remove device

	// Handle capabilities query specially in order to show the structure