// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

#include "precomp.h"

#include "WaitBlock.h"
#include "WaitQueue.h"

#include "ApiSorter.h"

#include "..\host\globals.h"
#include "..\host\utils.hpp"

#include "..\interactivity\inc\ServiceLocator.hpp"

// Routine Description:
// - Initializes a ConsoleWaitBlock
// - ConsoleWaitBlocks will self-manage their position in their two queues.
// - They will push themselves into the tail and store the iterator for constant deletion time later.
// Arguments:
// - pProcessQueue - The queue attached to the client process ID that requested this action
// - pObjectQueue - The queue attached to the console object that will service the action when data arrives
// - pWaitReplyMessage - The original API message related to the client process's service request
// - pWaiter - The context to return to later when the wait is satisfied.
ConsoleWaitBlock::ConsoleWaitBlock(_In_ ConsoleWaitQueue* const pProcessQueue,
                                   _In_ ConsoleWaitQueue* const pObjectQueue,
                                   const CONSOLE_API_MSG* const pWaitReplyMessage,
                                   _In_ IWaitRoutine* const pWaiter) :
    _pProcessQueue(THROW_HR_IF_NULL(E_INVALIDARG, pProcessQueue)),
    _pObjectQueue(THROW_HR_IF_NULL(E_INVALIDARG, pObjectQueue)),
    _pWaiter(THROW_HR_IF_NULL(E_INVALIDARG, pWaiter))
{
    _itProcessQueue = _pProcessQueue->_blocks.insert(_pProcessQueue->_blocks.end(), this);
    _itObjectQueue = _pObjectQueue->_blocks.insert(_pObjectQueue->_blocks.end(), this);

    _WaitReplyMessage = *pWaitReplyMessage;

    // We will write the original message back (with updated out parameters/payload) when the request is finally serviced.
    if (pWaitReplyMessage->Complete.Write.Data != nullptr)
    {
        _WaitReplyMessage.Complete.Write.Data = &_WaitReplyMessage.u;
    }
}

// Routine Description:
// - Destroys a ConsolewaitBlock
// - On deletion, ConsoleWaitBlocks will erase themselves from the process and object queues in
//   constant time with the iterator acquired on construction.
ConsoleWaitBlock::~ConsoleWaitBlock()
{
    _pProcessQueue->_blocks.erase(_itProcessQueue);
    _pObjectQueue->_blocks.erase(_itObjectQueue);

    if (_pWaiter != nullptr)
    {
        delete _pWaiter;
    }
}

// Routine Description:
// - Creates and enqueues a new wait for later callback when a routine cannot be serviced at this time.
// - Will extract the process ID and the target object, enqueuing in both to know when to callback
// Arguments:
// - pWaitReplyMessage - The original API message from the client asking for servicing
// - pWaiter - The context/callback information to restore and dispatch the call later.
// Return Value:
// - S_OK if queued and ready to go. Appropriate HRESULT value if it failed.
[[nodiscard]] HRESULT ConsoleWaitBlock::s_CreateWait(_Inout_ CONSOLE_API_MSG* const pWaitReplyMessage,
                                                     _In_ IWaitRoutine* const pWaiter)
{
    ConsoleProcessHandle* const ProcessData = pWaitReplyMessage->GetProcessHandle();
    FAIL_FAST_IF_NULL(ProcessData);

    ConsoleWaitQueue* const pProcessQueue = ProcessData->pWaitBlockQueue.get();

    ConsoleHandleData* const pHandleData = pWaitReplyMessage->GetObjectHandle();
    FAIL_FAST_IF_NULL(pHandleData);

    ConsoleWaitQueue* pObjectQueue = nullptr;
    LOG_IF_FAILED(pHandleData->GetWaitQueue(&pObjectQueue));
    FAIL_FAST_IF_NULL(pObjectQueue);

    ConsoleWaitBlock* pWaitBlock;
    try
    {
        pWaitBlock = new ConsoleWaitBlock(pProcessQueue,
                                          pObjectQueue,
                                          pWaitReplyMessage,
                                          pWaiter);
    }
    catch (...)
    {
        const HRESULT hr = wil::ResultFromCaughtException();
        pWaitReplyMessage->SetReplyStatus(NTSTATUS_FROM_HRESULT(hr));
        return hr;
    }

    return S_OK;
}

// Routine Description:
// - Used to trigger the callback routine inside this wait block.
// Arguments:
// - TerminationReason - A reason to tell the callback to terminate early or 0 if it should operate normally.
// Return Value:
// - True if the routine was able to successfully return data (or terminate). False otherwise.
bool ConsoleWaitBlock::Notify(const WaitTerminationReason TerminationReason)
{
    bool fRetVal;

    NTSTATUS status;
    size_t NumBytes = 0;
    DWORD dwControlKeyState;
    bool fIsUnicode = true;

    std::deque<std::unique_ptr<IInputEvent>> outEvents;
    // TODO: MSFT 14104228 - get rid of this void* and get the data
    // out of the read wait object properly.
    void* pOutputData = nullptr;
    // 1. Get unicode status of notify call based on message type.
    // We still need to know the Unicode status on reads as they will be converted after the wait operation.
    // Writes will have been converted before hitting the wait state.
    switch (_WaitReplyMessage.msgHeader.ApiNumber)
    {
    case API_NUMBER_GETCONSOLEINPUT:
    {
        CONSOLE_GETCONSOLEINPUT_MSG* a = &(_WaitReplyMessage.u.consoleMsgL1.GetConsoleInput);
        fIsUnicode = !!a->Unicode;
        pOutputData = &outEvents;
        break;
    }
    case API_NUMBER_READCONSOLE:
    {
        CONSOLE_READCONSOLE_MSG* a = &(_WaitReplyMessage.u.consoleMsgL1.ReadConsole);
        fIsUnicode = !!a->Unicode;
        break;
    }
    case API_NUMBER_WRITECONSOLE:
    {
        CONSOLE_WRITECONSOLE_MSG* a = &(_WaitReplyMessage.u.consoleMsgL1.WriteConsole);
        fIsUnicode = !!a->Unicode;
        break;
    }
    default:
    {
        FAIL_FAST_HR(E_NOTIMPL); // we shouldn't be getting a wait/notify on API numbers we don't support.
        break;
    }
    }

    // 2. If we have a waiter, dispatch to it.
    if (_pWaiter->Notify(TerminationReason, fIsUnicode, &status, &NumBytes, &dwControlKeyState, pOutputData))
    {
        // 3. If the wait was successful, set reply info and attach any additional return information that this request type might need.
        _WaitReplyMessage.SetReplyStatus(status);
        _WaitReplyMessage.SetReplyInformation(NumBytes);

        if (API_NUMBER_GETCONSOLEINPUT == _WaitReplyMessage.msgHeader.ApiNumber)
        {
            // ReadConsoleInput/PeekConsoleInput has this extra reply
            // information with the number of records, not number of
            // bytes.
            CONSOLE_GETCONSOLEINPUT_MSG* a = &(_WaitReplyMessage.u.consoleMsgL1.GetConsoleInput);

            void* buffer;
            ULONG cbBuffer;
            if (FAILED(_WaitReplyMessage.GetOutputBuffer(&buffer, &cbBuffer)))
            {
                return false;
            }

            INPUT_RECORD* const pRecordBuffer = static_cast<INPUT_RECORD* const>(buffer);
            a->NumRecords = static_cast<ULONG>(outEvents.size());
            for (size_t i = 0; i < a->NumRecords; ++i)
            {
                if (outEvents.empty())
                {
                    break;
                }
                pRecordBuffer[i] = outEvents.front()->ToInputRecord();
                outEvents.pop_front();
            }
        }
        else if (API_NUMBER_READCONSOLE == _WaitReplyMessage.msgHeader.ApiNumber)
        {
            // ReadConsole has this extra reply information with the control key state.
            CONSOLE_READCONSOLE_MSG* a = &(_WaitReplyMessage.u.consoleMsgL1.ReadConsole);
            a->ControlKeyState = dwControlKeyState;
            a->NumBytes = gsl::narrow<ULONG>(NumBytes);

            // - This routine is called when a ReadConsole or ReadFile request is about to be completed.
            // - It sets the number of bytes written as the information to be written with the completion status and,
            //   if CTRL+Z processing is enabled and a CTRL+Z is detected, switches the number of bytes read to zero.
            if (a->ProcessControlZ != FALSE &&
                a->NumBytes > 0 &&
                _WaitReplyMessage.State.OutputBuffer != nullptr &&
                *(PUCHAR)_WaitReplyMessage.State.OutputBuffer == 0x1a)
            {
                // On changing this, we also need to notify the Reply Information because it was stowed above into the reply packet.
                a->NumBytes = 0;
                // Setting the reply length to 0 and returning successfully from a blocked wait
                // will imply that the user has reached "End of File" on a raw read file stream.
                _WaitReplyMessage.SetReplyInformation(0);
            }
        }
        else if (API_NUMBER_WRITECONSOLE == _WaitReplyMessage.msgHeader.ApiNumber)
        {
            CONSOLE_WRITECONSOLE_MSG* a = &(_WaitReplyMessage.u.consoleMsgL1.WriteConsoleW);
            a->NumBytes = gsl::narrow<ULONG>(NumBytes);
        }

        LOG_IF_FAILED(_WaitReplyMessage.ReleaseMessageBuffers());

        LOG_IF_FAILED(Microsoft::Console::Interactivity::ServiceLocator::LocateGlobals().pDeviceComm->CompleteIo(&_WaitReplyMessage.Complete));

        fRetVal = true;
    }
    else
    {
        // If fThreadDying is TRUE we need to make sure that we removed the pWaitBlock from the list (which we don't do on this branch).
        FAIL_FAST_IF(!(WI_IsFlagClear(TerminationReason, WaitTerminationReason::ThreadDying)));
        fRetVal = false;
    }

    return fRetVal;
}