Reap Azure Bastion tunnel process tree on exit/interrupt

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")

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")

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
 	}

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")

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)

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 {

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,21 +147,35 @@ 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,
 		"--resource-port", "22",
 		"--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)
 	}