From 874814794bff65d365e8064bb3c8aed4dabc7861 Mon Sep 17 00:00:00 2001
From: bruvzg <7645683+bruvzg@users.noreply.github.com>
Date: Thu, 25 Apr 2024 12:56:07 +0300
Subject: [PATCH] [Windows] Improve console IO redirection.

---
 platform/windows/console_wrapper_windows.cpp |  4 +++
 platform/windows/os_windows.cpp              | 34 ++++++++++++++++----
 platform/windows/windows_terminal_logger.cpp | 22 +++----------
 3 files changed, 35 insertions(+), 25 deletions(-)

diff --git a/platform/windows/console_wrapper_windows.cpp b/platform/windows/console_wrapper_windows.cpp
index de751580b78..133711a9eaf 100644
--- a/platform/windows/console_wrapper_windows.cpp
+++ b/platform/windows/console_wrapper_windows.cpp
@@ -136,6 +136,10 @@ int main(int argc, char *argv[]) {
 	STARTUPINFOW si;
 	ZeroMemory(&si, sizeof(si));
 	si.cb = sizeof(si);
+	si.dwFlags = STARTF_USESTDHANDLES;
+	si.hStdInput = GetStdHandle(STD_INPUT_HANDLE);
+	si.hStdOutput = GetStdHandle(STD_OUTPUT_HANDLE);
+	si.hStdError = GetStdHandle(STD_ERROR_HANDLE);
 
 	WCHAR new_command_line[32767];
 	_snwprintf_s(new_command_line, 32767, _TRUNCATE, L"%ls %ls", exe_name, PathGetArgsW(GetCommandLineW()));
diff --git a/platform/windows/os_windows.cpp b/platform/windows/os_windows.cpp
index abed93d414b..157702655ea 100644
--- a/platform/windows/os_windows.cpp
+++ b/platform/windows/os_windows.cpp
@@ -115,7 +115,24 @@ void RedirectStream(const char *p_file_name, const char *p_mode, FILE *p_cpp_str
 }
 
 void RedirectIOToConsole() {
+	// Save current handles.
+	HANDLE h_stdin = GetStdHandle(STD_INPUT_HANDLE);
+	HANDLE h_stdout = GetStdHandle(STD_OUTPUT_HANDLE);
+	HANDLE h_stderr = GetStdHandle(STD_ERROR_HANDLE);
+
 	if (AttachConsole(ATTACH_PARENT_PROCESS)) {
+		// Restore redirection (Note: if not redirected it's NULL handles not INVALID_HANDLE_VALUE).
+		if (h_stdin != 0) {
+			SetStdHandle(STD_INPUT_HANDLE, h_stdin);
+		}
+		if (h_stdout != 0) {
+			SetStdHandle(STD_OUTPUT_HANDLE, h_stdout);
+		}
+		if (h_stderr != 0) {
+			SetStdHandle(STD_ERROR_HANDLE, h_stderr);
+		}
+
+		// Update file handles.
 		RedirectStream("CONIN$", "r", stdin, STD_INPUT_HANDLE);
 		RedirectStream("CONOUT$", "w", stdout, STD_OUTPUT_HANDLE);
 		RedirectStream("CONOUT$", "w", stderr, STD_ERROR_HANDLE);
@@ -173,10 +190,6 @@ void OS_Windows::initialize() {
 	add_error_handler(&error_handlers);
 #endif
 
-#ifndef WINDOWS_SUBSYSTEM_CONSOLE
-	RedirectIOToConsole();
-#endif
-
 	FileAccess::make_default<FileAccessWindows>(FileAccess::ACCESS_RESOURCES);
 	FileAccess::make_default<FileAccessWindows>(FileAccess::ACCESS_USERDATA);
 	FileAccess::make_default<FileAccessWindows>(FileAccess::ACCESS_FILESYSTEM);
@@ -1521,10 +1534,10 @@ void OS_Windows::unset_environment(const String &p_var) const {
 }
 
 String OS_Windows::get_stdin_string() {
-	WCHAR buff[1024];
+	char buff[1024];
 	DWORD count = 0;
-	if (ReadConsoleW(GetStdHandle(STD_INPUT_HANDLE), buff, 1024, &count, nullptr)) {
-		return String::utf16((const char16_t *)buff, count);
+	if (ReadFile(GetStdHandle(STD_INPUT_HANDLE), buff, 1024, &count, nullptr)) {
+		return String::utf8((const char *)buff, count);
 	}
 
 	return String();
@@ -1908,6 +1921,13 @@ String OS_Windows::get_system_ca_certificates() {
 OS_Windows::OS_Windows(HINSTANCE _hInstance) {
 	hInstance = _hInstance;
 
+#ifndef WINDOWS_SUBSYSTEM_CONSOLE
+	RedirectIOToConsole();
+#endif
+
+	SetConsoleOutputCP(CP_UTF8);
+	SetConsoleCP(CP_UTF8);
+
 	CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED);
 
 #ifdef WASAPI_ENABLED
diff --git a/platform/windows/windows_terminal_logger.cpp b/platform/windows/windows_terminal_logger.cpp
index 6881a755969..6fc33afe996 100644
--- a/platform/windows/windows_terminal_logger.cpp
+++ b/platform/windows/windows_terminal_logger.cpp
@@ -53,26 +53,12 @@ void WindowsTerminalLogger::logv(const char *p_format, va_list p_list, bool p_er
 	}
 	buf[len] = 0;
 
-	int wlen = MultiByteToWideChar(CP_UTF8, 0, buf, len, nullptr, 0);
-	if (wlen < 0) {
-		return;
-	}
-
-	wchar_t *wbuf = (wchar_t *)memalloc((len + 1) * sizeof(wchar_t));
-	ERR_FAIL_NULL_MSG(wbuf, "Out of memory.");
-	MultiByteToWideChar(CP_UTF8, 0, buf, len, wbuf, wlen);
-	wbuf[wlen] = 0;
-
-	if (p_err) {
-		fwprintf(stderr, L"%ls", wbuf);
-	} else {
-		wprintf(L"%ls", wbuf);
-	}
-
-	memfree(wbuf);
+	DWORD written = 0;
+	HANDLE h = p_err ? GetStdHandle(STD_ERROR_HANDLE) : GetStdHandle(STD_OUTPUT_HANDLE);
+	WriteFile(h, &buf[0], len, &written, nullptr);
 
 #ifdef DEBUG_ENABLED
-	fflush(stdout);
+	FlushFileBuffers(h);
 #endif
 }