// Power request handler for Generic driver
// Copyright (C) 1999 by Walter Oney
// All rights reserved

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

NTSTATUS DefaultPowerHandler(PGENERIC_EXTENSION pdx, PIRP Irp);
NTSTATUS WaitWakeCompletionRoutine(PDEVICE_OBJECT junk, PIRP Irp, PGENERIC_EXTENSION pdx);
VOID WaitWakeCallback(PDEVICE_OBJECT junk, UCHAR MinorFunction, POWER_STATE state,
	PGENERIC_EXTENSION pdx, PIO_STATUS_BLOCK pstatus);

enum POWSTATE {
	InitialState = 0,				// initial state of FSM
	SysPowerUpPending,				// system power-up IRP forwarded
	SubPowerUpPending,				// waiting for nested device power up to finish
	SubPowerDownPending,			// waiting from device to power down before forwarding system power-down IRP
	SysPowerDownPending,			// waiting for system power-down IRP to finish
	DevPowerUpPending,				// waiting for device power-up IRP
	DevPowerDownPending,			// waiting for device power-down IRP
	ContextSavePending,				// context save is underway
	ContextRestorePending,			// context restore is underway
	DevQueryUpPending,				// device query for power-up pending
	DevQueryDownPending,			// device query for power-down pending
	QueueStallPending,				// waiting for device to be idle
	SaveSeqPending,					// waiting to get sequence numbers after context save
	RestoreSeqPending,				// waiting to get sequence numbers before context restore
	PassiveCompletePending,			// waiting for PASSIVE_LEVEL callback to complete IRP
	FinalState,						// final state of FSM
	NUMPOWSTATES,
	};

enum POWEVENT {
	NewIrp = 0,						// new query/set IRP
	MainIrpComplete,				// the main IRP has finished
	AsyncNotify,					// some other event has occurred
	NUMPOWEVENTS,
	};

typedef struct _POWCONTEXT {
	LONG id;						// unique sequence number for this IRP
	LONG eventcount;				// number of events generated for this IRP
	PGENERIC_EXTENSION pdx;			// our own device extension
	PIRP irp;						// the IRP we're processing
	enum POWSTATE state;			// current state of FSM
	NTSTATUS status;				// completion status for main IRP
	DEVICE_POWER_STATE devstate;	// device power state to use
	POWER_SEQUENCE sequence;		// sequence numbers retrieved by ditto
	DEVICE_POWER_STATE oldpower;	// previous device power state
	UCHAR MinorFunction;			// minor function to use in requested power IRP
	BOOLEAN UnstallQueue;			// unstall queue when main IRP finishes
	} POWCONTEXT, *PPOWCONTEXT;

NTSTATUS HandlePowerEvent(PPOWCONTEXT ctx, enum POWEVENT event);

static LONG ctxcount = 0;			// counter for POWCONTEXT structures

// #define VERBOSETRACE				// for debugging HandlePowerEvent

#if DBG

	#ifdef VERBOSETRACE
		#define POWTRACE(x) DbgPrint x
	#else
		#define POWTRACE(x)
	#endif

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

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

	static char* devstate[] = {
		"PowerDeviceUnspecified",
		"PowerDeviceD0",
		"PowerDeviceD1",
		"PowerDeviceD2",
		"PowerDeviceD3",
		"PowerDeviceMaximum",
		};

#else
	#define POWTRACE(x)
#endif // DBG
									  
///////////////////////////////////////////////////////////////////////////////

GENERICAPI NTSTATUS GENERIC_EXPORT GenericDispatchPower(PGENERIC_EXTENSION pdx, PIRP Irp)
	{							// DispatchPower
	NTSTATUS status = STATUS_SUCCESS;
	if (pdx->RemoveLock)
		{
		status = IoAcquireRemoveLock(pdx->RemoveLock, Irp);
		if (!NT_SUCCESS(status))
			{
			PoStartNextPowerIrp(Irp);
			return CompleteRequest(Irp, status);
			}
		}

	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
	ASSERT(stack->MajorFunction == IRP_MJ_POWER);

	ULONG fcn = stack->MinorFunction;

	if (fcn == IRP_MN_SET_POWER || fcn == IRP_MN_QUERY_POWER)
		{						// handle set/query

#if DBG
		if (stack->Parameters.Power.Type == SystemPowerState)
			KdPrint(("%s - POWER Request (%s), S-state = %s\n", pdx->DebugName, fcnname[fcn], sysstate[stack->Parameters.Power.State.SystemState]));
		else
			KdPrint(("%s - POWER Request (%s), D-state = %s\n", pdx->DebugName, fcnname[fcn], devstate[stack->Parameters.Power.State.DeviceState]));
#endif // DBG

		// Create a context structure and launch the finite state machine that will process
		// this IRP asynchronously. The initial call to HandlePowerEvent should return
		// STATUS_PENDING. The FSM will eventually complete the IRP.

		PPOWCONTEXT ctx = (PPOWCONTEXT) ExAllocatePool(NonPagedPool, sizeof(POWCONTEXT));
		if (!ctx)
			{
			KdPrint(("%s - Can't allocate power context structure\n", pdx->DebugName));
			PoStartNextPowerIrp(Irp);
			status = CompleteRequest(Irp, STATUS_INSUFFICIENT_RESOURCES);
			}
		else
			{				// process this IRP
			RtlZeroMemory(ctx, sizeof(POWCONTEXT));
			ctx->pdx = pdx;
			ctx->irp = Irp;
			ctx->id = InterlockedIncrement(&ctxcount);
			status = HandlePowerEvent(ctx, NewIrp);
			}				// process this IRP
		}						// handle set/query

	else
		{						// handle other power request
	#if DBG
		if (fcn < arraysize(fcnname))
			KdPrint(("%s - POWER Request (%s)\n", pdx->DebugName, fcnname[fcn]));
		else
			KdPrint(("%s - POWER Request ?? (0x%X)\n", fcn));
	#endif

		// Install a completion routine for a WAIT_WAKE so we're sure to nullify our
		// cached pointer before the IRP disappears.

		if (fcn == IRP_MN_WAIT_WAKE)
			{					// wait-wake IRP
			PoStartNextPowerIrp(Irp);
			IoCopyCurrentIrpStackLocationToNext(Irp);
			IoSetCompletionRoutine(Irp, (PIO_COMPLETION_ROUTINE) WaitWakeCompletionRoutine, pdx, TRUE, TRUE, TRUE);
			status = PoCallDriver(pdx->LowerDeviceObject, Irp);
			}					// wait-wake IRP

		// Simply forward any other power IRP. At the present time, the only
		// kind it could be is IRP_MN_POWER_SEQUENCE, which probably no-one
		// actually uses.

		else
			status = DefaultPowerHandler(pdx, Irp);
		}						// handle other power request
	
	if (pdx->RemoveLock)
		IoReleaseRemoveLock(pdx->RemoveLock, Irp);

	return status;
	}							// DispatchPower

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

GENERICAPI NTSTATUS GENERIC_EXPORT GenericHandlePowerIoctl(PGENERIC_EXTENSION pdx, PIRP Irp)
	{							// GenericHandlePowerIoctl
	ASSERT(KeGetCurrentIrql() == PASSIVE_LEVEL);
	ULONG info = 0;
	NTSTATUS status = STATUS_SUCCESS;

	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
	if (stack->MajorFunction != IRP_MJ_DEVICE_CONTROL)
		return STATUS_INVALID_DEVICE_REQUEST;

	ULONG cbin = stack->Parameters.DeviceIoControl.InputBufferLength;
	ULONG cbout = stack->Parameters.DeviceIoControl.OutputBufferLength;
	ULONG code = stack->Parameters.DeviceIoControl.IoControlCode;

	switch (code)
		{						// process request

	case IOCTL_GET_POWERINFO:				// code == 0xFF0
		{						// IOCTL_GET_POWERINFO
		if (cbout < sizeof(POWERINFO))
			{
			status = STATUS_INVALID_PARAMETER;
			break;
			}
		
		PPOWERINFO pip = (PPOWERINFO) Irp->AssociatedIrp.SystemBuffer;

		// Get timeout constants from the registry

		status = GetPowerInfoFromRegistry(pdx, pip);
		if (!NT_SUCCESS(status))
			{					// initialize parameters
			pip->ulConservation = 0;
			pip->ulPerformance = 0;
			status = STATUS_SUCCESS;
			}					// initialize parameters

		// Determine current device state

		pip->bCanIdle = (pdx->Flags & GENERIC_IDLE_DETECT) != 0;
		pip->bCanWake = pdx->devcaps.SystemWake != PowerSystemUnspecified && pdx->devcaps.DeviceWake != PowerDeviceUnspecified;
		pip->bWakeup = pdx->WakeupEnabled;
		pip->bIdle = pdx->devpower > PowerDeviceD0;
		
		info = sizeof(POWERINFO);
		break;
		}						// IOCTL_GET_POWERINFO

	case IOCTL_SET_POWERINFO:				// code == 0xFF1
		{						// IOCTL_SET_POWERINFO
		if (cbin < sizeof(POWERINFO))
			{
			status = STATUS_INVALID_PARAMETER;
			break;
			}

		// If we're currently idle, restore power. This is in case we've just been told to
		// never idle in the power mode the machine is currently in (which we can't easily
		// determine).

		GenericWakeupFromIdle(pdx, TRUE);

		PPOWERINFO pip = (PPOWERINFO) Irp->AssociatedIrp.SystemBuffer;
		status = WritePowerInfoToRegistry(pdx, (PPOWERINFO) Irp->AssociatedIrp.SystemBuffer);
		if (NT_SUCCESS(status))
			status = ImplementPowerPolicy(pdx, (PPOWERINFO) Irp->AssociatedIrp.SystemBuffer);
		break;
		}						// IOCTL_SET_POWERINFO

	case IOCTL_IDLENOW:						// code == 0xFF2
		{						// IOCTL_IDLENOW
		status = GenericIdleDevice(pdx, PowerDeviceD3);
		break;
		}						// IOCTL_IDLENOW

	default:
		status = STATUS_INVALID_DEVICE_REQUEST;
		break;

		}						// process request

	Irp->IoStatus.Information = info;
	Irp->IoStatus.Status = status;
	return status;
	}							// GenericHandlePowerIoctl

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

GENERICAPI NTSTATUS GENERIC_EXPORT GenericIdleDevice(PGENERIC_EXTENSION pdx, DEVICE_POWER_STATE state, BOOLEAN wait /* = FALSE */)
	{							// GenericIdleDevice
	if (pdx->state != WORKING && pdx->state != STOPPED)
		return STATUS_SUCCESS;
	NTSTATUS status = SendDeviceSetPower(pdx, state, wait);
	if (status == STATUS_PENDING)
		status = STATUS_SUCCESS;
	return status;
	}							// GenericIdleDevice

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

GENERICAPI VOID GENERIC_EXPORT GenericMarkDeviceBusy(PGENERIC_EXTENSION pdx)
	{							// GenericMarkDeviceBusy
	if (pdx->idlecount)
		PoSetDeviceBusy(pdx->idlecount);
	}							// GenericMarkDeviceBusy

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

GENERICAPI VOID GENERIC_EXPORT GenericRegisterForIdleDetection(PGENERIC_EXTENSION pdx, ULONG ConservationTimeout,
	ULONG PerformanceTimeout, DEVICE_POWER_STATE state)
	{							// GenericRegisterForIdleDetection
	ASSERT(KeGetCurrentIrql() < DISPATCH_LEVEL);	// required by PoRegisterDeviceForIdleDetection
	pdx->cto = ConservationTimeout;
	pdx->pto = PerformanceTimeout;
	pdx->idlestate = state;
	pdx->idlecount = PoRegisterDeviceForIdleDetection(pdx->Pdo, ConservationTimeout, PerformanceTimeout, state);
	}							// GenericRegisterForIdleDetection

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

GENERICAPI VOID GENERIC_EXPORT GenericSaveRestoreComplete(PVOID context)
	{							// GenericSaveRestoreComplete
	HandlePowerEvent((PPOWCONTEXT) context, AsyncNotify);
	}							// GenericSaveRestoreComplete

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

GENERICAPI NTSTATUS GENERIC_EXPORT GenericWakeupControl(PGENERIC_EXTENSION pdx, enum WAKEFUNCTION wf)
	{							// GenericWakeupControl

	NTSTATUS status = STATUS_SUCCESS;

	switch (wf)
		{						// perform requested function

	case EnableWakeup:
		pdx->WakeupEnabled = TRUE;
		goto manage;

	case DisableWakeup:
		pdx->WakeupEnabled = FALSE;
		goto manage;

	case ManageWaitWake:
	manage:
		{						// ManageWaitWake
		if (pdx->WakeupEnabled)
			{					// wakeup is enabled

			// We're often called at StartDevice time before a capabilities query has happened.
			// Beginning in v. 2.01, we therefore attempt to get the capabilities now if we can.

			GenericGetDeviceCapabilities(pdx);	// do capabilities query if necessary and if at PASSIVE_LEVEL

			// If capabilities don't provide for wakeup, don't do anything. Also,
			// we mustn't issue a WAIT_WAKE if the device is not in D0.