Restore contents when a screen info is closed (microsoft#17853)

annihilatorrrr · Sep 6, 2024 · 4eb06fe · 4eb06fe
1 parent d2c3cfd
commit 4eb06fe
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 43 deletions.
diff --git a/src/host/VtIo.cpp b/src/host/VtIo.cpp
@@ -6,6 +6,7 @@
 
 #include <til/unicode.h>
 
+#include "directio.h"
 #include "handle.h" // LockConsole
 #include "output.h" // CloseConsoleProcessState
 #include "../interactivity/inc/ServiceLocator.hpp"
@@ -790,3 +791,51 @@ void VtIo::Writer::WriteInfos(til::point target, std::span<const CHAR_INFO> info
         } while (--repeat);
     }
 }
+
+void VtIo::Writer::WriteScreenInfo(SCREEN_INFORMATION& newContext, til::size oldSize) const
+{
+    const auto area = static_cast<size_t>(oldSize.width * oldSize.height);
+
+    auto& main = newContext.GetMainBuffer();
+    auto& alt = newContext.GetActiveBuffer();
+    const auto hasAltBuffer = &alt != &main;
+
+    // TODO GH#5094: This could use xterm's XTWINOPS "\e[8;<height>;<width>t" escape sequence here.
+    if (oldSize != main.GetBufferSize().Dimensions())
+    {
+        THROW_IF_NTSTATUS_FAILED(main.ResizeTraditional(oldSize));
+        main.SetViewportSize(&oldSize);
+    }
+    if (hasAltBuffer && oldSize != alt.GetBufferSize().Dimensions())
+    {
+        THROW_IF_NTSTATUS_FAILED(alt.ResizeTraditional(oldSize));
+        alt.SetViewportSize(&oldSize);
+    }
+
+    const auto request = Viewport::FromDimensions({}, oldSize);
+    Viewport read;
+    til::small_vector<CHAR_INFO, 1024> infos;
+    infos.resize(area, CHAR_INFO{ L' ', FOREGROUND_BLUE | FOREGROUND_GREEN | FOREGROUND_RED });
+
+    const auto dumpScreenInfo = [&](SCREEN_INFORMATION& screenInfo) {
+        THROW_IF_FAILED(ReadConsoleOutputWImplHelper(screenInfo, infos, request, read));
+        for (til::CoordType i = 0; i < oldSize.height; i++)
+        {
+            WriteInfos({ 0, i }, { infos.begin() + i * oldSize.width, static_cast<size_t>(oldSize.width) });
+        }
+
+        WriteCUP(screenInfo.GetTextBuffer().GetCursor().GetPosition());
+        WriteAttributes(screenInfo.GetAttributes());
+        WriteDECTCEM(screenInfo.GetTextBuffer().GetCursor().IsVisible());
+        WriteDECAWM(WI_IsFlagSet(screenInfo.OutputMode, ENABLE_WRAP_AT_EOL_OUTPUT));
+    };
+
+    WriteASB(false);
+    dumpScreenInfo(main);
+
+    if (hasAltBuffer)
+    {
+        WriteASB(true);
+        dumpScreenInfo(alt);
+    }
+}
diff --git a/src/host/VtIo.hpp b/src/host/VtIo.hpp
@@ -44,6 +44,7 @@ namespace Microsoft::Console::VirtualTerminal
             void WriteWindowTitle(std::wstring_view title) const;
             void WriteAttributes(const TextAttribute& attributes) const;
             void WriteInfos(til::point target, std::span<const CHAR_INFO> infos) const;
+            void WriteScreenInfo(SCREEN_INFORMATION& newContext, til::size oldSize) const;
 
         private:
             VtIo* _io = nullptr;

diff --git a/src/host/getset.cpp b/src/host/getset.cpp
@@ -487,49 +487,8 @@ void ApiRoutines::SetConsoleActiveScreenBufferImpl(SCREEN_INFORMATION& newContex
 
         if (auto writer = gci.GetVtWriter())
         {
-            const auto viewport = gci.GetActiveOutputBuffer().GetBufferSize();
-            const auto size = viewport.Dimensions();
-            const auto area = static_cast<size_t>(viewport.Width() * viewport.Height());
-
-            auto& main = newContext.GetMainBuffer();
-            auto& alt = newContext.GetActiveBuffer();
-            const auto hasAltBuffer = &alt != &main;
-
-            // TODO GH#5094: This could use xterm's XTWINOPS "\e[8;<height>;<width>t" escape sequence here.
-            THROW_IF_NTSTATUS_FAILED(main.ResizeTraditional(size));
-            main.SetViewportSize(&size);
-            if (hasAltBuffer)
-            {
-                THROW_IF_NTSTATUS_FAILED(alt.ResizeTraditional(size));
-                alt.SetViewportSize(&size);
-            }
-
-            Viewport read;
-            til::small_vector<CHAR_INFO, 1024> infos;
-            infos.resize(area, CHAR_INFO{ L' ', FOREGROUND_BLUE | FOREGROUND_GREEN | FOREGROUND_RED });
-
-            const auto dumpScreenInfo = [&](SCREEN_INFORMATION& screenInfo) {
-                THROW_IF_FAILED(ReadConsoleOutputWImpl(screenInfo, infos, viewport, read));
-                for (til::CoordType i = 0; i < size.height; i++)
-                {
-                    writer.WriteInfos({ 0, i }, { infos.begin() + i * size.width, static_cast<size_t>(size.width) });
-                }
-
-                writer.WriteCUP(screenInfo.GetTextBuffer().GetCursor().GetPosition());
-                writer.WriteAttributes(screenInfo.GetAttributes());
-                writer.WriteDECTCEM(screenInfo.GetTextBuffer().GetCursor().IsVisible());
-                writer.WriteDECAWM(WI_IsFlagSet(screenInfo.OutputMode, ENABLE_WRAP_AT_EOL_OUTPUT));
-            };
-
-            writer.WriteASB(false);
-            dumpScreenInfo(main);
-
-            if (hasAltBuffer)
-            {
-                writer.WriteASB(true);
-                dumpScreenInfo(alt);
-            }
-
+            const auto oldSize = gci.GetActiveOutputBuffer().GetBufferSize().Dimensions();
+            writer.WriteScreenInfo(newContext, oldSize);
             writer.Submit();
         }
 

diff --git a/src/server/ObjectHandle.cpp b/src/server/ObjectHandle.cpp
@@ -273,7 +273,18 @@ INPUT_READ_HANDLE_DATA* ConsoleHandleData::GetClientInput() const
     LOG_IF_FAILED(pScreenInfo->FreeIoHandle(this));
     if (!pScreenInfo->HasAnyOpenHandles())
     {
+        auto& gci = Microsoft::Console::Interactivity::ServiceLocator::LocateGlobals().getConsoleInformation();
+        const auto oldSize = gci.GetActiveOutputBuffer().GetBufferSize().Dimensions();
+        auto writer = gci.GetVtWriter();
+
         SCREEN_INFORMATION::s_RemoveScreenBuffer(pScreenInfo);
+
+        if (writer && gci.HasActiveOutputBuffer())
+        {
+            auto& newContext = gci.GetActiveOutputBuffer();
+            writer.WriteScreenInfo(newContext, oldSize);
+            writer.Submit();
+        }
     }
 
     return S_OK;