From ea693eacb6a1a8a5b60aef554d7d256542ad4952 Mon Sep 17 00:00:00 2001 From: Aaron Klotz Date: Thu, 17 Aug 2023 14:47:01 -0600 Subject: [PATCH] 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 --- cmd/derper/depaware.txt | 1 + cmd/tailscale/depaware.txt | 2 +- go.mod | 2 +- go.sum | 4 +-- util/winutil/mksyscall.go | 1 + util/winutil/winutil.go | 24 +++++++++++++ util/winutil/winutil_notwindows.go | 2 ++ util/winutil/winutil_windows.go | 57 ++++++++++++++++++++++++++++++ util/winutil/zsyscall_windows.go | 11 +++++- 9 files changed, 99 insertions(+), 5 deletions(-) diff --git a/cmd/derper/depaware.txt b/cmd/derper/depaware.txt index e105d4a83..2de22cf65 100644 --- a/cmd/derper/depaware.txt +++ b/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+ diff --git a/cmd/tailscale/depaware.txt b/cmd/tailscale/depaware.txt index 1c8260329..e8bbc5f5c 100644 --- a/cmd/tailscale/depaware.txt +++ b/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 diff --git a/go.mod b/go.mod index 24043a3e2..1b319b9e0 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 28c905d4a..bf88160a4 100644 --- a/go.sum +++ b/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= diff --git a/util/winutil/mksyscall.go b/util/winutil/mksyscall.go index c0a9b2082..3c5515ee0 100644 --- a/util/winutil/mksyscall.go +++ b/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 diff --git a/util/winutil/winutil.go b/util/winutil/winutil.go index 9ccc9e085..3ec3f7c99 100644 --- a/util/winutil/winutil.go +++ b/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) +} diff --git a/util/winutil/winutil_notwindows.go b/util/winutil/winutil_notwindows.go index 0f2c4a83a..c9a292aae 100644 --- a/util/winutil/winutil_notwindows.go +++ b/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 } diff --git a/util/winutil/winutil_windows.go b/util/winutil/winutil_windows.go index 89fc543db..ed516ce6b 100644 --- a/util/winutil/winutil_windows.go +++ b/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 +} diff --git a/util/winutil/zsyscall_windows.go b/util/winutil/zsyscall_windows.go index 8c899232f..77e9f36c8 100644 --- a/util/winutil/zsyscall_windows.go +++ b/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 +}