From 21e13c60c87dec9d00f1187cd01733416ddfde03 Mon Sep 17 00:00:00 2001
From: Twan van Laarhoven <twanvl@gmail.com>
Date: Mon, 20 Apr 2020 23:27:00 +0200
Subject: [PATCH] Fix: win32 cli wrapper: restore original console attributes

---
 src/cli/win32_cli_wrapper.cpp | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/cli/win32_cli_wrapper.cpp b/src/cli/win32_cli_wrapper.cpp
index f776b6ee..687251c5 100644
--- a/src/cli/win32_cli_wrapper.cpp
+++ b/src/cli/win32_cli_wrapper.cpp
@@ -45,6 +45,7 @@ DWORD WINAPI TransferThread(Transfer*);
 BOOL WINAPI HandleCtrlEvent(DWORD type);
 void CopyFileBuffer(HANDLE output, char* buffer, DWORD size);
 void InitEscapeTranslation(HANDLE console);
+void FinishEscapeTranslation(HANDLE console);
 void PerformEscapeCode(HANDLE console, char command, int argc, int argv[]);
 
 /// The child process
@@ -54,6 +55,7 @@ PROCESS_INFORMATION child_process_info;
 HANDLE in_mine,  in_theirs,  in_real;
 HANDLE out_mine, out_theirs, out_real;
 HANDLE err_mine, err_theirs, err_real;
+bool need_redirection;
 
 // ----------------------------------------------------------------------------- : Main function
 
@@ -61,7 +63,7 @@ const char* redirect_flags[] = {"-?","--help","-v","--version","--cli","-c","--e
 
 int main(int argc, char** argv) {
   // determine whether we need to wrap console i/o
-  bool need_redirection = false;
+  need_redirection = false;
   for (int i = 1 ; i < argc ; ++i) {
     for (int j = 0 ; j < sizeof(redirect_flags)/sizeof(redirect_flags[0]) ; ++j) {
       if (strcmp(argv[i],redirect_flags[j]) == 0) {
@@ -116,6 +118,7 @@ int main(int argc, char** argv) {
     out_real = GetStdHandle(STD_OUTPUT_HANDLE);
     err_real = GetStdHandle(STD_ERROR_HANDLE);
     InitEscapeTranslation(out_real);
+    InitEscapeTranslation(err_real);
     
     // start threads
     Transfer tranfer_in  = {in_real,  in_mine,  false};
@@ -135,6 +138,7 @@ int main(int argc, char** argv) {
   
   // start the child program
   if (!CreateProcess(app_path,command_line,NULL,NULL,TRUE,0,NULL,NULL,&child_startup_info,&child_process_info)) {
+    fprintf(stderr, "Unable to start child process.\n");
     ExitProcess(1);
   }
   
@@ -147,6 +151,10 @@ int main(int argc, char** argv) {
   }
   
   // That's all folks!
+  if (need_redirection) {
+    FinishEscapeTranslation(out_real);
+    FinishEscapeTranslation(err_real);
+  }
   return exit_code;
 }
 
@@ -253,6 +261,12 @@ void InitEscapeTranslation(HANDLE console) {
   original_attributes = screen_buffer.wAttributes;
 }
 
+// Cleanup for escape translation: put things back to normal
+void FinishEscapeTranslation(HANDLE console) {
+  CONSOLE_SCREEN_BUFFER_INFO screen_buffer;
+  SetConsoleTextAttribute(console, original_attributes);
+}
+
 /// Perform an escape code
 void PerformEscapeCode(HANDLE console, char command, int argc, int argv[]) {
   switch (command) {