/**********************************************************************
 * 
 *  Toby Opferman
 *
 *  Example Driver
 *
 *  This example is for educational purposes only.  I license this source
 *  out for use in learning how to write a device driver.
 *
 *  Copyright (c) 2005, All Rights Reserved  
 **********************************************************************/


#define _X86_ 

#include <wdm.h>
#include <tdi.h>
#include <tdikrnl.h>
#include <netdrv.h>
#include "kmem.h"
#include "tdiexample.h"
#include "handleirp.h"
#include "tdifuncs.h"


typedef enum
{
    CS_NOT_CONNECTED,
    CS_CONNECTING,
    CS_CONNECTED,    
    CS_DISCONNECTING
                                           
} CONNECTION_STATUS;

typedef struct _TDI_EXAMPLE_CONTEXT
{
    TDI_HANDLE TdiHandle;
    CONNECTION_STATUS csConnectionState;
    PIRPLISTHEAD pReadIrpListHead;
    PIRPLISTHEAD pWriteIrpListHead;
    KMUTEX kConnectionLock;
    KEVENT kWriteIrpReady;
    KEVENT kWakeWriteIrpThread;
    KEVENT kInitEvent;
    NTSTATUS NtThreadStatus;
    BOOLEAN bWriteThreadAlive;       
    PFILE_OBJECT pWriteThread;

} TDI_EXAMPLE_CONTEXT, *PTDI_EXAMPLE_CONTEXT;



/**********************************************************************
 * Internal Functions
 **********************************************************************/
 
 VOID TdiExample_CancelRoutine(PDEVICE_OBJECT DeviceObject, PIRP pIrp); 
 VOID TdiExample_IrpCleanUp(PIRP pIrp, PVOID pContext);
 NTSTATUS TdiExample_Connect(PTDI_EXAMPLE_CONTEXT pTdiExampleContext, PVOID pAddressContext, UINT uiLength);
 NTSTATUS TdiExample_Disconnect(PTDI_EXAMPLE_CONTEXT pTdiExampleContext);
 VOID TdiExample_WorkItem(PDEVICE_OBJECT  DeviceObject, PVOID  Context);
 VOID TdiExample_NetworkWriteThread(PVOID StartContext);
 NTSTATUS TdiExample_ClientEventReceive(PVOID TdiEventContext, CONNECTION_CONTEXT ConnectionContext, ULONG ReceiveFlags, ULONG  BytesIndicated, ULONG  BytesAvailable, ULONG  *BytesTaken, PVOID  Tsdu, PIRP  *IoRequestPacket);

#pragma alloc_text(PAGE, TdiExample_NetworkWriteThread)
#pragma alloc_text(PAGE, TdiExample_WorkItem)
#pragma alloc_text(PAGE, TdiExample_Disconnect)
#pragma alloc_text(PAGE, TdiExample_IrpCleanUp)
/* #pragma alloc_text(PAGE, TdiExample_CancelRoutine) */
#pragma alloc_text(PAGE, TdiExample_IoControlInternal)
#pragma alloc_text(PAGE, TdiExample_Create) 
#pragma alloc_text(PAGE, TdiExample_Close) 
#pragma alloc_text(PAGE, TdiExample_IoControl) 
#pragma alloc_text(PAGE, TdiExample_Read)
#pragma alloc_text(PAGE, TdiExample_Write)
#pragma alloc_text(PAGE, TdiExample_UnSupportedFunction)
#pragma alloc_text(PAGE, TdiExample_Connect)
/* #pragma alloc_text(PAGE, TdiExample_ClientEventReceive) */
                                
#define STATUS_CONTINUE_COMPLETION  STATUS_SUCCESS
#define TDIEXAMPLE_POOL_TAG         ((ULONG)'EidT')
#define READ_IRPLIST_POOL_TAG       ((ULONG)'RprI')
#define WRITE_IRPLIST_POOL_TAG      ((ULONG)'WprI')


/**********************************************************************
 * 
 *  TdiExample_Create
 *
 *    This is called when an instance of this driver is created (CreateFile)
 *
 **********************************************************************/
NTSTATUS TdiExample_Create(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_INSUFFICIENT_RESOURCES;
    PIO_STACK_LOCATION pIoStackIrp = NULL;
    PTDI_EXAMPLE_CONTEXT pTdiExampleContext = NULL;

    DbgPrint("TdiExample_Create Called \r\n");


    pIoStackIrp = IoGetCurrentIrpStackLocation(Irp);

    /* 
     * We need to allocate our instance context for this handle.  This
     * data structure will be associated with this user mode handle.
     *
     * We are allocating it in Non-Paged Pool so we can use it while holding
     * Spin Locks in our IRP Queue.
     *
     */

    pTdiExampleContext = (PTDI_EXAMPLE_CONTEXT)KMem_AllocateNonPagedMemory(sizeof(TDI_EXAMPLE_CONTEXT), TDIEXAMPLE_POOL_TAG);

    if(pTdiExampleContext)
    {
         DbgPrint("TdiExample_Create Allocate = 0x%0x \r\n", pTdiExampleContext);
        /*
         * We have created a library called "HandleIrp" which helps to queue our Pending IRP's.
         * The first thing we do is create the context for this list.
         */
        pTdiExampleContext->pReadIrpListHead = HandleIrp_CreateIrpList(READ_IRPLIST_POOL_TAG);
        pTdiExampleContext->pWriteIrpListHead = HandleIrp_CreateIrpList(WRITE_IRPLIST_POOL_TAG);

        if(pTdiExampleContext->pWriteIrpListHead && pTdiExampleContext->pReadIrpListHead)
        {
            /*
             * We need to maintain a state for this handle which will help us to know
             * if a Read or Write request should be ignored or performed for example.
             *
             * We obviously start with not connected.
             */
            pTdiExampleContext->csConnectionState = CS_NOT_CONNECTED;
    
            /*
             * We have our own TDI Client Driver library which allows us to simply
             * call functions to initate connections, etc.  The first thing we need to do
             * is create a TdiHandle context which can then be used in other TDI Library Calls.
             */
            NtStatus = TdiFuncs_InitializeTransportHandles(&pTdiExampleContext->TdiHandle);
    
            if(NT_SUCCESS(NtStatus))
            {
                NtStatus = TdiFuncs_SetEventHandler(pTdiExampleContext->TdiHandle.pfoTransport, TDI_EVENT_RECEIVE, TdiExample_ClientEventReceive, (PVOID)pTdiExampleContext);
            }

            if(NT_SUCCESS(NtStatus))
            {
                PIO_WORKITEM pIoWorkItem;

                /*
                 * We need to initialize our locks and events for this handle.  
                 *
                 * The Connection Lock is used to Synchronize access to connecting and disconnecting
                 *
                 * The Write IRP Ready event is used to signal the Writer thread
                 * that an IRP has been queued and it's waiting to be sent.
                 *
                 * The Wake Up Write Irp Thread is used to signal the thread to wake up.  It is generally
                 * used to signal the thread to exit.
                 *
                 * kInitEvent is used to synchronize initialization and notify when the work item has
                 * signaled that the thread has been created.
                 *
                 * The bWriteThreadAlive variable is used to tell the thread it should exit. 
                 */
                KeInitializeMutex(&pTdiExampleContext->kConnectionLock, 0);
                
                KeInitializeEvent(&pTdiExampleContext->kWriteIrpReady, SynchronizationEvent, FALSE);
                KeInitializeEvent(&pTdiExampleContext->kWakeWriteIrpThread, SynchronizationEvent, FALSE); 
                KeInitializeEvent(&pTdiExampleContext->kInitEvent, NotificationEvent, FALSE); 
                    
                pTdiExampleContext->bWriteThreadAlive = TRUE;   
                
                /* 
                 * We want to create a SYSTEM thread so we can handle our Writes Asynchronously.
                 * The problem is that we can't call PsCreateSystemThread() outside of the SYSTEM
                 * process on Windows 2000.  On Windows XP/2003 we can by using the object attributes,
                 * however we want this driver to run on Windows 2000 as well.  So we need to create
                 * a work item which will process some work on our behalf in the system context.
                 *
                 * We can use this to be in the context of SYSTEM to create our thread.
                 *
                 * IO_WORKITEM is an opaque data structure used by the system.
                 *
                 */

                pIoWorkItem = IoAllocateWorkItem(DeviceObject);

                if(pIoWorkItem)
                {

                    IoQueueWorkItem(pIoWorkItem, TdiExample_WorkItem, DelayedWorkQueue, (PVOID)pTdiExampleContext);
                    
                    /* 
                     * Wait for the work item to complete creating the thread.
                     */
                    KeWaitForSingleObject(&pTdiExampleContext->kInitEvent, Executive, KernelMode, FALSE, NULL);
                    
                    NtStatus = pTdiExampleContext->NtThreadStatus;

                    /*
                     * It is safe to free the work item when the Work Item function is called.  The work
                     * item is dequeued before the function is called.  It is not safe to free a work item
                     * that is currently queued.
                     */
                    IoFreeWorkItem(pIoWorkItem);

                    if(NT_SUCCESS(NtStatus))
                    {
                        /*
                         * Set the context of our FILE_OBJECT
                         */
                        pIoStackIrp->FileObject->FsContext = (PVOID)pTdiExampleContext;
                    }
                    else
                    {
                        HandleIrp_FreeIrpList(pTdiExampleContext->pReadIrpListHead);
                        HandleIrp_FreeIrpList(pTdiExampleContext->pWriteIrpListHead);
                        DbgPrint("TdiExample_Create Free = 0x%0x \r\n", pTdiExampleContext);

                        KMem_FreeNonPagedMemory(pTdiExampleContext);
                    }
                }
                else
                {
                    HandleIrp_FreeIrpList(pTdiExampleContext->pReadIrpListHead);
                    HandleIrp_FreeIrpList(pTdiExampleContext->pWriteIrpListHead);
                    DbgPrint("TdiExample_Create Free = 0x%0x \r\n", pTdiExampleContext);

                    KMem_FreeNonPagedMemory(pTdiExampleContext);
                }

            }
            else
            {
                HandleIrp_FreeIrpList(pTdiExampleContext->pReadIrpListHead);
                HandleIrp_FreeIrpList(pTdiExampleContext->pWriteIrpListHead);
                DbgPrint("TdiExample_Create Free = 0x%0x \r\n", pTdiExampleContext);

                KMem_FreeNonPagedMemory(pTdiExampleContext);
            }
        }
        else
        {
            if(pTdiExampleContext->pWriteIrpListHead)
            {
                HandleIrp_FreeIrpList(pTdiExampleContext->pWriteIrpListHead);
            }

            if(pTdiExampleContext->pReadIrpListHead)
            {
                HandleIrp_FreeIrpList(pTdiExampleContext->pReadIrpListHead);
            }
            DbgPrint("TdiExample_Create Free = 0x%0x \r\n", pTdiExampleContext);

            KMem_FreeNonPagedMemory(pTdiExampleContext);

        }
    }

    Irp->IoStatus.Status      = NtStatus;
    Irp->IoStatus.Information = 0;

    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    DbgPrint("TdiExample_Create Exit 0x%0x \r\n", NtStatus);

    return NtStatus;
}

/**********************************************************************
 * 
 *  TdiExample_WorkItem
 *
 *    This is executed in the context of the SYSTEM process.
 *
 **********************************************************************/
VOID TdiExample_WorkItem(PDEVICE_OBJECT  DeviceObject, PVOID  Context)
{
   PTDI_EXAMPLE_CONTEXT pTdiExampleContext = (PVOID)Context;
   HANDLE WriteThreadHandle;

   DbgPrint("TdiExample_WorkItem Enter \r\n");
    /*
     * We want to allow asynchronous Reads and Writes.  To handle our Writes we will have our
     * own thread in the SYSTEM process space.  This thread will simply be woken up when we have
     * Write's on the Queue and perform the write.
     *
     * If a write is issued Asynchronous the I/O Manager will return to user mode and allow this
     * thread to continue.  If it is not, even though we return from our write and post to this
     * thread the I/O Manager will wait for it to finish, just as we are doing when we issue TDI
     * IRPs.
     */
    pTdiExampleContext->NtThreadStatus = PsCreateSystemThread(&WriteThreadHandle, 0, NULL, NULL, NULL, TdiExample_NetworkWriteThread, (PVOID)pTdiExampleContext);
    
    if(NT_SUCCESS(pTdiExampleContext->NtThreadStatus))
    {
        /*
         * We would rather use the PFILE_OBJECT since it is a pointer to the object rather than a handle.  Handles
         * are only good from within the context of a particular process.
         */
        ObReferenceObjectByHandle(WriteThreadHandle, GENERIC_READ | GENERIC_WRITE, NULL, KernelMode, (PVOID *)&pTdiExampleContext->pWriteThread, NULL);
        ZwClose(WriteThreadHandle);
    }

    KeSetEvent(&pTdiExampleContext->kInitEvent, IO_NO_INCREMENT, FALSE);
}


/**********************************************************************
 * 
 *  TdiExample_NetworkWriteThread
 *
 *    This is executed in the context of the SYSTEM process.