From 82d77f22cde1cba1fdc1c2101b78f029341005d9 Mon Sep 17 00:00:00 2001
From: SnipUndercover <snipundercover@gmail.com>
Date: Sun, 2 Jul 2023 01:08:01 +0200
Subject: [PATCH] Add STDERR capturing in LogWriter

---
 Celeste.Mod.mm/Mod/Helpers/LogWriter.cs | 91 +++++++++++++++++--------
 Celeste.Mod.mm/Patches/Celeste.cs       | 17 +----
 MiniInstaller/Program.cs                | 10 +--
 3 files changed, 66 insertions(+), 52 deletions(-)

diff --git a/Celeste.Mod.mm/Mod/Helpers/LogWriter.cs b/Celeste.Mod.mm/Mod/Helpers/LogWriter.cs
index 4fc7e2f25..d53345838 100644
--- a/Celeste.Mod.mm/Mod/Helpers/LogWriter.cs
+++ b/Celeste.Mod.mm/Mod/Helpers/LogWriter.cs
@@ -1,53 +1,88 @@
-using System.IO;
+using System;
+using System.Collections.Generic;
+using System.IO;
 using System.Text;
 
 namespace Celeste.Mod.Helpers {
-    public class LogWriter : TextWriter {
+    public class LogWriter : IDisposable {
 
-        public TextWriter STDOUT;
+        public OutputStreamCapture STDOUT;
+        public OutputStreamCapture STDERR;
         public TextWriter File;
 
-        public override Encoding Encoding {
-            get {
-                return STDOUT?.Encoding ?? File?.Encoding;
-            }
+        private bool captureStarted = false;
+
+        public LogWriter(TextWriter @out, TextWriter err, TextWriter file, bool beginCapture = true) {
+            STDOUT = new OutputStreamCapture(@out, this);
+            STDERR = new OutputStreamCapture(err, this);
+            File = file;
+            if (beginCapture)
+                BeginCapture();
+        }
+
+        public Encoding Encoding => STDOUT.Encoding ?? STDERR.Encoding ?? File.Encoding;
+
+        public void BeginCapture() {
+            if (captureStarted)
+                return;
+            Console.SetOut(STDOUT);
+            Console.SetError(STDERR);
+            captureStarted = true;
+        }
+
+        public void EndCapture() {
+            if (!captureStarted)
+                return;
+            Console.SetOut(STDOUT.Stream);
+            Console.SetError(STDERR.Stream);
+            captureStarted = false;
         }
 
+        public void Dispose() {
+            EndCapture();
+            // TextWriter.Close has no effect if the stream is a console stream
+            STDOUT?.Close();
+            STDERR?.Close();
+            File?.Close();
+        }
+    }
+
+    public class OutputStreamCapture : TextWriter {
+        public TextWriter Stream;
+        public LogWriter Writer;
+
+        public OutputStreamCapture(TextWriter stream, LogWriter writer) {
+            Stream = stream;
+            Writer = writer;
+        }
+
+        public override Encoding Encoding => Stream.Encoding;
+
         public override void Write(string value) {
-            STDOUT?.Write(value);
-            File?.Write(value);
-            File?.Flush();
+            Stream.Write(value);
+            Writer.File.Write(value);
         }
 
         public override void WriteLine(string value) {
-            STDOUT?.WriteLine(value);
-            File?.WriteLine(value);
-            File?.Flush();
+            Stream.WriteLine(value);
+            Writer.File.WriteLine(value);
         }
 
         public override void Write(char value) {
-            STDOUT?.Write(value);
-            File?.Write(value);
-            File?.Flush();
+            Stream.Write(value);
+            Writer.File.Write(value);
         }
 
         public override void Write(char[] buffer, int index, int count) {
-            STDOUT?.Write(buffer, index, count);
-            File?.Write(buffer, index, count);
-            File?.Flush();
+            Stream.Write(buffer, index, count);
+            Writer.File.Write(buffer, index, count);
         }
 
         public override void Flush() {
-            STDOUT?.Flush();
-            File?.Flush();
-        }
-
-        public override void Close() {
-            STDOUT?.Close();
-            STDOUT = null;
-            File?.Close();
-            File = null;
+            Stream.Flush();
+            Writer.File.Flush();
         }
 
+        public override void Close() => Stream.Close();
     }
 }
diff --git a/Celeste.Mod.mm/Patches/Celeste.cs b/Celeste.Mod.mm/Patches/Celeste.cs
index 6e6171b80..7379768f8 100644
--- a/Celeste.Mod.mm/Patches/Celeste.cs
+++ b/Celeste.Mod.mm/Patches/Celeste.cs
@@ -172,21 +172,8 @@ public static void Main(string[] args) {
 
             using (Stream fileStream = new FileStream(logfile, FileMode.OpenOrCreate, FileAccess.Write, FileShare.ReadWrite | FileShare.Delete))
             using (StreamWriter fileWriter = new StreamWriter(fileStream, Console.OutputEncoding))
-            using (LogWriter logWriter = new LogWriter {
-                STDOUT = Console.Out,
-                File = fileWriter
-            }) {
-                try {
-                    Console.SetOut(logWriter);
-
-                    MainInner(args);
-                } finally {
-                    if (logWriter.STDOUT != null) {
-                        Console.SetOut(logWriter.STDOUT);
-                        logWriter.STDOUT = null;
-                    }
-                }
-            }
+            using (LogWriter logWriter = new LogWriter(Console.Out, Console.Error, fileWriter))
+                MainInner(args);
 
         }
 
diff --git a/MiniInstaller/Program.cs b/MiniInstaller/Program.cs
index a6334721c..f90db2bf8 100644
--- a/MiniInstaller/Program.cs
+++ b/MiniInstaller/Program.cs
@@ -84,12 +84,7 @@ public static int Main(string[] args) {
                 File.Delete(PathLog);
             using (Stream fileStream = File.OpenWrite(PathLog))
             using (StreamWriter fileWriter = new StreamWriter(fileStream, Console.OutputEncoding))
-            using (LogWriter logWriter = new LogWriter {
-                STDOUT = Console.Out,
-                File = fileWriter
-            }) {
-                Console.SetOut(logWriter);
-
+            using (LogWriter logWriter = new LogWriter(Console.Out, Console.Error, fileWriter)) {
                 try {
 
                     if (!IsMonoVersionCompatible()) {
@@ -148,9 +143,6 @@ public static int Main(string[] args) {
                     Environment.SetEnvironmentVariable("MONOMOD_MODS", "");
                     Environment.SetEnvironmentVariable("MONOMOD_DEPENDENCY_MISSING_THROW", "");
                 }
-
-                Console.SetOut(logWriter.STDOUT);
-                logWriter.STDOUT = null;
             }
 
             return 0;