text: // Copyright (c) Microsoft Corporation.
  // Licensed under the MIT license.
  
  #include "precomp.h"
  #include "writeData.hpp"
  
  #include "_stream.h"
  #include "../types/inc/convert.hpp"
  
  #include "../interactivity/inc/ServiceLocator.hpp"
  
  // Routine Description:
  // - Creates a new write data object for used in servicing write console requests
  // Arguments:
  // - siContext - The output buffer to write text data to
  // - pwchContext - The string information that the client application sent us to be written
  // - cbContext - Byte count of the string above
  // - uiOutputCodepage - When the wait is completed, we *might* have to convert the byte count
  //                      back into a specific codepage if the initial call was an A call.
  //                      We need to remember what output codepage was set at the moment in time
  //                      when the write was delayed as it might change by the time it is serviced.
  // Return Value:
  // - THROW: Throws if space cannot be allocated to copy the given string
  WriteData::WriteData(SCREEN_INFORMATION& siContext,
                       _In_reads_bytes_(cbContext) wchar_t* const pwchContext,
                       const size_t cbContext,
                       const UINT uiOutputCodepage,
                       const bool requiresVtQuirk) :
      IWaitRoutine(ReplyDataType::Write),
      _siContext(siContext),
      _pwchContext(THROW_IF_NULL_ALLOC(reinterpret_cast<wchar_t*>(new byte[cbContext]))),
      _cbContext(cbContext),
      _uiOutputCodepage(uiOutputCodepage),
      _requiresVtQuirk(requiresVtQuirk),
      _fLeadByteCaptured(false),
      _fLeadByteConsumed(false),
      _cchUtf8Consumed(0)
  {
      memmove(_pwchContext, pwchContext, _cbContext);
  }
  
  // Routine Description:
  // - Destroys the write data object
  // - Frees the string copy we made on creation
  WriteData::~WriteData()
  {
      if (nullptr != _pwchContext)
      {
          delete[] _pwchContext;
      }
  }
  
  // Routine Description:
  // - Stores some additional information about lead byte adjustments from the conversion
  //   in WriteConsoleA before the real WriteConsole processing (always W) is reached
  //   so we can restore an accurate A byte count at the very end when the wait is serviced.
  // Arguments:
  // - fLeadByteCaptured - A lead byte was removed from the string before converted it and saved it.
  //                       We need to report to the original caller that we "wrote" the byte
  //                       even though it is held in escrow for the next call because it was
  //                       the last character in the stream.
  // - fLeadByteConsumed - We had a lead byte in escrow from the previous call that we stitched onto the
  //                       front of the input string even though the caller didn't write it in this call.
  //                       We need to report the byte count back to the caller without including this byte
  //                       in the calculation as it wasn't a part of what was given in this exact call.
  // Return Value:
  // - <none>
  void WriteData::SetLeadByteAdjustmentStatus(const bool fLeadByteCaptured,
                                              const bool fLeadByteConsumed)
  {
      _fLeadByteCaptured = fLeadByteCaptured;
      _fLeadByteConsumed = fLeadByteConsumed;
  }
  
  // Routine Description:
  // - For UTF-8 codepages, remembers how many bytes that the UTF-8 parser said it consumed from the input stream.
  //   This will allow us to give back the correct value after the wait routine Notify services the data later.
  // Arguments:
  // - cchUtf8Consumed - Count of characters consumed by the UTF-8 parser off the input stream to generate the
  //                     wide character string that is stowed in this object for consumption in the notify routine later.
  // Return Value:
  // - <none>
  void WriteData::SetUtf8ConsumedCharacters(const size_t cchUtf8Consumed)
  {
      _cchUtf8Consumed = cchUtf8Consumed;
  }
  
  // Routine Description:
  // - Called back at a later time to resume the writing operation when the output object becomes unblocked.
  // Arguments:
  // - TerminationReason - if this routine is called because a ctrl-c or ctrl-break was seen, this argument
  //                      contains CtrlC or CtrlBreak. If the owning thread is exiting, it will have ThreadDying. Otherwise 0.
  // - fIsUnicode - Input data was in UCS-2 unicode or it needs to be converted with the current Output Codepage
  // - pReplyStatus - The status code to return to the client application that originally called the API (before it was queued to wait)
  // - pNumBytes - The number of bytes of data that the server/driver will need to transmit back to the client process
  // - pControlKeyState - Unused for write operations. Set to 0.
  // - pOutputData - not used.
  // - true if the wait is done and result buffer/status code can be sent back to the client.
  // - false if we need to continue to wait because the output object blocked again
  bool WriteData::Notify(const WaitTerminationReason TerminationReason,
                         const bool fIsUnicode,
                         _Out_ NTSTATUS* const pReplyStatus,
                         _Out_ size_t* const pNumBytes,
                         _Out_ DWORD* const pControlKeyState,
                         _Out_ void* const /*pOutputData*/)
  {
      *pNumBytes = _cbContext;
      *pControlKeyState = 0;
  
      if (WI_IsFlagSet(TerminationReason, WaitTerminationReason::ThreadDying))
      {
          *pReplyStatus = STATUS_THREAD_IS_TERMINATING;
          return true;
      }
  
      // if we get to here, this routine was called by the input
      // thread, which grabs the current console lock.
  
      // This routine should be called by a thread owning the same lock on the
      // same console as we're reading from.
  
      FAIL_FAST_IF(!(Microsoft::Console::Interactivity::ServiceLocator::LocateGlobals().getConsoleInformation().IsConsoleLocked()));
  
      std::unique_ptr<WriteData> waiter;
      size_t cbContext = _cbContext;
      NTSTATUS Status = DoWriteConsole(_pwchContext,
                                       &cbContext,
                                       _siContext,
                                       _requiresVtQuirk,
                                       waiter);
  
      if (Status == CONSOLE_STATUS_WAIT)
      {
          // an extra waiter will be created by DoWriteConsole, but we're already a waiter so discard it.
          waiter.reset();
          return false;
      }
  
      // There's extra work to do to correct the byte counts if the original call was an A-version call.
      // We always process and hold text in the waiter as W-version text, but the A call is expecting
      // a byte value in its own codepage of how much we have written in that codepage.
      if (!fIsUnicode)
      {
          if (CP_UTF8 != _uiOutputCodepage)
          {
              // At this level with WriteConsole, everything is byte counts, so change back to char counts for
              // GetALengthFromW to work correctly.
              const size_t cchContext = cbContext / sizeof(wchar_t);
  
              // For non-UTF-8 codepages, we need to back convert the amount consumed and then
              // correlate that with any lead bytes we may have kept for later or reintroduced
              // from previous calls.
              size_t cchTextBufferRead = 0;
  
              // Start by counting the number of A bytes we used in printing our W string to the screen.
              try
              {
                  cchTextBufferRead = GetALengthFromW(_uiOutputCodepage, { _pwchContext, cchContext });
              }
              CATCH_LOG();
  
              // If we captured a byte off the string this time around up above, it means we didn't feed
              // it into the WriteConsoleW above, and therefore its consumption isn't accounted for
              // in the count we just made. Add +1 to compensate.
              if (_fLeadByteCaptured)
              {
                  cchTextBufferRead++;
              }
  
              // If we consumed an internally-stored lead byte this time around up above, it means that we
              // fed a byte into WriteConsoleW that wasn't a part of this particular call's request.
              // We need to -1 to compensate and tell the caller the right number of bytes consumed this request.
              if (_fLeadByteConsumed)
              {
                  cchTextBufferRead--;
              }
  
              cbContext = cchTextBufferRead;
          }
          else
          {
              // For UTF-8, we were told exactly how many valid bytes were consumed before we got into the wait state.
              // Just give that value back now.
              cbContext = _cchUtf8Consumed;
          }
      }
  
      *pNumBytes = cbContext;
      *pReplyStatus = Status;
      return true;
  }
  
  void WriteData::MigrateUserBuffersOnTransitionToBackgroundWait(const void* /*oldBuffer*/, void* /*newBuffer*/)
  {
      // WriteData does not hold API message buffers across a wait
  }
  