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 <[email protected]>

cmd/derper/depaware.txt

@@ -13,6 +13,7 @@ tailscale.com/cmd/derper dependencies: (generated by github.com/tailscale/depawa
         github.com/beorn7/perks/quantile                             from github.com/prometheus/client_golang/prometheus
      💣 github.com/cespare/xxhash/v2                                 from github.com/prometheus/client_golang/prometheus
    L    github.com/coreos/go-iptables/iptables                       from tailscale.com/util/linuxfw
+   W 💣 github.com/dblohm7/wingoes                                   from tailscale.com/util/winutil
         github.com/fxamacker/cbor/v2                                 from tailscale.com/tka
         github.com/golang/groupcache/lru                             from tailscale.com/net/dnscache
         github.com/golang/protobuf/proto                             from github.com/matttproud/golang_protobuf_extensions/pbutil+

cmd/tailscale/depaware.txt

@@ -11,7 +11,7 @@ tailscale.com/cmd/tailscale dependencies: (generated by github.com/tailscale/dep
    W    github.com/alexbrainman/sspi/internal/common                 from github.com/alexbrainman/sspi/negotiate
    W 💣 github.com/alexbrainman/sspi/negotiate                       from tailscale.com/net/tshttpproxy
    L    github.com/coreos/go-iptables/iptables                       from tailscale.com/util/linuxfw
-   W 💣 github.com/dblohm7/wingoes                                   from tailscale.com/util/winutil/authenticode
+   W 💣 github.com/dblohm7/wingoes                                   from tailscale.com/util/winutil/authenticode+
    W 💣 github.com/dblohm7/wingoes/pe                                from tailscale.com/util/winutil/authenticode
         github.com/fxamacker/cbor/v2                                 from tailscale.com/tka
         github.com/golang/groupcache/lru                             from tailscale.com/net/dnscache

go.mod

@@ -18,7 +18,7 @@ require (
 	github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf
 	github.com/creack/pty v1.1.18
 	github.com/dave/jennifer v1.6.1
-	github.com/dblohm7/wingoes v0.0.0-20230803162905-5c6286bb8c6e
+	github.com/dblohm7/wingoes v0.0.0-20230821191801-fc76608aecf0
 	github.com/dsnet/try v0.0.3
 	github.com/evanw/esbuild v0.14.53
 	github.com/frankban/quicktest v1.14.5

go.sum

@@ -222,8 +222,8 @@ github.com/dave/jennifer v1.6.1/go.mod h1:nXbxhEmQfOZhWml3D1cDK5M1FLnMSozpbFN/m3
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/dblohm7/wingoes v0.0.0-20230803162905-5c6286bb8c6e h1:tTRuQNnXKO6Ffu62nk9bnnPx/m+IyNMdFFfzsETyRO8=
-github.com/dblohm7/wingoes v0.0.0-20230803162905-5c6286bb8c6e/go.mod h1:6NCrWM5jRefaG7iN0iMShPalLsljHWBh9v1zxM2f8Xs=
+github.com/dblohm7/wingoes v0.0.0-20230821191801-fc76608aecf0 h1:/dgKwHVTI0J+A0zd/BHOF2CTn1deN0735cJrb+w2hbE=
+github.com/dblohm7/wingoes v0.0.0-20230821191801-fc76608aecf0/go.mod h1:6NCrWM5jRefaG7iN0iMShPalLsljHWBh9v1zxM2f8Xs=
 github.com/denis-tingaikin/go-header v0.4.3 h1:tEaZKAlqql6SKCY++utLmkPLd6K8IBM20Ha7UVm+mtU=
 github.com/denis-tingaikin/go-header v0.4.3/go.mod h1:0wOCWuN71D5qIgE2nz9KrKmuYBAC2Mra5RassOIQ2/c=
 github.com/docker/cli v23.0.5+incompatible h1:ufWmAOuD3Vmr7JP2G5K3cyuNC4YZWiAsuDEvFVVDafE=

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
+}