From 6aae5b48ea52edef6c8bec4864ba140587471b5e 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
Closes #475

Co-authored-by: Filippo Valsorda <hi@filippo.io>
---
 cmd/age/tui.go         | 12 ++++++++--
 cmd/age/tui_windows.go | 50 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 cmd/age/tui_windows.go

diff --git a/cmd/age/tui.go b/cmd/age/tui.go
index ac364f0..3247d6f 100644
--- a/cmd/age/tui.go
+++ b/cmd/age/tui.go
@@ -53,6 +53,11 @@ func errorWithHint(error string, hints ...string) {
 	os.Exit(1)
 }
 
+// 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) {
@@ -63,13 +68,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 0000000..2f4f8eb
--- /dev/null
+++ b/cmd/age/tui_windows.go
@@ -0,0 +1,50 @@
+// 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 (
+	"errors"
+	"os"
+	"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")
+
+	if err := withTerminal(func(in, out *os.File) error {
+		var mode uint32
+		if err := syscall.GetConsoleMode(syscall.Handle(out.Fd()), &mode); err != nil {
+			return err
+		}
+
+		mode |= ENABLE_PROCESSED_OUTPUT
+		mode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING
+
+		// If the SetConsoleMode function fails, the return value is zero.
+		// See https://learn.microsoft.com/en-us/windows/console/setconsolemode#return-value.
+		if ret, _, _ := setConsoleMode.Call(out.Fd(), uintptr(mode)); ret == 0 {
+			return errors.New("SetConsoleMode failed")
+		}
+
+		return nil
+	}); err != nil {
+		avoidTerminalEscapeSequences = true
+	}
+}