Properly implement WaitForData for ReadConsoleInput

src/host/ApiRoutines.h

@@ -55,6 +55,7 @@ class ApiRoutines : public IApiRoutines
                                               INPUT_READ_HANDLE_DATA& readHandleState,
                                               const bool IsUnicode,
                                               const bool IsPeek,
+                                              const bool IsWaitAllowed,
                                               std::unique_ptr<IWaitRoutine>& waiter) noexcept override;
 
     [[nodiscard]] HRESULT ReadConsoleImpl(IConsoleInputObject& context,

src/host/directio.cpp

@@ -45,6 +45,7 @@ using Microsoft::Console::Interactivity::ServiceLocator;
 // - ppWaiter - If we have to wait (not enough data to fill client
 // buffer), this contains context that will allow the server to
 // restore this call later.
+// - IsWaitAllowed - Whether an async read via CONSOLE_STATUS_WAIT is permitted.
 // Return Value:
 // - STATUS_SUCCESS - If data was found and ready for return to the client.
 // - CONSOLE_STATUS_WAIT - If we didn't have enough data or needed to
@@ -56,6 +57,7 @@ using Microsoft::Console::Interactivity::ServiceLocator;
                                                        INPUT_READ_HANDLE_DATA& readHandleState,
                                                        const bool IsUnicode,
                                                        const bool IsPeek,
+                                                       const bool IsWaitAllowed,
                                                        std::unique_ptr<IWaitRoutine>& waiter) noexcept
 {
     try
@@ -73,7 +75,7 @@ using Microsoft::Console::Interactivity::ServiceLocator;
         const auto Status = inputBuffer.Read(outEvents,
                                              eventReadCount,
                                              IsPeek,
-                                             true,
+                                             IsWaitAllowed,
                                              IsUnicode,
                                              false);
 

src/host/inputBuffer.cpp

@@ -492,8 +492,7 @@ size_t InputBuffer::Prepend(const std::span<const INPUT_RECORD>& inEvents)
             return STATUS_SUCCESS;
         }
 
-        _vtInputShouldSuppress = true;
-        auto resetVtInputSuppress = wil::scope_exit([&]() { _vtInputShouldSuppress = false; });
+        const auto wakeup = _wakeupReadersOnExit();
 
         // read all of the records out of the buffer, then write the
         // prepend ones, then write the original set. We need to do it
@@ -503,35 +502,15 @@ size_t InputBuffer::Prepend(const std::span<const INPUT_RECORD>& inEvents)
         std::deque<INPUT_RECORD> existingStorage;
         existingStorage.swap(_storage);
 
-        // We will need this variable to pass to _WriteBuffer so it can attempt to determine wait status.
-        // However, because we swapped the storage out from under it with an empty deque, it will always
-        // return true after the first one (as it is filling the newly emptied backing deque.)
-        // Then after the second one, because we've inserted some input, it will always say false.
-        auto unusedWaitStatus = false;
-
         // write the prepend records
         size_t prependEventsWritten;
-        _WriteBuffer(inEvents, prependEventsWritten, unusedWaitStatus);
-        FAIL_FAST_IF(!(unusedWaitStatus));
+        _WriteBuffer(inEvents, prependEventsWritten);
 
         for (const auto& event : existingStorage)
         {
             _storage.push_back(event);
         }
 
-        // We need to set the wait event if there were 0 events in the
-        // input queue when we started.
-        // Because we did interesting manipulation of the wait queue
-        // in order to prepend, we can't trust what _WriteBuffer said
-        // and instead need to set the event if the original backing
-        // buffer (the one we swapped out at the top) was empty
-        // when this whole thing started.
-        if (existingStorage.empty())
-        {
-            ServiceLocator::LocateGlobals().hInputEvent.SetEvent();
-        }
-        WakeUpReadersWaitingForData();
-
         return prependEventsWritten;
     }
     catch (...)
@@ -575,21 +554,9 @@ size_t InputBuffer::Write(const std::span<const INPUT_RECORD>& inEvents)
             return 0;
         }
 
-        _vtInputShouldSuppress = true;
-        auto resetVtInputSuppress = wil::scope_exit([&]() { _vtInputShouldSuppress = false; });
-
-        // Write to buffer.
+        const auto wakeup = _wakeupReadersOnExit();
         size_t EventsWritten;
-        bool SetWaitEvent;
-        _WriteBuffer(inEvents, EventsWritten, SetWaitEvent);
-
-        if (SetWaitEvent)
-        {
-            ServiceLocator::LocateGlobals().hInputEvent.SetEvent();
-        }
-
-        // Alert any writers waiting for space.
-        WakeUpReadersWaitingForData();
+        _WriteBuffer(inEvents, EventsWritten);
         return EventsWritten;
     }
     catch (...)
@@ -607,16 +574,8 @@ try
         return;
     }
 
-    const auto initiallyEmptyQueue = _storage.empty();
-
+    const auto wakeup = _wakeupReadersOnExit();
     _writeString(text);
-
-    if (initiallyEmptyQueue && !_storage.empty())
-    {
-        ServiceLocator::LocateGlobals().hInputEvent.SetEvent();
-    }
-
-    WakeUpReadersWaitingForData();
 }
 CATCH_LOG()
 
@@ -625,29 +584,26 @@ CATCH_LOG()
 // input buffer and the next application will suddenly get a "\x1b[I" sequence in their input. See GH#13238.
 void InputBuffer::WriteFocusEvent(bool focused) noexcept
 {
+    const auto wakeup = _wakeupReadersOnExit();
+
     if (IsInVirtualTerminalInputMode())
     {
         if (const auto out = _termInput.HandleFocus(focused))
         {
-            _HandleTerminalInputCallback(*out);
+            _writeString(*out);
         }
     }
     else
     {
-        // This is a mini-version of Write().
-        const auto wasEmpty = _storage.empty();
         _storage.push_back(SynthesizeFocusEvent(focused));
-        if (wasEmpty)
-        {
-            ServiceLocator::LocateGlobals().hInputEvent.SetEvent();
-        }
-        WakeUpReadersWaitingForData();
     }
 }
 
 // Returns true when mouse input started. You should then capture the mouse and produce further events.
 bool InputBuffer::WriteMouseEvent(til::point position, const unsigned int button, const short keyState, const short wheelDelta)
 {
+    const auto wakeup = _wakeupReadersOnExit();
+
     if (IsInVirtualTerminalInputMode())
     {
         // This magic flag is "documented" at https://msdn.microsoft.com/en-us/library/windows/desktop/ms646301(v=vs.85).aspx
@@ -669,7 +625,7 @@ bool InputBuffer::WriteMouseEvent(til::point position, const unsigned int button
 
         if (const auto out = _termInput.HandleMouse(position, button, keyState, wheelDelta, state))
         {
-            _HandleTerminalInputCallback(*out);
+            _writeString(*out);
             return true;
         }
     }
@@ -690,25 +646,38 @@ static bool IsPauseKey(const KEY_EVENT_RECORD& event)
     return ctrlButNotAlt && event.wVirtualKeyCode == L'S';
 }
 
+void InputBuffer::_wakeupReadersImpl(bool initiallyEmpty)
+{
+    if (!_storage.empty())
+    {
+        // It would be fine to call SetEvent() unconditionally,
+        // but technically we only need to ResetEvent() if the buffer is empty,
+        // and SetEvent() once it stopped being so, which is what this code does.
+        if (initiallyEmpty)
+        {
+            ServiceLocator::LocateGlobals().hInputEvent.SetEvent();
+        }
+
+        WakeUpReadersWaitingForData();
+    }
+}
+
 // Routine Description:
 // - Coalesces input events and transfers them to storage queue.
 // Arguments:
 // - inRecords - The events to store.
 // - eventsWritten - The number of events written since this function
 // was called.
-// - setWaitEvent - on exit, true if buffer became non-empty.
 // Return Value:
 // - None
 // Note:
 // - The console lock must be held when calling this routine.
 // - will throw on failure
-void InputBuffer::_WriteBuffer(const std::span<const INPUT_RECORD>& inEvents, _Out_ size_t& eventsWritten, _Out_ bool& setWaitEvent)
+void InputBuffer::_WriteBuffer(const std::span<const INPUT_RECORD>& inEvents, _Out_ size_t& eventsWritten)
 {
     auto& gci = ServiceLocator::LocateGlobals().getConsoleInformation();
 
     eventsWritten = 0;
-    setWaitEvent = false;
-    const auto initiallyEmptyQueue = _storage.empty();
     const auto initialInEventsSize = inEvents.size();
     const auto vtInputMode = IsInVirtualTerminalInputMode();
 
@@ -739,7 +708,7 @@ void InputBuffer::_WriteBuffer(const std::span<const INPUT_RECORD>& inEvents, _O
             // GH#11682: TerminalInput::HandleKey can handle both KeyEvents and Focus events seamlessly
             if (const auto out = _termInput.HandleKey(inEvent))
             {
-                _HandleTerminalInputCallback(*out);
+                _writeString(*out);
                 eventsWritten++;
                 continue;
             }
@@ -759,10 +728,6 @@ void InputBuffer::_WriteBuffer(const std::span<const INPUT_RECORD>& inEvents, _O
         _storage.push_back(inEvent);
         ++eventsWritten;
     }
-    if (initiallyEmptyQueue && !_storage.empty())
-    {
-        setWaitEvent = true;
-    }
 }
 
 // Routine Description::
@@ -826,36 +791,6 @@ bool InputBuffer::IsInVirtualTerminalInputMode() const
     return WI_IsFlagSet(InputMode, ENABLE_VIRTUAL_TERMINAL_INPUT);
 }
 
-// Routine Description:
-// - Handler for inserting key sequences into the buffer when the terminal emulation layer
-//   has determined a key can be converted appropriately into a sequence of inputs
-// Arguments:
-// - inEvents - Series of input records to insert into the buffer
-// Return Value:
-// - <none>
-void InputBuffer::_HandleTerminalInputCallback(const TerminalInput::StringType& text)
-{
-    try
-    {
-        if (text.empty())
-        {
-            return;
-        }
-
-        _writeString(text);
-
-        if (!_vtInputShouldSuppress)
-        {
-            ServiceLocator::LocateGlobals().hInputEvent.SetEvent();
-            WakeUpReadersWaitingForData();
-        }
-    }
-    catch (...)
-    {
-        LOG_HR(wil::ResultFromCaughtException());
-    }
-}
-
 void InputBuffer::_writeString(const std::wstring_view& text)
 {
     for (const auto& wch : text)

src/host/inputBuffer.hpp

@@ -84,17 +84,20 @@ class InputBuffer final : public ConsoleObjectHeader
     bool _writePartialByteSequenceAvailable = false;
     Microsoft::Console::VirtualTerminal::TerminalInput _termInput;
 
-    // This flag is used in _HandleTerminalInputCallback
-    // If the InputBuffer leads to a _HandleTerminalInputCallback call,
-    //    we should suppress the wakeup functions.
-    // Otherwise, we should be calling them.
-    bool _vtInputShouldSuppress{ false };
+    // Wakes up readers waiting for data to be in the input buffer.
+    auto _wakeupReadersOnExit() noexcept
+    {
+        const auto initiallyEmpty = _storage.empty();
+        return wil::scope_exit([this, initiallyEmpty]() {
+            _wakeupReadersImpl(initiallyEmpty);
+        });
+    }
 
+    void _wakeupReadersImpl(bool initiallyEmpty);
     void _switchReadingMode(ReadingMode mode);
     void _switchReadingModeSlowPath(ReadingMode mode);
-    void _WriteBuffer(const std::span<const INPUT_RECORD>& inRecords, _Out_ size_t& eventsWritten, _Out_ bool& setWaitEvent);
+    void _WriteBuffer(const std::span<const INPUT_RECORD>& inRecords, _Out_ size_t& eventsWritten);
     bool _CoalesceEvent(const INPUT_RECORD& inEvent) noexcept;
-    void _HandleTerminalInputCallback(const Microsoft::Console::VirtualTerminal::TerminalInput::StringType& text);
     void _writeString(const std::wstring_view& text);
 
 #ifdef UNIT_TESTING

src/host/readDataDirect.cpp

@@ -107,7 +107,7 @@ try
     *pReplyStatus = _pInputBuffer->Read(_outEvents,
                                         amountToRead,
                                         false,
-                                        false,
+                                        true,
                                         fIsUnicode,
                                         false);
 

src/host/ut_host/InputBufferTests.cpp

@@ -575,31 +575,6 @@ class InputBufferTests
                                            false));
     }
 
-    TEST_METHOD(WritingToEmptyBufferSignalsWaitEvent)
-    {
-        InputBuffer inputBuffer;
-        auto record = MakeKeyEvent(true, 1, L'a', 0, L'a', 0);
-        auto inputEvent = record;
-        size_t eventsWritten;
-        auto waitEvent = false;
-        inputBuffer.Flush();
-        // write one event to an empty buffer
-        InputEventQueue storage;
-        storage.push_back(std::move(inputEvent));
-        inputBuffer._WriteBuffer(storage, eventsWritten, waitEvent);
-        VERIFY_IS_TRUE(waitEvent);
-        // write another, it shouldn't signal this time
-        auto record2 = MakeKeyEvent(true, 1, L'b', 0, L'b', 0);
-        auto inputEvent2 = record2;
-        // write another event to a non-empty buffer
-        waitEvent = false;
-        storage.clear();
-        storage.push_back(std::move(inputEvent2));
-        inputBuffer._WriteBuffer(storage, eventsWritten, waitEvent);
-
-        VERIFY_IS_FALSE(waitEvent);
-    }
-
     TEST_METHOD(StreamReadingDeCoalesces)
     {
         InputBuffer inputBuffer;

src/server/ApiDispatchers.cpp

@@ -182,6 +182,7 @@ constexpr T saturate(auto val)
         *pInputReadHandleData,
         a->Unicode,
         fIsPeek,
+        fIsWaitAllowed,
         waiter);
 
     // We must return the number of records in the message payload (to alert the client)
@@ -191,30 +192,13 @@ constexpr T saturate(auto val)
     size_t cbWritten;
     LOG_IF_FAILED(SizeTMult(outEvents.size(), sizeof(INPUT_RECORD), &cbWritten));
 
-    if (nullptr != waiter.get())
+    if (waiter)
     {
-        // In some circumstances, the read may have told us to wait because it didn't have data,
-        // but the client explicitly asked us to return immediate. In that case, we'll convert the
-        // wait request into a "0 bytes found, OK".
-
-        if (fIsWaitAllowed)
-        {
-            hr = ConsoleWaitQueue::s_CreateWait(m, waiter.release());
-            if (SUCCEEDED(hr))
-            {
-                *pbReplyPending = TRUE;
-                hr = CONSOLE_STATUS_WAIT;
-            }
-        }
-        else
+        hr = ConsoleWaitQueue::s_CreateWait(m, waiter.release());
+        if (SUCCEEDED(hr))
         {
-            // If wait isn't allowed and the routine generated a
-            // waiter, say there was nothing to be
-            // retrieved right now.
-            // The waiter will be auto-freed in the smart pointer.
-
-            cbWritten = 0;
-            hr = S_OK;
+            *pbReplyPending = TRUE;
+            hr = CONSOLE_STATUS_WAIT;
         }
     }
     else

src/server/IApiRoutines.h

@@ -63,6 +63,7 @@ class __declspec(novtable) IApiRoutines
                                                       INPUT_READ_HANDLE_DATA& readHandleState,
                                                       const bool IsUnicode,
                                                       const bool IsPeek,
+                                                      const bool IsWaitAllowed,
                                                       std::unique_ptr<IWaitRoutine>& waiter) noexcept = 0;
 
     [[nodiscard]] virtual HRESULT ReadConsoleImpl(IConsoleInputObject& context,