/*++
Copyright (c) Microsoft Corporation.  All rights reserved.

    THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF ANY
    KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
    IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A PARTICULAR
    PURPOSE.

Module Name:

    trivial.c

Abstract:   Derived from cancel.c (which demonstrates the new cancel-safe queue mechanism).
            Adapted from the original cancel sample (KB Q188276) available in MSDN.
--*/

#include <ntddk.h>

#include "..\inc\trivialioctl.h"

#define TRIVIAL_DEVICE_NAME_U     L"\\Device\\Trivial"
#define TRIVIAL_DOS_DEVICE_NAME_U L"\\DosDevices\\Trivial"

typedef struct _TRACE_BUFFER
{
    CHAR      Buffer[16384];
    SIZE_T    Length;
} TRACE_BUFFER, *PTRACE_BUFFER;

TRACE_BUFFER TraceBuffer;

typedef struct _DEVICE_EXTENSION{

   LARGE_INTEGER  RegistryCallbackCookie;
   KMUTEX         TraceBufferMutex;
    
}  DEVICE_EXTENSION, *PDEVICE_EXTENSION;

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

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

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

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

VOID
TrivialUnload(
    IN PDRIVER_OBJECT DriverObject
    );

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

void
TrivialpAcquire(PKMUTEX pKMutex);

void
TrivialpRelease(PKMUTEX pKMutex);

NTSTATUS TrivialRegistryCallback(
    IN PVOID CallbackContext,
    IN PVOID Argument1,
    IN PVOID Argument2
    );

#ifdef ALLOC_PRAGMA
#pragma alloc_text( INIT, DriverEntry )
#pragma alloc_text( PAGE, TrivialCreateClose)
#pragma alloc_text( PAGE, TrivialRead)
#pragma alloc_text( PAGE, TrivialDeviceControl)
#pragma alloc_text( PAGE, TrivialUnload)
#pragma alloc_text( PAGE, TrivialRegistryCallback)
#pragma alloc_text( PAGE, TrivialpAcquire)
#pragma alloc_text( PAGE, TrivialpRelease)

#endif // ALLOC_PRAGMA

NTSTATUS
DriverEntry(
    IN PDRIVER_OBJECT  DriverObject,
    IN PUNICODE_STRING RegistryPath
    )
/*++

Routine Description:

    Installable driver initialization entry point.
    This entry point is called directly by the I/O system.

Arguments:

    DriverObject - pointer to the driver object

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

Return Value:

    STATUS_SUCCESS if successful,
    STATUS_UNSUCCESSFUL otherwise

--*/
{
    NTSTATUS            status = STATUS_SUCCESS;
    UNICODE_STRING      unicodeDeviceName;   
    UNICODE_STRING      unicodeDosDeviceName;  
    PDEVICE_OBJECT      deviceObject;
    PDEVICE_EXTENSION   devExtension;
    HANDLE              threadHandle;

    UNREFERENCED_PARAMETER (RegistryPath);

    //
    // Initialize UNICODE_STRING structures with NT and DOS device names
    //
    (void) RtlInitUnicodeString( &unicodeDeviceName, TRIVIAL_DEVICE_NAME_U);
    (void) RtlInitUnicodeString( &unicodeDosDeviceName, TRIVIAL_DOS_DEVICE_NAME_U );

    //
    // Create DOS device
    //
    status = IoCreateDevice(
                DriverObject,
                sizeof(DEVICE_EXTENSION),
                &unicodeDeviceName,
                FILE_DEVICE_UNKNOWN,
                0,
                (BOOLEAN) FALSE,
                &deviceObject
                );

    if (!NT_SUCCESS(status))
    {
        return status;
    }

    //
    // Create the DOS name for the device
    //
    status = IoCreateSymbolicLink(
                (PUNICODE_STRING) &unicodeDosDeviceName,
                (PUNICODE_STRING) &unicodeDeviceName
                );

    if (!NT_SUCCESS(status))
    {
        IoDeleteDevice(deviceObject);
        return status;
    }

    // Initialize our data lock
    devExtension = deviceObject->DeviceExtension;

    KeInitializeMutex(&devExtension->TraceBufferMutex, 0);

    //
    // Initialize our IRP handling and unload functions
    //
    DriverObject->MajorFunction[IRP_MJ_CREATE]            = TrivialCreateClose;
    DriverObject->MajorFunction[IRP_MJ_READ]              = TrivialRead;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]    = TrivialDeviceControl;
    DriverObject->MajorFunction[IRP_MJ_CLOSE]             = TrivialCreateClose;
    DriverObject->MajorFunction[IRP_MJ_CLEANUP]           = TrivialCleanup;

    DriverObject->DriverUnload = TrivialUnload;

    // Arrange for I/O Manager to copy from the system buffer to the user buffer
    deviceObject->Flags |= DO_BUFFERED_IO;

    return status;
}



NTSTATUS
TrivialCreateClose(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    )
/*++

Routine Description:

   Process the Create and close IRPs sent to this device.

Arguments:

   DeviceObject - pointer to a device object.

   Irp - pointer to an I/O Request Packet.

Return Value:

      NT Status code

--*/
{
    PIO_STACK_LOCATION   irpStack;
    NTSTATUS             status = STATUS_SUCCESS;

    PDEVICE_EXTENSION    devExtension = DeviceObject->DeviceExtension;

    PAGED_CODE ();

    //
    // Get a pointer to the current location in the Irp.
    //

    irpStack = IoGetCurrentIrpStackLocation(Irp);

    switch(irpStack->MajorFunction)
    {
        case IRP_MJ_CREATE:
            // 
            // The dispatch routine for IRP_MJ_CREATE is called when a 
            // file object associated with the device is created. 
            // This is typically because of a call to CreateFile() in 
            // a user-mode program or because a another driver is 
            // layering itself over a this driver. A driver is 
            // required to supply a dispatch routine for IRP_MJ_CREATE.
            //
            Irp->IoStatus.Information = 0;

            //
            // Register for registry callbacks
            //
            status = CmRegisterCallback(TrivialRegistryCallback, devExtension, &devExtension->RegistryCallbackCookie);
            break;

        case IRP_MJ_CLOSE:
            //
            // The IRP_MJ_CLOSE dispatch routine is called when a file object
            // opened on the driver is being removed from the system; that is,
            // all file object handles have been closed and the reference count
            // of the file object is down to 0. 
            //
            Irp->IoStatus.Information = 0;
            break;

        default:
            status = STATUS_INVALID_PARAMETER;
            break;
    }

    //
    // Save Status for return and complete Irp
    //
    Irp->IoStatus.Status = status;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return status;
}

NTSTATUS
TrivialCleanup(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
)
/*++

Routine Description:
    This dispatch routine is called when the last handle (in
    the whole system) to a file object is closed. In other words, the open
    handle count for the file object goes to 0. A driver that holds pending
    IRPs internally must implement a routine for IRP_MJ_CLEANUP. When the
    routine is called, the driver should cancel all the pending IRPs that
    belong to the file object identified by the IRP_MJ_CLEANUP call. In other
    words, it should cancel all the IRPs that have the same file-object pointer
    as the one supplied in the current I/O stack location of the IRP for the
    IRP_MJ_CLEANUP call. Of course, IRPs belonging to other file objects should
    not be canceled. Also, if an outstanding IRP is completed immediately, the
    driver does not have to cancel it.

Arguments:

    DeviceObject     -- pointer to the device object
    Irp             -- pointer to the requesing Irp

Return Value:

    STATUS_SUCCESS   -- if the poll succeeded,
--*/
{

    PDEVICE_EXTENSION     devExtension;
    LIST_ENTRY             tempQueue;   
    PLIST_ENTRY            thisEntry;
    PIRP                   pendingIrp;
    PIO_STACK_LOCATION    pendingIrpStack, irpStack;

    //irpStack = IoGetCurrentIrpStackLocation(Irp);
    devExtension = DeviceObject->DeviceExtension;

    //
    // Unregister for registry callbacks
    //
    (void) CmUnRegisterCallback(devExtension->RegistryCallbackCookie);
    devExtension->RegistryCallbackCookie.QuadPart = 0;

    //
    // Finally complete the cleanup IRP
    //
    Irp->IoStatus.Information = 0;
    Irp->IoStatus.Status = STATUS_SUCCESS;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return STATUS_SUCCESS;
}

VOID
TrivialUnload(
    IN PDRIVER_OBJECT DriverObject
    )
/*++

Routine Description:

    Free all the allocated resources, etc.

Arguments:

    DriverObject - pointer to a driver object.

Return Value:

    VOID
--*/
{
    PDEVICE_OBJECT       deviceObject = DriverObject->DeviceObject;
    UNICODE_STRING       uniWin32NameString;
    PDEVICE_EXTENSION    devExtension = deviceObject->DeviceExtension;

    PAGED_CODE();

    //
    // Delete our DOS link and then our device
    //
    RtlInitUnicodeString( &uniWin32NameString, TRIVIAL_DOS_DEVICE_NAME_U );
    IoDeleteSymbolicLink( &uniWin32NameString );

    IoDeleteDevice( deviceObject );
 
    return;
}

void
TrivialpAcquire(PKMUTEX pKMutex)
{
    (void) KeWaitForSingleObject (pKMutex, Executive, KernelMode, FALSE, NULL);
}

void
TrivialpRelease(PKMUTEX pKMutex)
{
    KeReleaseMutex (pKMutex, FALSE);
}

NTSTATUS
TrivialRead (
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
 )
 /*++
     Routine Description:
  
           Read dispatch routine
           
     Arguments:
  
         DeviceObject - pointer to a device object.
         Irp          - pointer to current Irp
  
     Return Value:
  
         NT status code.
  
--*/
{
    NTSTATUS            status;
    PDEVICE_EXTENSION   devExtension;
    PIO_STACK_LOCATION  irpStack;

    PCHAR               iobuffer;

    PAGED_CODE();

    //
    // Get a pointer to the device extension and irp stack.
    //
    devExtension = DeviceObject->DeviceExtension;
    irpStack = IoGetCurrentIrpStackLocation(Irp);


    TrivialpAcquire(&devExtension->TraceBufferMutex);

    //
    // First make sure there is enough room in our buffer.
    //
    if (irpStack->Parameters.Read.Length < TraceBuffer.Length)
    {
        Irp->IoStatus.Status = status = STATUS_BUFFER_TOO_SMALL;
        Irp->IoStatus.Information  = 0;

    } else {
        iobuffer = (PCHAR)Irp->AssociatedIrp.SystemBuffer;
        RtlCopyMemory ( iobuffer, &TraceBuffer.Buffer, TraceBuffer.Length );

        Irp->IoStatus.Status = status = STATUS_SUCCESS;
        Irp->IoStatus.Information  = TraceBuffer.Length;

        TraceBuffer.Length = 0;
    }

    TrivialpRelease(&devExtension->TraceBufferMutex);
    IoCompleteRequest (Irp, IO_NO_INCREMENT);
    return status;
}

NTSTATUS
TrivialDeviceControl (
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
 )
 /*++
     Routine Description:
  
           DeviceControl dispatch routine
           
     Arguments:
  
         DeviceObject - pointer to a device object.
         Irp          - pointer to current Irp
  
     Return Value:
  
         NT status code.
  
--*/
{
    NTSTATUS            status;
    ULONG               information = 0;

    PDEVICE_EXTENSION   devExtension;
    PIO_STACK_LOCATION  irpStack;
    LARGE_INTEGER       currentTime;

    ULONG               inBufLength; // Input buffer length
    ULONG               outBufLength; // Output buffer length
    PCHAR               inBuf, outBuf; // pointer to Input and output buffer


    PAGED_CODE();

    //
    // Get a pointer to the device extension and irp stack.
    //
    devExtension = DeviceObject->DeviceExtension;
    irpStack = IoGetCurrentIrpStackLocation(Irp);

    //
    // We use METHOD_BUFFERED for all our IOCTLs in this driver
    //
    inBufLength = irpStack->Parameters.DeviceIoControl.InputBufferLength;
    outBufLength = irpStack->Parameters.DeviceIoControl.OutputBufferLength;
    inBuf = Irp->AssociatedIrp.SystemBuffer;
    outBuf = Irp->AssociatedIrp.SystemBuffer;


    //
    // Switch on the specified IOCTL code
    //
    switch ( irpStack->Parameters.DeviceIoControl.IoControlCode )
    {
    case IOCTL_TRIVIAL_GETLENGTH: 
        if (outBufLength == sizeof(TraceBuffer.Length)) {
            information = sizeof(TraceBuffer.Length);
            RtlCopyMemory (outBuf, &TraceBuffer.Length, information);
            status = STATUS_SUCCESS;
        }
        break;

    default:
        status = STATUS_INVALID_DEVICE_REQUEST;
        break;
    }

    Irp->IoStatus.Information  = information;
    Irp->IoStatus.Status = status;
    IoCompleteRequest (Irp, IO_NO_INCREMENT);
    return status;
}


void
TrivialpLogRegistryCallback ( PKMUTEX pKMutex, PWCHAR p, ULONG nbytes )
{
    PTRACE_BUFFER pBuffer       = &TraceBuffer;
    ULONG         maxLength = sizeof(TraceBuffer.Buffer);
    ULONG         nchars;
    ULONG         newLength;
    PCHAR         q;

    if (nbytes & 1) return;

    nchars = nbytes / sizeof(WCHAR);

    TrivialpAcquire(pKMutex);

    //
    // Copy if we have enough space.
    //
    newLength = TraceBuffer.Length + nchars + sizeof(CHAR);
    if (newLength <= maxLength) {

        q = TraceBuffer.Buffer + TraceBuffer.Length;
        while (nchars > 0) {
            *q++ = ( *p++ & 0xff );   // Unicode -> Ansi
            nchars--;
        }
        *q++ = '\n';

        TraceBuffer.Length = newLength;
    }

if (newLength < maxLength) TraceBuffer.Buffer[newLength] = '\0';    //////////////////// XXX:
    TrivialpRelease(pKMutex);
}

NTSTATUS TrivialRegistryCallback(
    IN PVOID CallbackContext,
    IN PVOID Argument1,
    IN PVOID Argument2
    )
{
    REG_NOTIFY_CLASS Type          = (REG_NOTIFY_CLASS) Argument1;
    PDEVICE_EXTENSION devExtension = (PDEVICE_EXTENSION) CallbackContext;

    switch( Type ) {
    case RegNtPreCreateKeyEx:
        {
            PREG_CREATE_KEY_INFORMATION pCreate = (PREG_CREATE_KEY_INFORMATION) Argument2;
            //
            // Code to handle Pre - CreateKey
            //
        }
        break;

    case RegNtPreDeleteKey:
        {
            PREG_DELETE_KEY_INFORMATION  pDelete = (PREG_DELETE_KEY_INFORMATION)Argument2;
            //
            // Code to handle NtDeleteKey
            //
        }
        break;

    case RegNtPreSetValueKey:
        {
            PREG_SET_VALUE_KEY_INFORMATION pSetValue = (PREG_SET_VALUE_KEY_INFORMATION)Argument2;
            //
            // Code to handle NtSetValueKey
            //
            TrivialpLogRegistryCallback(&devExtension->TraceBufferMutex, pSetValue->ValueName->Buffer, pSetValue->ValueName->Length);
        }
        break;

    case RegNtPreDeleteValueKey:
        {
            PREG_DELETE_VALUE_KEY_INFORMATION  pDeteteValue =  (PREG_DELETE_VALUE_KEY_INFORMATION)Argument2;
            //
            // Code to handle NtDeleteValueKey
            //
        }
        break;

    default:
        //
        // we are not interested to hook these APIs. Let the call pass through and return STATUS_SUCCESS 
        //
        break;
    }
    
    return STATUS_SUCCESS;
}