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>
This commit is contained in:
Aaron Klotz 2023-08-17 14:47:01 -06:00
parent 3a652d7761
commit ea693eacb6
9 changed files with 99 additions and 5 deletions

View File

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

View File

@ -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

2
go.mod
View File

@ -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

4
go.sum
View File

@ -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=

View File

@ -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

View File

@ -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)
}

View File

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

View File

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

View File

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