ensure command cancellation works

Author: theodore-s-beers

checks/cli.go

@@ -85,6 +85,7 @@ func runCLICommandWithLimits(
 		cmd = exec.CommandContext(ctx, "sh", "-c", finalCommand)
 	}
 
+	configureCommandCancellation(cmd)
 	cmd.Env = append(os.Environ(), "LANG=en_US.UTF-8")
 	cmd.WaitDelay = commandWaitDelay
 	cancelForOutputLimit := func() {
@@ -94,6 +95,8 @@ func runCLICommandWithLimits(
 	stderr := newBoundedBuffer(maxOutputBytesPerStream, cancelForOutputLimit)
 	cmd.Stdout = stdout
 	cmd.Stderr = stderr
+	stopSignalForwarding := forwardSignalsToCommand(cmd)
+	defer stopSignalForwarding()
 	err := cmd.Run()
 	if ee, ok := err.(*exec.ExitError); ok {
 		result.ExitCode = ee.ExitCode()

checks/command_process_other.go

@@ -0,0 +1,11 @@
+//go:build !aix && !darwin && !dragonfly && !freebsd && !linux && !netbsd && !openbsd && !solaris && !windows
+
+package checks
+
+import "os/exec"
+
+func configureCommandCancellation(cmd *exec.Cmd) {}
+
+func forwardSignalsToCommand(cmd *exec.Cmd) func() {
+	return func() {}
+}

checks/command_process_unix.go

@@ -0,0 +1,63 @@
+//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris
+
+package checks
+
+import (
+	"errors"
+	"os"
+	"os/exec"
+	"os/signal"
+	"sync"
+	"syscall"
+)
+
+func configureCommandCancellation(cmd *exec.Cmd) {
+	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+	cmd.Cancel = func() error { return killCommandProcessGroup(cmd) }
+}
+
+func forwardSignalsToCommand(cmd *exec.Cmd) func() {
+	signals := make(chan os.Signal, 1)
+	done := make(chan struct{})
+	signal.Notify(signals, os.Interrupt, syscall.SIGTERM)
+	var forwardOnce sync.Once
+	forward := func(received os.Signal) {
+		forwardOnce.Do(func() {
+			_ = killCommandProcessGroup(cmd)
+			signal.Reset(received)
+			if unixSignal, ok := received.(syscall.Signal); ok {
+				_ = syscall.Kill(os.Getpid(), unixSignal)
+			}
+		})
+	}
+
+	go func() {
+		select {
+		case received := <-signals:
+			forward(received)
+		case <-done:
+		}
+	}()
+
+	return func() {
+		signal.Stop(signals)
+		select {
+		case received := <-signals:
+			forward(received)
+		default:
+		}
+		close(done)
+	}
+}
+
+func killCommandProcessGroup(cmd *exec.Cmd) error {
+	if cmd.Process == nil {
+		return os.ErrProcessDone
+	}
+
+	err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
+	if errors.Is(err, syscall.ESRCH) {
+		return os.ErrProcessDone
+	}
+	return err
+}

checks/command_process_unix_test.go

@@ -0,0 +1,52 @@
+//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris
+
+package checks
+
+import (
+	"errors"
+	"strconv"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	api "github.com/bootdotdev/bootdev/client"
+)
+
+func TestRunCLICommandTimeoutKillsDescendants(t *testing.T) {
+	result := runCLICommandWithLimits(
+		api.CLIStepCLICommand{Command: "sleep 30 & echo $!; wait"},
+		map[string]string{},
+		100*time.Millisecond,
+		1024,
+	)
+	if !strings.Contains(result.Err, "command timed out") {
+		t.Fatalf("command error = %q, want timeout error", result.Err)
+	}
+
+	pid, err := strconv.Atoi(strings.TrimSpace(result.Stdout))
+	if err != nil {
+		t.Fatalf("child PID output = %q: %v", result.Stdout, err)
+	}
+	childAlive := true
+	t.Cleanup(func() {
+		if childAlive {
+			_ = syscall.Kill(pid, syscall.SIGKILL)
+		}
+	})
+
+	deadline := time.Now().Add(time.Second)
+	for time.Now().Before(deadline) {
+		err := syscall.Kill(pid, 0)
+		if errors.Is(err, syscall.ESRCH) {
+			childAlive = false
+			return
+		}
+		if err != nil {
+			t.Fatalf("check child process %d: %v", pid, err)
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	t.Fatalf("child process %d survived command cancellation", pid)
+}

checks/command_process_windows.go

@@ -0,0 +1,38 @@
+//go:build windows
+
+package checks
+
+import (
+	"errors"
+	"os"
+	"os/exec"
+	"strconv"
+)
+
+func configureCommandCancellation(cmd *exec.Cmd) {
+	cmd.Cancel = func() error {
+		if cmd.Process == nil {
+			return os.ErrProcessDone
+		}
+
+		treeKill := exec.Command(
+			"taskkill.exe",
+			"/PID", strconv.Itoa(cmd.Process.Pid),
+			"/T",
+			"/F",
+		)
+		if err := treeKill.Run(); err == nil {
+			return nil
+		}
+
+		err := cmd.Process.Kill()
+		if errors.Is(err, os.ErrProcessDone) {
+			return os.ErrProcessDone
+		}
+		return err
+	}
+}
+
+func forwardSignalsToCommand(cmd *exec.Cmd) func() {
+	return func() {}
+}