From dfcc8f3c6218b8d747bdd4bed6f7b6d1996bb376 Mon Sep 17 00:00:00 2001
From: Leonard Hecker <lhecker@microsoft.com>
Date: Tue, 29 Jul 2025 20:48:33 +0200
Subject: [PATCH] Fix use-after-free when disabling the ASB (#19186)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Closes #17515

## Validation Steps Performed
* Disable the ASB while there's a pending cooked read
* Type some text
* No crash ✅
---
 src/host/readDataCooked.cpp |  5 +++++
 src/host/readDataCooked.hpp |  1 +
 src/host/screenInfo.cpp     |  3 +++
 src/server/IWaitRoutine.h   |  4 ++++
 src/server/WaitBlock.cpp    |  5 +++++
 src/server/WaitBlock.h      |  1 +
 src/server/WaitQueue.cpp    | 16 ++++++++++++++++
 src/server/WaitQueue.h      |  2 ++
 8 files changed, 37 insertions(+)

diff --git a/src/host/readDataCooked.cpp b/src/host/readDataCooked.cpp
index c7d82fae0b..dd3d246645 100644
--- a/src/host/readDataCooked.cpp
+++ b/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.
diff --git a/src/host/readDataCooked.hpp b/src/host/readDataCooked.hpp
index 6132ab0c37..eada0fe8e7 100644
--- a/src/host/readDataCooked.hpp
+++ b/src/host/readDataCooked.hpp
@@ -19,6 +19,7 @@ public:
                      _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,
diff --git a/src/host/screenInfo.cpp b/src/host/screenInfo.cpp
index 3fe3b7b859..a2857eae55 100644
--- a/src/host/screenInfo.cpp
+++ b/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;
diff --git a/src/server/IWaitRoutine.h b/src/server/IWaitRoutine.h
index 59ffa905bd..b213c26fc4 100644
--- a/src/server/IWaitRoutine.h
+++ b/src/server/IWaitRoutine.h
@@ -42,6 +42,10 @@ public:
     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,
diff --git a/src/server/WaitBlock.cpp b/src/server/WaitBlock.cpp
index a871a52fc6..d1485f7068 100644
--- a/src/server/WaitBlock.cpp
+++ b/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
diff --git a/src/server/WaitBlock.h b/src/server/WaitBlock.h
index 560eaee91f..38e8f51ebd 100644
--- a/src/server/WaitBlock.h
+++ b/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,
diff --git a/src/server/WaitQueue.cpp b/src/server/WaitQueue.cpp
index b6231ea46f..d714c65011 100644
--- a/src/server/WaitQueue.cpp
+++ b/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:
diff --git a/src/server/WaitQueue.h b/src/server/WaitQueue.h
index be8a871ade..a23c7059c9 100644
--- a/src/server/WaitQueue.h
+++ b/src/server/WaitQueue.h
@@ -37,6 +37,8 @@ public:
     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);