How can I be notified about system state transition before they happen?

A PnP WDM driver gets notified about a system power transition when it receives a

IRP_MN_SET_POWER/SystemPowerState

PIRP. By the time you get this PIRP, the system power transition has already been
committed and some devices may have already been powered down.
Nearly 100% of the time, this is good enough, but if you require another device to
be powered on to flush data (where that other device is not a child of your device),
you need some other notification.

Alternatively, let's say that you are in a
non power pageable stack (the top device in the stack does not have the
DO_POWER_PAGABLE bit set in Flags) and you want
to write out a value to the registry during power down...but you can't do it. Why?
Because you cannot cause paging I/O once you recieve a
IRP_MN_SET_POWER/SystemPowerState PIRP if you are
a non power pagable stack (because the paging path might have already been powered
off and you would wait forever for the page in of the memory to complete). So, again, you must
use some other notification mechanism to know when to write out the value.

Enter the Ex

callback
, a not so well known object in the kernel. The Ex callback
object is a generic callback manages the registration of multiple listening callbacks and
allows for the calling into those listening callbacks at any time. The object
manages all the tricky problems (like handling a deregistration in the middle of
a callback being invoked) for you as well. In the system power case, the kernel
has created a public callback object,
\Callback\PowerState which is invoked before
the system power state transition has been committed. By registering for this
callback, you can perform your registry write or other actions in the callback
without worry. By the way, NT4 legacy style drivers can also register for this notification and do work before the system shuts down. Here is some sample code that shows you how to register and
implement the callback since it is not 100% intuitive.

Device extension

 
typdef struct _DEV_EXT {
    PVOID CbRegistration;
    PCALLBACK_OBJECT CbObject;
} DEV_EXT, *PDEV_EXT;

Registration

 
NTSTATUS
RegisterPowerStateChangeCallback(
    PDEV_EXT DevExt
    )
{
    OBJECT_ATTRIBUTES oa;
    UNICODE_STRING string;
    NTSTATUS status;

    RtlInitUnicodeString(&string, L"\\Callback\\PowerState");
    InitializeObjectAttributes(&oa,
                               &string,
                               OBJ_CASE_INSENSITIVE,
                               NULL,
                               NULL);

    //
    // Create a callback object, but we do not want to be the first ones
    // to create it (it should be created way before we load in NTOS
    // anyways) so we pass FALSE for the 3rd parameter.
    //
    status = ExCreateCallback(&DevExt->CbObject, &oa, FALSE, TRUE);

    if (NT_SUCCESS(status)) {
        DevExt->CbRegistration = ExRegisterCallback(DevExt->CbObject,
                                                    PowerStateCallback,
                                                    DevExt);
        if (DevExt->CbRegistration == NULL) {
            ObDereferenceObject(DevExt->CbObject);
            DevExt->CbObject = NULL;

            status = STATUS_UNSUCCESSFUL;
        }
    }

    return status;
}

Deregistration

 
VOID
DeregisterPowerStateChangeCallback(
    PDEV_EXT DevExt
    )
{
    if (DevExt->CbRegistration != NULL) {
        ExUnregisterCallback(DevExt->CbRegistration);
        DevExt->CbRegistration = NULL;
    }

    if (DevExt->CbObject != NULL) {
        ObDereferenceObject(DevExt->CbObject);
        DevExt->CbObject = NULL;
    }
}

Callback

 
VOID
PowerStateCallback(
    __in PVOID Context,
    __in PVOID Argument1,
    __in PVOID Argument2
    )
{
    PDEV_EXT pDevExt;

    pDevExt = (PDEV_EXT) Context;

    if (Argument1 != (PVOID) PO_CB_SYSTEM_STATE_LOCK) {
        return;
    }

    if (Argument2 == (PVOID) 0) {
        //
        // Exiting S0...
        //
    }
    else if (Argument2 == (PVOID) 1) {
        //
        // We have reentered S0 (note that we may have never left S0)
        //
    }
}