Fix use-after-free when disabling the ASB (microsoft#19186)

Closes microsoft#17515

## Validation Steps Performed
* Disable the ASB while there's a pending cooked read
* Type some text
* No crash ✅

(cherry picked from commit dfcc8f3)
Service-Card-Id: PVTI_lADOAF3p4s4AxadtzgdFoBM
Service-Version: 1.23

Author: lhecker

src/host/readDataCooked.cpp

@@ -176,6 +176,11 @@ COOKED_READ_DATA::COOKED_READ_DATA(_In_ InputBuffer* const pInputBuffer,
     }
 }
 
+const SCREEN_INFORMATION* COOKED_READ_DATA::GetScreenBuffer() const noexcept
+{
+    return &_screenInfo;
+}
+
 // Routine Description:
 // - This routine is called to complete a cooked read that blocked in ReadInputBuffer.
 // - The context of the read was saved in the CookedReadData structure.

src/host/readDataCooked.hpp

@@ -19,6 +19,7 @@ class COOKED_READ_DATA final : public ReadData
                      _In_ std::wstring_view initialData,
                      _In_ ConsoleProcessHandle* pClientProcess);
 
+    const SCREEN_INFORMATION* GetScreenBuffer() const noexcept override;
     void MigrateUserBuffersOnTransitionToBackgroundWait(const void* oldBuffer, void* newBuffer) noexcept override;
 
     bool Notify(WaitTerminationReason TerminationReason,

src/host/screenInfo.cpp

@@ -192,6 +192,9 @@ void SCREEN_INFORMATION::s_InsertScreenBuffer(_In_ SCREEN_INFORMATION* const pSc
 void SCREEN_INFORMATION::s_RemoveScreenBuffer(_In_ SCREEN_INFORMATION* const pScreenInfo)
 {
     auto& gci = ServiceLocator::LocateGlobals().getConsoleInformation();
+
+    gci.GetActiveInputBuffer()->WaitQueue.CancelWaitersForScreenBuffer(pScreenInfo);
+
     if (pScreenInfo == gci.ScreenBuffers)
     {
         gci.ScreenBuffers = pScreenInfo->Next;

src/server/IWaitRoutine.h

@@ -42,6 +42,10 @@ class IWaitRoutine
     IWaitRoutine& operator=(const IWaitRoutine&) & = delete;
     IWaitRoutine& operator=(IWaitRoutine&&) & = delete;
 
+    virtual const SCREEN_INFORMATION* GetScreenBuffer() const noexcept
+    {
+        return nullptr;
+    }
     virtual void MigrateUserBuffersOnTransitionToBackgroundWait(const void* oldBuffer, void* newBuffer) = 0;
 
     virtual bool Notify(const WaitTerminationReason TerminationReason,

src/server/WaitBlock.cpp

@@ -75,6 +75,11 @@ ConsoleWaitBlock::~ConsoleWaitBlock()
     delete _pWaiter;
 }
 
+const SCREEN_INFORMATION* ConsoleWaitBlock::GetScreenBuffer() const noexcept
+{
+    return _pWaiter->GetScreenBuffer();
+}
+
 // 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

src/server/WaitBlock.h

@@ -30,6 +30,7 @@ class ConsoleWaitBlock
 public:
     ~ConsoleWaitBlock();
 
+    const SCREEN_INFORMATION* GetScreenBuffer() const noexcept;
     bool Notify(const WaitTerminationReason TerminationReason);
 
     [[nodiscard]] static HRESULT s_CreateWait(_Inout_ CONSOLE_API_MSG* const pWaitReplymessage,

src/server/WaitQueue.cpp

@@ -103,6 +103,22 @@ bool ConsoleWaitQueue::NotifyWaiters(const bool fNotifyAll,
     return fResult;
 }
 
+void ConsoleWaitQueue::CancelWaitersForScreenBuffer(const SCREEN_INFORMATION* pScreenInfo)
+{
+    for (auto it = _blocks.begin(); it != _blocks.end();)
+    {
+        const auto nextIt = std::next(it); // we have to capture next before it is potentially erased
+        auto* pWaitBlock = *it;
+
+        if (pWaitBlock->GetScreenBuffer() == pScreenInfo)
+        {
+            _NotifyBlock(pWaitBlock, WaitTerminationReason::HandleClosing);
+        }
+
+        it = nextIt;
+    }
+}
+
 // Routine Description:
 // - A helper to delete successfully notified callbacks
 // Arguments:

src/server/WaitQueue.h

@@ -37,6 +37,8 @@ class ConsoleWaitQueue
     bool NotifyWaiters(const bool fNotifyAll,
                        const WaitTerminationReason TerminationReason);
 
+    void CancelWaitersForScreenBuffer(const SCREEN_INFORMATION* pScreenInfo);
+
     [[nodiscard]] static HRESULT s_CreateWait(_Inout_ CONSOLE_API_MSG* const pWaitReplyMessage,
                                               _In_ IWaitRoutine* const pWaiter);