From 4570c49e6ebd21c1fe0a7769ecd91e0fe618340a Mon Sep 17 00:00:00 2001
From: Simone Ragusa <hi@interrato.dev>
Date: Mon, 26 Dec 2022 23:16:07 +0100
Subject: [PATCH] cmd/age: fix terminal escape sequences on Windows

If possible, we enable virtual terminal processing, which is necessary
for using terminal escape sequences on instances of the Windows Console.
When enabling virtual terminal processing fails, we completely avoid
using escape sequences to prevent weird characters to be printed to the
console.

Fixes #474
---
 cmd/age/tui.go         | 12 +++++++++--
 cmd/age/tui_windows.go | 45 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)
 create mode 100644 cmd/age/tui_windows.go

diff --git a/cmd/age/tui.go b/cmd/age/tui.go
index 47a13604..1e5b6ee0 100644
--- a/cmd/age/tui.go
+++ b/cmd/age/tui.go
@@ -67,6 +67,11 @@ func exit(code int) {
 	os.Exit(code)
 }
 
+// avoidTerminalEscapeSequences is set if we need to avoid using escape
+// sequences to prevent weird characters being printed to the console. This will
+// happen on Windows when virtual terminal processing cannot be enabled.
+var avoidTerminalEscapeSequences bool
+
 // clearLine clears the current line on the terminal, or opens a new line if
 // terminal escape codes don't work.
 func clearLine(out io.Writer) {
@@ -77,13 +82,16 @@ func clearLine(out io.Writer) {
 	)
 
 	// First, open a new line, which is guaranteed to work everywhere. Then, try
-	// to erase the line above with escape codes.
+	// to erase the line above with escape codes, if possible.
 	//
 	// (We use CRLF instead of LF to work around an apparent bug in WSL2's
 	// handling of CONOUT$. Only when running a Windows binary from WSL2, the
 	// cursor would not go back to the start of the line with a simple LF.
 	// Honestly, it's impressive CONIN$ and CONOUT$ work at all inside WSL2.)
-	fmt.Fprintf(out, "\r\n"+CPL+EL)
+	fmt.Fprintf(out, "\r\n")
+	if !avoidTerminalEscapeSequences {
+		fmt.Fprintf(out, CPL+EL)
+	}
 }
 
 // withTerminal runs f with the terminal input and output files, if available.
diff --git a/cmd/age/tui_windows.go b/cmd/age/tui_windows.go
new file mode 100644
index 00000000..381d6c31
--- /dev/null
+++ b/cmd/age/tui_windows.go
@@ -0,0 +1,45 @@
+// Copyright 2022 The age Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"syscall"
+
+	"golang.org/x/sys/windows"
+)
+
+// Some instances of the Windows Console (e.g., cmd.exe and Windows PowerShell)
+// do not have the virtual terminal processing enabled, which is necessary to
+// make terminal escape sequences work. For this reason the clearLine function
+// may not properly work. Here we enable the virtual terminal processing, if
+// possible.
+//
+// See https://learn.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences.
+func init() {
+	const (
+		ENABLE_PROCESSED_OUTPUT            uint32 = 0x1
+		ENABLE_VIRTUAL_TERMINAL_PROCESSING uint32 = 0x4
+	)
+
+	kernel32DLL := windows.NewLazySystemDLL("Kernel32.dll")
+	setConsoleMode := kernel32DLL.NewProc("SetConsoleMode")
+
+	var mode uint32
+	err := syscall.GetConsoleMode(syscall.Stdout, &mode)
+	if err != nil {
+		// Terminal escape sequences may work at this point, but we can't know.
+		avoidTerminalEscapeSequences = true
+		return
+	}
+
+	mode |= ENABLE_PROCESSED_OUTPUT
+	mode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING
+	ret, _, _ := setConsoleMode.Call(uintptr(syscall.Stdout), uintptr(mode))
+	// If the SetConsoleMode function fails, the return value is zero.
+	// See https://learn.microsoft.com/en-us/windows/console/setconsolemode#return-value.
+	if ret == 0 {
+		avoidTerminalEscapeSequences = true
+	}
+}