util/winutil: add RegisterForRestart, allowing programs to indicate their preferences to the Windows restart manager

In order for the installer to restart the GUI correctly post-upgrade, we
need the GUI to be able to register its restart preferences.

This PR adds API support for doing so. I'm adding it to OSS so that it
is available should we need to do any such registrations on OSS binaries
in the future.

Updates https://github.com/tailscale/corp/issues/13998

Signed-off-by: Aaron Klotz <aaron@tailscale.com>

util/winutil/mksyscall.go

@@ -7,3 +7,4 @@ package winutil
 //go:generate go run golang.org/x/tools/cmd/goimports -w zsyscall_windows.go
 
 //sys queryServiceConfig2(hService windows.Handle, infoLevel uint32, buf *byte, bufLen uint32, bytesNeeded *uint32) (err error) [failretval==0] = advapi32.QueryServiceConfig2W
+//sys registerApplicationRestart(cmdLineExclExeName *uint16, flags uint32) (ret wingoes.HRESULT) = kernel32.RegisterApplicationRestart

util/winutil/winutil.go

@@ -75,3 +75,27 @@ func IsSIDValidPrincipal(uid string) bool {
 func LookupPseudoUser(uid string) (*user.User, error) {
 	return lookupPseudoUser(uid)
 }
+
+// RegisterForRestartOpts supplies options to RegisterForRestart.
+type RegisterForRestartOpts struct {
+	RestartOnCrash   bool     // When true, this program will be restarted after a crash.
+	RestartOnHang    bool     // When true, this program will be restarted after a hang.
+	RestartOnUpgrade bool     // When true, this program will be restarted after an upgrade.
+	RestartOnReboot  bool     // When true, this program will be restarted after a reboot.
+	UseCmdLineArgs   bool     // When true, CmdLineArgs will be used as the program's arguments upon restart. Otherwise no arguments will be provided.
+	CmdLineArgs      []string // When UseCmdLineArgs == true, contains the command line arguments, excluding the executable name itself. If nil or empty, the arguments from the current process will be re-used.
+}
+
+// RegisterForRestart registers the current process' restart preferences with
+// the Windows Restart Manager. This enables the OS to intelligently restart
+// the calling executable as requested via opts. This should be called by any
+// programs which need to be restarted by the installer post-update.
+//
+// This function may be called multiple times; the opts from the most recent
+// call will override those from any previous invocations.
+//
+// This function will only work on GOOS=windows. Trying to run it on any other
+// OS will always return nil.
+func RegisterForRestart(opts RegisterForRestartOpts) error {
+	return registerForRestart(opts)
+}

util/winutil/winutil_notwindows.go

@@ -28,3 +28,5 @@ func lookupPseudoUser(uid string) (*user.User, error) {
 }
 
 func IsCurrentProcessElevated() bool { return false }
+
+func registerForRestart(opts RegisterForRestartOpts) error { return nil }

util/winutil/winutil_windows.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"os"
 	"os/exec"
 	"os/user"
 	"runtime"
@@ -15,6 +16,7 @@ import (
 	"time"
 	"unsafe"
 
+	"github.com/dblohm7/wingoes"
 	"golang.org/x/sys/windows"
 	"golang.org/x/sys/windows/registry"
 )
@@ -551,3 +553,58 @@ func findHomeDirInRegistry(uid string) (dir string, err error) {
 	}
 	return dir, nil
 }
+
+const (
+	_RESTART_NO_CRASH  = 1
+	_RESTART_NO_HANG   = 2
+	_RESTART_NO_PATCH  = 4
+	_RESTART_NO_REBOOT = 8
+)
+
+func registerForRestart(opts RegisterForRestartOpts) error {
+	var flags uint32
+
+	if !opts.RestartOnCrash {
+		flags |= _RESTART_NO_CRASH
+	}
+	if !opts.RestartOnHang {
+		flags |= _RESTART_NO_HANG
+	}
+	if !opts.RestartOnUpgrade {
+		flags |= _RESTART_NO_PATCH
+	}
+	if !opts.RestartOnReboot {
+		flags |= _RESTART_NO_REBOOT
+	}
+
+	var cmdLine *uint16
+	if opts.UseCmdLineArgs {
+		if len(opts.CmdLineArgs) == 0 {
+			// re-use our current args, excluding the exe name itself
+			opts.CmdLineArgs = os.Args[1:]
+		}
+
+		var b strings.Builder
+		for _, arg := range opts.CmdLineArgs {
+			if b.Len() > 0 {
+				b.WriteByte(' ')
+			}
+			b.WriteString(windows.EscapeArg(arg))
+		}
+
+		if b.Len() > 0 {
+			var err error
+			cmdLine, err = windows.UTF16PtrFromString(b.String())
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	hr := registerApplicationRestart(cmdLine, flags)
+	if e := wingoes.ErrorFromHRESULT(hr); e.Failed() {
+		return e
+	}
+
+	return nil
+}

util/winutil/zsyscall_windows.go

@@ -6,6 +6,7 @@ import (
 	"syscall"
 	"unsafe"
 
+	"github.com/dblohm7/wingoes"
 	"golang.org/x/sys/windows"
 )
 
@@ -39,8 +40,10 @@ func errnoErr(e syscall.Errno) error {
 
 var (
 	modadvapi32 = windows.NewLazySystemDLL("advapi32.dll")
+	modkernel32 = windows.NewLazySystemDLL("kernel32.dll")
 
-	procQueryServiceConfig2W = modadvapi32.NewProc("QueryServiceConfig2W")
+	procQueryServiceConfig2W       = modadvapi32.NewProc("QueryServiceConfig2W")
+	procRegisterApplicationRestart = modkernel32.NewProc("RegisterApplicationRestart")
 )
 
 func queryServiceConfig2(hService windows.Handle, infoLevel uint32, buf *byte, bufLen uint32, bytesNeeded *uint32) (err error) {
@@ -50,3 +53,9 @@ func queryServiceConfig2(hService windows.Handle, infoLevel uint32, buf *byte, b
 	}
 	return
 }
+
+func registerApplicationRestart(cmdLineExclExeName *uint16, flags uint32) (ret wingoes.HRESULT) {
+	r0, _, _ := syscall.Syscall(procRegisterApplicationRestart.Addr(), 2, uintptr(unsafe.Pointer(cmdLineExclExeName)), uintptr(flags), 0)
+	ret = wingoes.HRESULT(r0)
+	return
+}