/*++

Copyright (c) 1999 - 2002  Microsoft Corporation

Module Name:

    context.c

Abstract:

    This is the stream nd stream handle context module of the kernel mode
    context sample filter driver


Environment:

    Kernel mode


--*/


#include "pch.h"


#ifdef ALLOC_PRAGMA
#pragma alloc_text(PAGE, CtxFindOrCreateFileContext)
#pragma alloc_text(PAGE, CtxCreateFileContext)
#pragma alloc_text(PAGE, CtxFindOrCreateStreamContext)
#pragma alloc_text(PAGE, CtxCreateStreamContext)
#pragma alloc_text(PAGE, CtxUpdateNameInStreamContext)
#pragma alloc_text(PAGE, CtxCreateOrReplaceStreamHandleContext)
#pragma alloc_text(PAGE, CtxCreateStreamHandleContext)
#pragma alloc_text(PAGE, CtxUpdateNameInStreamHandleContext)
#endif




NTSTATUS
CtxFindOrCreateFileContext (
    __in PFLT_CALLBACK_DATA Cbd,
    __in BOOLEAN CreateIfNotFound,
    __in_opt PUNICODE_STRING FileName,
    __deref_out PCTX_FILE_CONTEXT *FileContext,
    __out_opt PBOOLEAN ContextCreated
    )
/*++

Routine Description:

    This routine finds the file context for the target file.
    Optionally, if the context does not exist this routing creates
    a new one and attaches the context to the file.

Arguments:

    Cbd                   - Supplies a pointer to the callbackData which
                            declares the requested operation.
    CreateIfNotFound      - Supplies if the file context must be created if missing
    FileName              - Supplies the file name
    FileContext           - Returns the file context
    ContextCreated        - Returns if a new context was created

Return Value:

    Status

--*/
{
    NTSTATUS status;
    PCTX_FILE_CONTEXT fileContext;
    PCTX_FILE_CONTEXT oldFileContext;

    PAGED_CODE();

    *FileContext = NULL;
    if (ContextCreated != NULL) *ContextCreated = FALSE;

    //
    //  First try to get the file context.
    //

    DebugTrace( DEBUG_TRACE_FILE_CONTEXT_OPERATIONS,
                ("[Ctx]: Trying to get file context (FileObject = %p, Instance = %p)\n",
                 Cbd->Iopb->TargetFileObject,
                 Cbd->Iopb->TargetInstance) );

    status = FltGetFileContext( Cbd->Iopb->TargetInstance,
                                  Cbd->Iopb->TargetFileObject,
                                  &fileContext );

    //
    //  If the call failed because the context does not exist
    //  and the user wants to creat a new one, the create a
    //  new context
    //

    if (!NT_SUCCESS( status ) &&
        (status == STATUS_NOT_FOUND) &&
        CreateIfNotFound) {


        //
        //  Create a file context
        //

        DebugTrace( DEBUG_TRACE_FILE_CONTEXT_OPERATIONS,
                    ("[Ctx]: Creating file context (FileObject = %p, Instance = %p)\n",
                     Cbd->Iopb->TargetFileObject,
                     Cbd->Iopb->TargetInstance) );

        status = CtxCreateFileContext( FileName, &fileContext );

        if (!NT_SUCCESS( status )) {

            DebugTrace( DEBUG_TRACE_ERROR | DEBUG_TRACE_FILE_CONTEXT_OPERATIONS,
                        ("[Ctx]: Failed to create file context with status 0x%x. (FileObject = %p, Instance = %p)\n",
                        status,
                        Cbd->Iopb->TargetFileObject,
                        Cbd->Iopb->TargetInstance) );

            return status;
        }


        //
        //  Set the new context we just allocated on the file object
        //

        DebugTrace( DEBUG_TRACE_FILE_CONTEXT_OPERATIONS,
                    ("[Ctx]: Setting file context %p (FileObject = %p, Instance = %p)\n",
                     fileContext,
                     Cbd->Iopb->TargetFileObject,
                     Cbd->Iopb->TargetInstance) );

        status = FltSetFileContext( Cbd->Iopb->TargetInstance,
                                      Cbd->Iopb->TargetFileObject,
                                      FLT_SET_CONTEXT_KEEP_IF_EXISTS,
                                      fileContext,
                                      &oldFileContext );

        if (!NT_SUCCESS( status )) {

            DebugTrace( DEBUG_TRACE_FILE_CONTEXT_OPERATIONS,
                        ("[Ctx]: Failed to set file context with status 0x%x. (FileObject = %p, Instance = %p)\n",
                        status,
                        Cbd->Iopb->TargetFileObject,
                        Cbd->Iopb->TargetInstance) );
            //
            //  We release the context here because FltSetFileContext failed
            //
            //  If FltSetFileContext succeeded then the context will be returned
            //  to the caller. The caller will use the context and then release it
            //  when he is done with the context.
            //

            DebugTrace( DEBUG_TRACE_FILE_CONTEXT_OPERATIONS,
                        ("[Ctx]: Releasing file context %p (FileObject = %p, Instance = %p)\n",
                         fileContext,
                         Cbd->Iopb->TargetFileObject,
                         Cbd->Iopb->TargetInstance) );

            FltReleaseContext( fileContext );

            if (status != STATUS_FLT_CONTEXT_ALREADY_DEFINED) {

                //
                //  FltSetFileContext failed for a reason other than the context already
                //  existing on the file. So the object now does not have any context set
                //  on it. So we return failure to the caller.
                //

                DebugTrace( DEBUG_TRACE_ERROR | DEBUG_TRACE_FILE_CONTEXT_OPERATIONS,
                            ("[Ctx]: Failed to set file context with status 0x%x != STATUS_FLT_CONTEXT_ALREADY_DEFINED. (FileObject = %p, Instance = %p)\n",
                            status,
                            Cbd->Iopb->TargetFileObject,
                            Cbd->Iopb->TargetInstance) );

                return status;
            }

            //
            //  Race condition. Someone has set a context after we queried it.
            //  Use the already set context instead
            //

            DebugTrace( DEBUG_TRACE_FILE_CONTEXT_OPERATIONS,
                        ("[Ctx]: File context already defined. Retaining old file context %p (FileObject = %p, Instance = %p)\n",
                         oldFileContext,
                         Cbd->Iopb->TargetFileObject,
                         Cbd->Iopb->TargetInstance) );

            //
            //  Return the existing context. Note that the new context that we allocated has already been
            //  realeased above.
            //

            fileContext = oldFileContext;
            status = STATUS_SUCCESS;

        } else {

            if (ContextCreated != NULL) *ContextCreated = TRUE;
        }
    }

    *FileContext = fileContext;

    return status;
}


NTSTATUS
CtxCreateFileContext (
    __in PUNICODE_STRING FileName,
    __deref_out PCTX_FILE_CONTEXT *FileContext
    )
/*++

Routine Description:

    This routine creates a new file context

Arguments:

    FileName            - Supplies the file name
    FileContext         - Returns the file context

Return Value:

    Status

--*/
{
    NTSTATUS status;
    PCTX_FILE_CONTEXT fileContext;

    PAGED_CODE();

    //
    //  Allocate a file context
    //

    DebugTrace( DEBUG_TRACE_FILE_CONTEXT_OPERATIONS,
                ("[Ctx]: Allocating file context \n") );

    status = FltAllocateContext( Globals.Filter,
                                 FLT_FILE_CONTEXT,
                                 CTX_FILE_CONTEXT_SIZE,
                                 PagedPool,
                                 &fileContext );

    if (!NT_SUCCESS( status )) {

        DebugTrace( DEBUG_TRACE_FILE_CONTEXT_OPERATIONS | DEBUG_TRACE_ERROR,
                    ("[Ctx]: Failed to allocate file context with status 0x%x \n",
                     status) );
        return status;
    }

    //
    //  Initialize the newly created context
    //

    //
    //  Allocate and copy off the file name
    //

    fileContext->FileName.MaximumLength = FileName->Length;
    status = CtxAllocateUnicodeString( &fileContext->FileName );
    if (NT_SUCCESS( status )) {

        RtlCopyUnicodeString( &fileContext->FileName, FileName );
    }

    *FileContext = fileContext;

    return STATUS_SUCCESS;
}


NTSTATUS
CtxFindOrCreateStreamContext (
    __in PFLT_CALLBACK_DATA Cbd,
    __in BOOLEAN CreateIfNotFound,
    __deref_out PCTX_STREAM_CONTEXT *StreamContext,
    __out_opt PBOOLEAN ContextCreated
    )
/*++

Routine Description:

    This routine finds the stream context for the target stream.
    Optionally, if the context does not exist this routing creates
    a new one and attaches the context to the stream.

Arguments:

    Cbd                   - Supplies a pointer to the callbackData which
                            declares the requested operation.
    CreateIfNotFound      - Supplies if the stream must be created if missing
    StreamContext         - Returns the stream context
    ContextCreated        - Returns if a new context was created

Return Value:

    Status

--*/
{
    NTSTATUS status;
    PCTX_STREAM_CONTEXT streamContext;
    PCTX_STREAM_CONTEXT oldStreamContext;

    PAGED_CODE();

    *StreamContext = NULL;
    if (ContextCreated != NULL) *ContextCreated = FALSE;

    //
    //  First try to get the stream context.
    //

    DebugTrace( DEBUG_TRACE_STREAM_CONTEXT_OPERATIONS,
                ("[Ctx]: Trying to get stream context (FileObject = %p, Instance = %p)\n",
                 Cbd->Iopb->TargetFileObject,
                 Cbd->Iopb->TargetInstance) );

    status = FltGetStreamContext( Cbd->Iopb->TargetInstance,
                                  Cbd->Iopb->TargetFileObject,
                                  &streamContext );

    //
    //  If the call failed because the context does not exist
    //  and the user wants to creat a new one, the create a
    //  new context
    //

    if (!NT_SUCCESS( status ) &&
        (status == STATUS_NOT_FOUND) &&
        CreateIfNotFound) {


        //
        //  Create a stream context
        //

        DebugTrace( DEBUG_TRACE_STREAM_CONTEXT_OPERATIONS,
                    ("[Ctx]: Creating stream context (FileObject = %p, Instance = %p)\n",
                     Cbd->Iopb->TargetFileObject,
                     Cbd->Iopb->TargetInstance) );

        status = CtxCreateStreamContext( &streamContext );

        if (!NT_SUCCESS( status )) {

            DebugTrace( DEBUG_TRACE_ERROR | DEBUG_TRACE_STREAM_CONTEXT_OPERATIONS,
                        ("[Ctx]: Failed to create stream context with status 0x%x. (FileObject = %p, Instance = %p)\n",
                        status,
                        Cbd->Iopb->TargetFileObject,
                        Cbd->Iopb->TargetInstance) );

            return status;
        }


        //
        //  Set the new context we just allocated on the file object
        //

        DebugTrace( DEBUG_TRACE_STREAM_CONTEXT_OPERATIONS,
                    ("[Ctx]: Setting stream context %p (FileObject = %p, Instance = %p)\n",
                     streamContext,
                     Cbd->Iopb->TargetFileObject,
                     Cbd->Iopb->TargetInstance) );

        status = FltSetStreamContext( Cbd->Iopb->TargetInstance,
                                      Cbd->Iopb->TargetFileObject,
                                      FLT_SET_CONTEXT_KEEP_IF_EXISTS,
                                      streamContext,
                                      &oldStreamContext );

        if (!NT_SUCCESS( status )) {

            DebugTrace( DEBUG_TRACE_STREAM_CONTEXT_OPERATIONS,
                        ("[Ctx]: Failed to set stream context with status 0x%x. (FileObject = %p, Instance = %p)\n",
                        status,
                        Cbd->Iopb->TargetFileObject,
                        Cbd->Iopb->TargetInstance) );
            //
            //  We release the context here because FltSetStreamContext failed
            //
            //  If FltSetStreamContext succeeded then the context will be returned
            //  to the caller. The caller will use the context and then release it
            //  when he is done with the context.
            //

            DebugTrace( DEBUG_TRACE_STREAM_CONTEXT_OPERATIONS,
                        ("[Ctx]: Releasing stream context %p (FileObject = %p, Instance = %p)\n",
                         streamContext,
                         Cbd->Iopb->TargetFileObject,
                         Cbd->Iopb->TargetInstance) );

            FltReleaseContext( streamContext );

            if (status != STATUS_FLT_CONTEXT_ALREADY_DEFINED) {

                //
                //  FltSetStreamContext failed for a reason other than the context already
                //  existing on the stream. So the object now does not have any context set
                //  on it. So we return failure to the caller.
                //

                DebugTrace( DEBUG_TRACE_ERROR | DEBUG_TRACE_STREAM_CONTEXT_OPERATIONS,
                            ("[Ctx]: Failed to set stream context with status 0x%x != STATUS_FLT_CONTEXT_ALREADY_DEFINED. (FileObject = %p, Instance = %p)\n",
                            status,
                            Cbd->Iopb->TargetFileObject,
                            Cbd->Iopb->TargetInstance) );

                return status;
            }

            //
            //  Race condition. Someone has set a context after we queried it.
            //  Use the already set context instead
            //

            DebugTrace( DEBUG_TRACE_STREAM_CONTEXT_OPERATIONS,
                        ("[Ctx]: Stream context already defined. Retaining old stream context %p (FileObject = %p, Instance = %p)\n",
                         oldStreamContext,
                         Cbd->Iopb->TargetFileObject,
                         Cbd->Iopb->TargetInstance) );

            //
            //  Return the existing context. Note that the new context that we allocated has already been
            //  realeased above.
            //

            streamContext = oldStreamContext;
            status = STATUS_SUCCESS;

        } else {