go-worm/lib/instance/console_windows.go

//go:build windows

package instance

import (
	"os"
	"syscall"
)

var (
	kernel32             = syscall.NewLazyDLL("kernel32.dll")
	procFreeConsole      = kernel32.NewProc("FreeConsole")
	procGetConsoleWindow = kernel32.NewProc("GetConsoleWindow")
	procSetConsoleCtrl   = kernel32.NewProc("SetConsoleCtrlHandler")
	procAllocConsole     = kernel32.NewProc("AllocConsole")
	procGetStdHandle     = kernel32.NewProc("GetStdHandle")
	user32               = syscall.NewLazyDLL("user32.dll")
	procShowWindow       = user32.NewProc("ShowWindow")
	consoleCtrlCallback  uintptr
)

const (
	swHide         = 0
	ctrlCEvent     = 0
	ctrlBreakEvent = 1
	ctrlCloseEvent = 2
)

func init() {
	// Must run before main: a console-subsystem build otherwise flashes a
	// terminal during flag.Parse, which looks like a broken automation step.
	ProtectFromConsoleClose()
	HideConsole()
}

// HideConsole hides any console window and detaches from it.
func HideConsole() {
	hwnd, _, _ := procGetConsoleWindow.Call()
	if hwnd != 0 {
		_, _, _ = procShowWindow.Call(hwnd, swHide)
	}
	_, _, _ = procFreeConsole.Call()
	redirectStdioToNul()
}

// ShowConsole allocates a console for -foreground when built as a GUI binary.
func ShowConsole() {
	_, _, _ = procAllocConsole.Call()
	hin, _, _ := procGetStdHandle.Call(^uintptr(9))   // STD_INPUT_HANDLE
	hout, _, _ := procGetStdHandle.Call(^uintptr(10)) // STD_OUTPUT_HANDLE
	herr, _, _ := procGetStdHandle.Call(^uintptr(11)) // STD_ERROR_HANDLE
	os.Stdin = os.NewFile(hin, "stdin")
	os.Stdout = os.NewFile(hout, "stdout")
	os.Stderr = os.NewFile(herr, "stderr")
}

// ProtectFromConsoleClose ignores Ctrl+C/Break/close so Task Scheduler and
// explorer do not kill the process when the flashed console goes away.
func ProtectFromConsoleClose() {
	if consoleCtrlCallback == 0 {
		consoleCtrlCallback = syscall.NewCallback(consoleCtrlHandler)
	}
	_, _, _ = procSetConsoleCtrl.Call(consoleCtrlCallback, 1)
}

// AllowConsoleKill restores default Ctrl+C handling for -foreground.
func AllowConsoleKill() {
	if consoleCtrlCallback == 0 {
		return
	}
	_, _, _ = procSetConsoleCtrl.Call(consoleCtrlCallback, 0)
}

func consoleCtrlHandler(ctrlType uintptr) uintptr {
	switch uint32(ctrlType) {
	case ctrlCEvent, ctrlBreakEvent, ctrlCloseEvent:
		return 1
	default:
		return 0
	}
}

func redirectStdioToNul() {
	nul, err := os.OpenFile("NUL", os.O_RDWR, 0)
	if err != nil {
		return
	}
	os.Stdin = nul
	os.Stdout = nul
	os.Stderr = nul
}