diff --git a/cli/cmd/lab_reset.go b/cli/cmd/lab_reset.go index 58662997..cbd16f5e 100644 --- a/cli/cmd/lab_reset.go +++ b/cli/cmd/lab_reset.go @@ -487,7 +487,10 @@ func runLabReset(cmd *cobra.Command, args []string) error { if err != nil { return err } - ctx := context.Background() + // Signal-aware context from the root: `lab reset` shares provisionPlaybooks + // with `provision`, so it opens the same Bastion tunnel and needs the same + // cancellation path to tear it down on interrupt. + ctx := cmd.Context() skipPurge, _ := cmd.Flags().GetBool("skip-purge") skipProvision, _ := cmd.Flags().GetBool("skip-provision") diff --git a/cli/cmd/provision.go b/cli/cmd/provision.go index cb42216a..fe643a99 100644 --- a/cli/cmd/provision.go +++ b/cli/cmd/provision.go @@ -327,7 +327,10 @@ func runProvision(cmd *cobra.Command, args []string) error { if err != nil { return err } - ctx := context.Background() + // Signal-aware context from the root: Ctrl+C/SIGTERM cancels ctx so + // provisionPlaybooks unwinds and its deferred socksTunnel.Close() runs, + // instead of the process dying with the Bastion tunnel orphaned. + ctx := cmd.Context() playsFlag, _ := cmd.Flags().GetString("plays") fromFlag, _ := cmd.Flags().GetString("from") diff --git a/cli/cmd/root.go b/cli/cmd/root.go index ac8a3c05..4dba71a0 100644 --- a/cli/cmd/root.go +++ b/cli/cmd/root.go @@ -1,8 +1,11 @@ package cmd import ( + "context" "fmt" "os" + "os/signal" + "syscall" "github.com/dreadnode/dreadgoad/internal/config" "github.com/dreadnode/dreadgoad/internal/logging" @@ -52,8 +55,30 @@ func SetVersionInfo(version, commit, date string) { // Execute runs the root cobra command and returns any error encountered. // It is the entry point called from main. +// +// The command runs under a signal-aware context: Ctrl+C (SIGINT) or SIGTERM +// cancels ctx instead of hard-killing the process, so ctx-aware commands unwind +// and run their deferred cleanup (e.g. tearing down Bastion tunnels) rather than +// leaking child processes. A second signal force-quits. func Execute() error { - if err := rootCmd.Execute(); err != nil { + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + // NotifyContext only cancels ctx; it leaves the handler installed and drops + // every later signal on the floor, so a teardown that hangs would trap the + // user with no way out but SIGKILL. Watch the signals separately and hard + // exit on the second one. + forceQuit := make(chan os.Signal, 2) + signal.Notify(forceQuit, os.Interrupt, syscall.SIGTERM) + defer signal.Stop(forceQuit) + go func() { + <-forceQuit // first: ctx cancellation above drives the graceful unwind + <-forceQuit // second: caller is done waiting for cleanup + fmt.Fprintln(os.Stderr, "\ninterrupted again — exiting now; Bastion tunnels may be left running") + os.Exit(130) + }() + + if err := rootCmd.ExecuteContext(ctx); err != nil { fmt.Fprintln(os.Stderr, err) return err } diff --git a/cli/cmd/validate.go b/cli/cmd/validate.go index 131559f8..f3b467e1 100644 --- a/cli/cmd/validate.go +++ b/cli/cmd/validate.go @@ -117,7 +117,10 @@ func makeRunChecks(v *validate.Validator, p provider.Provider, quick bool) func( } func runValidate(cmd *cobra.Command, args []string) error { - ctx := context.Background() + // Signal-aware context from the root: Ctrl+C/SIGTERM cancels ctx so this + // function unwinds and the deferred Drain below runs (tearing down the + // Bastion tunnel) instead of the process dying with the tunnel orphaned. + ctx := cmd.Context() opts, err := validateOptsFromFlags(cmd) if err != nil { @@ -133,6 +136,13 @@ func runValidate(cmd *cobra.Command, args []string) error { return err } + // Guarantee provider teardown (Bastion tunnel + cached clients) on every + // exit path — normal return, error, or interrupt. Drain is idempotent, so + // the terminal Drain inside makeRunChecks on the happy path is harmless. + if d, ok := infra.Provider.(provider.Drainer); ok { + defer d.Drain() + } + useTUI := !opts.plain && term.IsTerminal(int(os.Stdout.Fd())) if opts.pollInterval > 0 && !useTUI { fmt.Fprintf(os.Stderr, "Warning: --poll is ignored without the live dashboard (TTY/--plain)\n") diff --git a/cli/internal/ansible/retry.go b/cli/internal/ansible/retry.go index 472c0126..a15b6f86 100644 --- a/cli/internal/ansible/retry.go +++ b/cli/internal/ansible/retry.go @@ -58,6 +58,13 @@ func RunPlaybookWithRetry(ctx context.Context, opts RetryOptions) error { retryForks := 2 // limit SSM concurrency to avoid session saturation for attempt := range opts.MaxRetries { + // An interrupt reaches ansible-playbook directly (shared foreground + // process group), so the attempt comes back as an ordinary failure and + // the error-strategy path below would log a retry it cannot perform. + // Bail on cancellation before interpreting any result. + if err := ctx.Err(); err != nil { + return err + } if attempt > 0 { log.Info("retry attempt", "attempt", attempt, "playbook", opts.Playbook) log.Info("waiting before retry", "delay", opts.RetryDelay) diff --git a/cli/internal/ansible/retry_test.go b/cli/internal/ansible/retry_test.go index c6481105..596dae5f 100644 --- a/cli/internal/ansible/retry_test.go +++ b/cli/internal/ansible/retry_test.go @@ -1,9 +1,41 @@ package ansible import ( + "bytes" + "context" + "errors" + "log/slog" + "strings" "testing" ) +// TestRunPlaybookWithRetryStopsOnCancelledContext pins the cancellation check +// at the top of the retry loop. Wiring `provision` to the root's signal-aware +// context means an interrupt now surfaces as an ordinary playbook failure, and +// without the check the loop interprets that failure and announces retries it +// cannot perform. The returned error is context.Canceled either way, so this +// asserts on the log: no attempt may be started once ctx is done. +func TestRunPlaybookWithRetryStopsOnCancelledContext(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + var logBuf bytes.Buffer + err := RunPlaybookWithRetry(ctx, RetryOptions{ + Playbook: "noop.yml", + Env: "test", + MaxRetries: 3, + Log: slog.New(slog.NewTextHandler(&logBuf, nil)), + }) + + if !errors.Is(err, context.Canceled) { + t.Fatalf("got err %v, want context.Canceled", err) + } + if got := logBuf.String(); strings.Contains(got, "starting playbook") || + strings.Contains(got, "retrying with") { + t.Fatalf("cancelled context still drove a retry attempt; log was:\n%s", got) + } +} + // TestBuildRetryLimit covers all branches of buildRetryLimit. func TestBuildRetryLimit(t *testing.T) { tests := []struct { diff --git a/cli/internal/azure/provision_tunnel.go b/cli/internal/azure/provision_tunnel.go index fae8da86..7bcd50ef 100644 --- a/cli/internal/azure/provision_tunnel.go +++ b/cli/internal/azure/provision_tunnel.go @@ -9,6 +9,8 @@ import ( "path/filepath" "strconv" "strings" + "sync" + "syscall" "time" "github.com/dreadnode/dreadgoad/internal/ludus" @@ -22,6 +24,7 @@ type ProvisionTunnel struct { socks *ludus.SOCKSTunnel bastionCmd *exec.Cmd localPort int + closeOnce sync.Once } // ProxyURL returns the SOCKS5 proxy URL Ansible's psrp connection plugin @@ -35,14 +38,85 @@ func (t *ProvisionTunnel) SOCKSAddr() string { } // Close terminates the SOCKS5 listener, the underlying SSH connection to the -// controller, and the spawned `az network bastion tunnel` subprocess. +// controller, and the spawned `az network bastion tunnel` subprocess tree. +// +// Teardown runs exactly once even if Close is called concurrently. That is not +// cosmetic: killBastionTunnel reaps via cmd.Wait, and two Wait calls racing on +// one exec.Cmd is a data race the detector flags. Callers reach Close through +// several paths (winrmRunner.close, the deferred Drain in `validate`, the +// deferred socksTunnel.Close in `provision`), so the guarantee lives here +// rather than depending on every caller staying serialized. func (t *ProvisionTunnel) Close() { - if t.socks != nil { - t.socks.Close() + t.closeOnce.Do(func() { + if t.socks != nil { + t.socks.Close() + } + killBastionTunnel(t.bastionCmd) + }) +} + +// killBastionTunnel reaps the whole `az network bastion tunnel` process tree. +// The `az` entry point is a shell wrapper that *spawns* a `python -m azure.cli` +// child, so killing only cmd.Process (the wrapper) leaves that child running — +// it reparents to init/launchd and the Bastion tunnel leaks. We start the +// command in its own process group (Setpgid) and signal the whole group here. +// killGracePeriod is how long the tunnel gets to honor SIGTERM before the +// group is SIGKILLed. +const killGracePeriod = 500 * time.Millisecond + +func killBastionTunnel(cmd *exec.Cmd) { + if cmd == nil || cmd.Process == nil { + return + } + pid := cmd.Process.Pid + + // Reap in the background so the wrapper leaves zombie state the moment it + // dies. This is what makes the polling below work at all: an unreaped + // zombie still answers kill(pid, 0), so waiting on the group without + // concurrently reaping would never observe the exit. + reaped := make(chan struct{}) + go func() { + _ = cmd.Wait() + close(reaped) + }() + + // pgid == pid proves Setpgid took effect. Without that check, a cmd started + // without SysProcAttr.Setpgid reports *our* group, and the negative-pid kill + // below would take down dreadgoad itself along with its foreground group. + pgid, err := syscall.Getpgid(pid) + if err != nil || pgid != pid { + _ = cmd.Process.Kill() + <-reaped + return } - if t.bastionCmd != nil && t.bastionCmd.Process != nil { - _ = t.bastionCmd.Process.Kill() - _ = t.bastionCmd.Wait() + + // Negative pid targets the entire process group (wrapper + python). + _ = syscall.Kill(-pgid, syscall.SIGTERM) + if !awaitGroupExit(pgid, reaped, killGracePeriod) { + _ = syscall.Kill(-pgid, syscall.SIGKILL) + } + <-reaped +} + +// awaitGroupExit polls until no member of pgid remains, or timeout elapses. +// Returns true if the group went away on its own, letting the caller skip the +// SIGKILL escalation — a tunnel that honors SIGTERM promptly costs a few +// milliseconds here instead of the full grace period. +func awaitGroupExit(pgid int, reaped <-chan struct{}, timeout time.Duration) bool { + deadline := time.Now().Add(timeout) + for { + select { + case <-reaped: + // Wrapper is reaped, so a surviving group means real stragglers. + if syscall.Kill(-pgid, 0) == syscall.ESRCH { + return true + } + default: + } + if !time.Now().Before(deadline) { + return false + } + time.Sleep(10 * time.Millisecond) } } @@ -73,7 +147,9 @@ func StartProvisionTunnel(ctx context.Context, c *Client, env string) (*Provisio return nil, fmt.Errorf("controller ephemeral key not found at expected path; was 'infra apply' run?") } - cmd := exec.Command("az", "network", "bastion", "tunnel", + // exec.CommandContext so a cancelled ctx (Ctrl+C / SIGTERM propagated to a + // signal-aware root context) kills the tunnel even if Close() is skipped. + cmd := exec.CommandContext(ctx, "az", "network", "bastion", "tunnel", "--name", bastion.Name, "--resource-group", bastion.ResourceGroup, "--target-resource-id", controller.ID, @@ -81,13 +157,25 @@ func StartProvisionTunnel(ctx context.Context, c *Client, env string) (*Provisio "--port", strconv.Itoa(localPort)) cmd.Stdout = os.Stderr cmd.Stderr = os.Stderr + // Own process group so Close() (and ctx-cancel) can reap the wrapper *and* + // its python child as one unit — see killBastionTunnel. + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + // Kill the whole group on ctx-cancel, not just the wrapper process. Same + // pgid == pid guard as killBastionTunnel: never negative-kill a group we + // haven't confirmed belongs to the child. + cmd.Cancel = func() error { + pid := cmd.Process.Pid + if pgid, err := syscall.Getpgid(pid); err == nil && pgid == pid { + return syscall.Kill(-pgid, syscall.SIGKILL) + } + return cmd.Process.Kill() + } if err := cmd.Start(); err != nil { return nil, fmt.Errorf("start bastion tunnel: %w", err) } if err := waitForLocalPort(ctx, localPort, 60*time.Second); err != nil { - _ = cmd.Process.Kill() - _ = cmd.Wait() + killBastionTunnel(cmd) return nil, fmt.Errorf("bastion tunnel never came up on :%d: %w", localPort, err) } @@ -101,8 +189,7 @@ func StartProvisionTunnel(ctx context.Context, c *Client, env string) (*Provisio } socks, err := ludus.StartSOCKSTunnel(sshCfg) if err != nil { - _ = cmd.Process.Kill() - _ = cmd.Wait() + killBastionTunnel(cmd) return nil, fmt.Errorf("start SOCKS5 over controller: %w", err) } diff --git a/cli/internal/azure/provision_tunnel_test.go b/cli/internal/azure/provision_tunnel_test.go new file mode 100644 index 00000000..ab2ef021 --- /dev/null +++ b/cli/internal/azure/provision_tunnel_test.go @@ -0,0 +1,158 @@ +//go:build !windows + +package azure + +import ( + "bufio" + "os" + "os/exec" + "strconv" + "strings" + "sync" + "syscall" + "testing" + "time" +) + +// processAlive reports whether pid still exists (signal 0 probes without +// delivering). A reaped process yields ESRCH. +func processAlive(pid int) bool { + return syscall.Kill(pid, 0) == nil +} + +// TestKillBastionTunnelReapsChildTree is the regression guard for the tunnel +// leak: the real `az network bastion tunnel` is a shell wrapper that spawns a +// python child, so killing only the wrapper leaves the child (and its tunnel) +// running. We reproduce that topology with `sh` (wrapper) spawning a +// backgrounded `sleep` (child), then assert killBastionTunnel reaps BOTH by +// signalling the whole process group. +func TestKillBastionTunnelReapsChildTree(t *testing.T) { + // sh backgrounds a long sleep (the "python child"), prints its PID, then + // waits — mirroring a wrapper that outlives nothing of its own but holds a + // child that must die with it. + cmd := exec.Command("sh", "-c", "sleep 120 & echo $!; wait") + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + + stdout, err := cmd.StdoutPipe() + if err != nil { + t.Fatalf("stdout pipe: %v", err) + } + if err := cmd.Start(); err != nil { + t.Fatalf("start: %v", err) + } + + // Read the grandchild (sleep) PID the wrapper printed. Read through the + // newline rather than trusting one Read to return the whole line, so a + // split write can't turn into a flaky parse. + line, err := bufio.NewReader(stdout).ReadString('\n') + if err != nil { + t.Fatalf("read child pid (got %q): %v", line, err) + } + childPID, err := strconv.Atoi(strings.TrimSpace(line)) + if err != nil { + t.Fatalf("parse child pid %q: %v", line, err) + } + + if !processAlive(childPID) { + t.Fatalf("precondition failed: child %d not alive after start", childPID) + } + + killBastionTunnel(cmd) + + // The grandchild reparents to init and is reaped shortly after SIGKILL. + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if !processAlive(childPID) { + return // reaped — the leak is fixed + } + time.Sleep(50 * time.Millisecond) + } + t.Fatalf("child process %d survived killBastionTunnel — tunnel would leak", childPID) +} + +// TestKillBastionTunnelNilSafe guards the early-error paths that may call +// Close() before the command was started. +func TestKillBastionTunnelNilSafe(t *testing.T) { + killBastionTunnel(nil) + killBastionTunnel(&exec.Cmd{}) // Process == nil +} + +// TestProvisionTunnelCloseIsRaceFree pins the closeOnce guard. killBastionTunnel +// reaps with cmd.Wait, and two Wait calls on one exec.Cmd is a data race — so +// concurrent Close must collapse to a single teardown. Run under -race; without +// closeOnce the detector reports a write/write race inside os/exec.(*Cmd).Wait. +func TestProvisionTunnelCloseIsRaceFree(t *testing.T) { + cmd := exec.Command("sleep", "120") + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + if err := cmd.Start(); err != nil { + t.Fatalf("start: %v", err) + } + // socks stays nil: this exercises the subprocess half, which is where the + // race lives. + tunnel := &ProvisionTunnel{bastionCmd: cmd} + + var wg sync.WaitGroup + for range 4 { + wg.Add(1) + go func() { + defer wg.Done() + tunnel.Close() + }() + } + wg.Wait() + + if processAlive(cmd.Process.Pid) { + t.Fatalf("process %d survived Close", cmd.Process.Pid) + } +} + +// pgidGuardEnv re-enters this test binary as a subprocess for the guard check +// below. A regression there SIGKILLs the caller's whole process group, so the +// dangerous half runs isolated in its own group rather than taking down +// `go test` (and the developer's shell) with it. +const pgidGuardEnv = "DREADGOAD_TEST_PGID_GUARD_CHILD" + +// pgidGuardOK is the exit code the child reports when it survived the kill. +const pgidGuardOK = 7 + +// TestKillBastionTunnelSpareOwnProcessGroup pins the `pgid == pid` guard in +// killBastionTunnel. Given a command started WITHOUT SysProcAttr.Setpgid, +// syscall.Getpgid returns the *caller's* group — so an unguarded +// kill(-pgid, SIGKILL) would take down dreadgoad itself. The guard must detect +// that and fall back to killing only the single process. +func TestKillBastionTunnelSpareOwnProcessGroup(t *testing.T) { + if os.Getenv(pgidGuardEnv) == "1" { + // Detach into our own process group so a regression's group-kill is + // contained to this subprocess. + if err := syscall.Setpgid(0, 0); err != nil { + os.Exit(3) + } + // No Setpgid here: the child inherits OUR pgid, which is exactly the + // condition the guard exists to catch. + victim := exec.Command("sleep", "120") + if err := victim.Start(); err != nil { + os.Exit(4) + } + killBastionTunnel(victim) + // Still executing => the guard held and we did not signal our own group. + os.Exit(pgidGuardOK) + } + + exe, err := os.Executable() + if err != nil { + t.Fatalf("locate test binary: %v", err) + } + cmd := exec.Command(exe, "-test.run=TestKillBastionTunnelSpareOwnProcessGroup") + cmd.Env = append(os.Environ(), pgidGuardEnv+"=1") + + err = cmd.Run() + code := cmd.ProcessState.ExitCode() + if code == pgidGuardOK { + return // guard held + } + if code == -1 { + t.Fatalf("subprocess was killed by a signal (%v) — killBastionTunnel "+ + "signalled its own process group; the pgid == pid guard is missing", cmd.ProcessState) + } + t.Fatalf("subprocess exited %d (err=%v), want %d", code, err, pgidGuardOK) +}