tailscale/util/winutil/subprocess_windows_test.go
Aaron Klotz db39a43f06 util/winutil: add support for restarting Windows processes in specific sessions
This PR is all about adding functionality that will enable the installer's
upgrade sequence to terminate processes belonging to the previous version,
and then subsequently restart instances belonging to the new version within
the session(s) corresponding to the processes that were killed.

There are multiple parts to this:

* We add support for the Restart Manager APIs, which allow us to query the
  OS for a list of processes locking specific files;
* We add the RestartableProcess and RestartableProcesses types that query
  additional information about the running processes that will allow us
  to correctly restart them in the future. These types also provide the
  ability to terminate the processes.
* We add the StartProcessInSession family of APIs that permit us to create
  new processes within specific sessions. This is needed in order to
  properly attach a new GUI process to the same RDP session and desktop that
  its previously-terminated counterpart would have been running in.
* I tweaked the winutil token APIs again.
* A lot of this stuff is pretty hard to test without a very elaborate
  harness, but I added a unit test for the most complicated part (though it
  requires LocalSystem to run).

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

Signed-off-by: Aaron Klotz <aaron@tailscale.com>
2023-11-30 14:04:27 -08:00

434 lines
10 KiB
Go

// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package winutil
import (
"bytes"
"errors"
"fmt"
"os"
"os/exec"
"path/filepath"
"runtime"
"strconv"
"strings"
"sync"
"testing"
"time"
)
// The code in this file is adapted from internal/testenv in the Go source tree
// and is used for writing tests that require spawning subprocesses.
var toRemove []string
func TestMain(m *testing.M) {
status := m.Run()
for _, file := range toRemove {
os.RemoveAll(file)
}
os.Exit(status)
}
var testprog struct {
sync.Mutex
dir string
target map[string]*buildexe
}
type buildexe struct {
once sync.Once
exe string
err error
}
func pathToTestProg(t *testing.T, binary string) string {
exe, err := buildTestProg(t, binary, "-buildvcs=false")
if err != nil {
t.Fatal(err)
}
return exe
}
func runTestProg(t *testing.T, binary, name string, env ...string) string {
exe, err := buildTestProg(t, binary, "-buildvcs=false")
if err != nil {
t.Fatal(err)
}
return runBuiltTestProg(t, exe, name, env...)
}
func startTestProg(t *testing.T, binary, name string, env ...string) {
exe, err := buildTestProg(t, binary, "-buildvcs=false")
if err != nil {
t.Fatal(err)
}
startBuiltTestProg(t, exe, name, env...)
}
func runBuiltTestProg(t *testing.T, exe, name string, env ...string) string {
cmd := exec.Command(exe, name)
cmd.Env = append(cmd.Env, env...)
if testing.Short() {
cmd.Env = append(cmd.Env, "RUNTIME_TEST_SHORT=1")
}
out, _ := runWithTimeout(t, cmd)
return string(out)
}
func startBuiltTestProg(t *testing.T, exe, name string, env ...string) {
cmd := exec.Command(exe, name)
cmd.Env = append(cmd.Env, env...)
if testing.Short() {
cmd.Env = append(cmd.Env, "RUNTIME_TEST_SHORT=1")
}
start(t, cmd)
}
var serializeBuild = make(chan bool, 2)
func buildTestProg(t *testing.T, binary string, flags ...string) (string, error) {
testprog.Lock()
if testprog.dir == "" {
dir, err := os.MkdirTemp("", "go-build")
if err != nil {
t.Fatalf("failed to create temp directory: %v", err)
}
testprog.dir = dir
toRemove = append(toRemove, dir)
}
if testprog.target == nil {
testprog.target = make(map[string]*buildexe)
}
name := binary
if len(flags) > 0 {
nameFlags := make([]string, 0, len(flags))
for _, flag := range flags {
nameFlags = append(nameFlags, strings.ReplaceAll(flag, "=", "_"))
}
name += "_" + strings.Join(nameFlags, "_")
}
target, ok := testprog.target[name]
if !ok {
target = &buildexe{}
testprog.target[name] = target
}
dir := testprog.dir
// Unlock testprog while actually building, so that other
// tests can look up executables that were already built.
testprog.Unlock()
target.once.Do(func() {
// Only do two "go build"'s at a time,
// to keep load from getting too high.
serializeBuild <- true
defer func() { <-serializeBuild }()
// Don't get confused if goToolPath calls t.Skip.
target.err = errors.New("building test called t.Skip")
exe := filepath.Join(dir, name+".exe")
t.Logf("running go build -o %s %s", exe, strings.Join(flags, " "))
cmd := exec.Command(goToolPath(t), append([]string{"build", "-o", exe}, flags...)...)
cmd.Dir = "testdata/" + binary
out, err := cmd.CombinedOutput()
if err != nil {
target.err = fmt.Errorf("building %s %v: %v\n%s", binary, flags, err, out)
} else {
target.exe = exe
target.err = nil
}
})
return target.exe, target.err
}
// goTool reports the path to the Go tool.
func goTool() (string, error) {
if !hasGoBuild() {
return "", errors.New("platform cannot run go tool")
}
exeSuffix := ".exe"
goroot, err := findGOROOT()
if err != nil {
return "", fmt.Errorf("cannot find go tool: %w", err)
}
path := filepath.Join(goroot, "bin", "go"+exeSuffix)
if _, err := os.Stat(path); err == nil {
return path, nil
}
goBin, err := exec.LookPath("go" + exeSuffix)
if err != nil {
return "", errors.New("cannot find go tool: " + err.Error())
}
return goBin, nil
}
// knownEnv is a list of environment variables that affect the operation
// of the Go command.
const knownEnv = `
AR
CC
CGO_CFLAGS
CGO_CFLAGS_ALLOW
CGO_CFLAGS_DISALLOW
CGO_CPPFLAGS
CGO_CPPFLAGS_ALLOW
CGO_CPPFLAGS_DISALLOW
CGO_CXXFLAGS
CGO_CXXFLAGS_ALLOW
CGO_CXXFLAGS_DISALLOW
CGO_ENABLED
CGO_FFLAGS
CGO_FFLAGS_ALLOW
CGO_FFLAGS_DISALLOW
CGO_LDFLAGS
CGO_LDFLAGS_ALLOW
CGO_LDFLAGS_DISALLOW
CXX
FC
GCCGO
GO111MODULE
GO386
GOAMD64
GOARCH
GOARM
GOBIN
GOCACHE
GOENV
GOEXE
GOEXPERIMENT
GOFLAGS
GOGCCFLAGS
GOHOSTARCH
GOHOSTOS
GOINSECURE
GOMIPS
GOMIPS64
GOMODCACHE
GONOPROXY
GONOSUMDB
GOOS
GOPATH
GOPPC64
GOPRIVATE
GOPROXY
GOROOT
GOSUMDB
GOTMPDIR
GOTOOLDIR
GOVCS
GOWASM
GOWORK
GO_EXTLINK_ENABLED
PKG_CONFIG
`
// goToolPath reports the path to the Go tool.
// It is a convenience wrapper around goTool.
// If the tool is unavailable goToolPath calls t.Skip.
// If the tool should be available and isn't, goToolPath calls t.Fatal.
func goToolPath(t testing.TB) string {
mustHaveGoBuild(t)
path, err := goTool()
if err != nil {
t.Fatal(err)
}
// Add all environment variables that affect the Go command to test metadata.
// Cached test results will be invalidate when these variables change.
// See golang.org/issue/32285.
for _, envVar := range strings.Fields(knownEnv) {
os.Getenv(envVar)
}
return path
}
// hasGoBuild reports whether the current system can build programs with “go build”
// and then run them with os.StartProcess or exec.Command.
func hasGoBuild() bool {
if os.Getenv("GO_GCFLAGS") != "" {
// It's too much work to require every caller of the go command
// to pass along "-gcflags="+os.Getenv("GO_GCFLAGS").
// For now, if $GO_GCFLAGS is set, report that we simply can't
// run go build.
return false
}
return true
}
// mustHaveGoBuild checks that the current system can build programs with “go build”
// and then run them with os.StartProcess or exec.Command.
// If not, mustHaveGoBuild calls t.Skip with an explanation.
func mustHaveGoBuild(t testing.TB) {
if os.Getenv("GO_GCFLAGS") != "" {
t.Skipf("skipping test: 'go build' not compatible with setting $GO_GCFLAGS")
}
if !hasGoBuild() {
t.Skipf("skipping test: 'go build' not available on %s/%s", runtime.GOOS, runtime.GOARCH)
}
}
// hasGoRun reports whether the current system can run programs with “go run.”
func hasGoRun() bool {
// For now, having go run and having go build are the same.
return hasGoBuild()
}
// mustHaveGoRun checks that the current system can run programs with “go run.”
// If not, mustHaveGoRun calls t.Skip with an explanation.
func mustHaveGoRun(t testing.TB) {
if !hasGoRun() {
t.Skipf("skipping test: 'go run' not available on %s/%s", runtime.GOOS, runtime.GOARCH)
}
}
var (
gorootOnce sync.Once
gorootPath string
gorootErr error
)
func findGOROOT() (string, error) {
gorootOnce.Do(func() {
gorootPath = runtime.GOROOT()
if gorootPath != "" {
// If runtime.GOROOT() is non-empty, assume that it is valid.
//
// (It might not be: for example, the user may have explicitly set GOROOT
// to the wrong directory, or explicitly set GOROOT_FINAL but not GOROOT
// and hasn't moved the tree to GOROOT_FINAL yet. But those cases are
// rare, and if that happens the user can fix what they broke.)
return
}
// runtime.GOROOT doesn't know where GOROOT is (perhaps because the test
// binary was built with -trimpath, or perhaps because GOROOT_FINAL was set
// without GOROOT and the tree hasn't been moved there yet).
//
// Since this is internal/testenv, we can cheat and assume that the caller
// is a test of some package in a subdirectory of GOROOT/src. ('go test'
// runs the test in the directory containing the packaged under test.) That
// means that if we start walking up the tree, we should eventually find
// GOROOT/src/go.mod, and we can report the parent directory of that.
cwd, err := os.Getwd()
if err != nil {
gorootErr = fmt.Errorf("finding GOROOT: %w", err)
return
}
dir := cwd
for {
parent := filepath.Dir(dir)
if parent == dir {
// dir is either "." or only a volume name.
gorootErr = fmt.Errorf("failed to locate GOROOT/src in any parent directory")
return
}
if base := filepath.Base(dir); base != "src" {
dir = parent
continue // dir cannot be GOROOT/src if it doesn't end in "src".
}
b, err := os.ReadFile(filepath.Join(dir, "go.mod"))
if err != nil {
if os.IsNotExist(err) {
dir = parent
continue
}
gorootErr = fmt.Errorf("finding GOROOT: %w", err)
return
}
goMod := string(b)
for goMod != "" {
var line string
line, goMod, _ = strings.Cut(goMod, "\n")
fields := strings.Fields(line)
if len(fields) >= 2 && fields[0] == "module" && fields[1] == "std" {
// Found "module std", which is the module declaration in GOROOT/src!
gorootPath = parent
return
}
}
}
})
return gorootPath, gorootErr
}
// runWithTimeout runs cmd and returns its combined output. If the
// subprocess exits with a non-zero status, it will log that status
// and return a non-nil error, but this is not considered fatal.
func runWithTimeout(t testing.TB, cmd *exec.Cmd) ([]byte, error) {
args := cmd.Args
if args == nil {
args = []string{cmd.Path}
}
var b bytes.Buffer
cmd.Stdout = &b
cmd.Stderr = &b
if err := cmd.Start(); err != nil {
t.Fatalf("starting %s: %v", args, err)
}
// If the process doesn't complete within 1 minute,
// assume it is hanging and kill it to get a stack trace.
p := cmd.Process
done := make(chan bool)
go func() {
scale := 2
if s := os.Getenv("GO_TEST_TIMEOUT_SCALE"); s != "" {
if sc, err := strconv.Atoi(s); err == nil {
scale = sc
}
}
select {
case <-done:
case <-time.After(time.Duration(scale) * time.Minute):
p.Signal(os.Kill)
// If SIGQUIT doesn't do it after a little
// while, kill the process.
select {
case <-done:
case <-time.After(time.Duration(scale) * 30 * time.Second):
p.Signal(os.Kill)
}
}
}()
err := cmd.Wait()
if err != nil {
t.Logf("%s exit status: %v", args, err)
}
close(done)
return b.Bytes(), err
}
// start runs cmd asynchronously and returns immediately.
func start(t testing.TB, cmd *exec.Cmd) {
args := cmd.Args
if args == nil {
args = []string{cmd.Path}
}
var b bytes.Buffer
cmd.Stdout = &b
cmd.Stderr = &b
if err := cmd.Start(); err != nil {
t.Fatalf("starting %s: %v", args, err)
}
}