fix: end the session on stop signals in run and proxy, and exit the way the child did

CONTRIBUTING.md

@@ -25,6 +25,16 @@ This runs: go vet, gofmt, build, unit tests (with race detector), public audit,
 4. Run `./scripts/run-tests.sh`
 5. Open a pull request
 
+## Optional live test
+
+`test-run-signal-live.sh` verifies that `shellroute run` forwards SIGTERM to the child process, ends the API session cleanly, and then exits by SIGTERM the way the child did. It creates one real paid session and requires `--live`:
+
+```bash
+./scripts/test-run-signal-live.sh --live [COUNTRY]
+```
+
+Not part of `run-tests.sh` or CI. Run manually before releasing signal-handling changes.
+
 ## DCO
 
 All commits must be signed off (`git commit -s`). This certifies you wrote the code or have the right to submit it under the Apache 2.0 license.

internal/cli/connect.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
-	"os/signal"
 	"syscall"
 	"time"
 
@@ -127,8 +126,13 @@ func runConnectHeadless(cfg *config.Config, country string) error {
 
 	client := api.New(cfg.APIURL, cfg.APIKey)
 
+	// Stop signals are handled from here on: during startup they abort the
+	// session, afterwards they end it.
+	ctx, cancel := context.WithCancel(context.Background())
+	sigs := NewSignalHandler(cancel, syscall.SIGINT, syscall.SIGTERM)
+
 	sess, err := session.Start(
-		context.Background(),
+		ctx,
 		client,
 		&api.SessionCreateRequest{
 			Country: country,
@@ -140,10 +144,16 @@ func runConnectHeadless(cfg *config.Config, country string) error {
 		session.StartOpts{Mode: "proxy"},
 	)
 	if err != nil {
+		sigs.Stop()
 		return handleSessionError(err)
 	}
 
+	if sigs.StartupSignal() != 0 {
+		return waitAndDisconnect(sess, sigs)
+	}
+
 	if sess.GetExitIP() == "" {
+		sigs.Stop()
 		sess.Stop()
 		display.Error("Connection failed — no working upstream. Try again.")
 		return fmt.Errorf("no exit IP")
@@ -156,13 +166,13 @@ func runConnectHeadless(cfg *config.Config, country string) error {
 		outputEnv(sess)
 	}
 
-	return waitAndDisconnect(sess)
+	return waitAndDisconnect(sess, sigs)
 }
 
-func waitAndDisconnect(sess *session.Session) error {
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
-	<-sigCh
+// waitAndDisconnect blocks until a stop signal, then ends the session.
+func waitAndDisconnect(sess *session.Session, sigs *SignalHandler) error {
+	sigs.Wait()
+	sigs.Stop()
 
 	fmt.Fprintln(os.Stderr)
 

internal/cli/run.go

@@ -104,9 +104,13 @@ func runRun(cmd *cobra.Command, args []string) error {
 
 	client := api.New(cfg.APIURL, cfg.APIKey)
 
-	// Start session
+	// Stop signals are handled from here on: during startup they abort the
+	// session, once the child runs they are forwarded to it.
+	ctx, cancel := context.WithCancel(context.Background())
+	sigs := NewSignalHandler(cancel, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
+
 	sess, err := session.Start(
-		context.Background(),
+		ctx,
 		client,
 		&api.SessionCreateRequest{
 			Country: country,
@@ -118,10 +122,18 @@ func runRun(cmd *cobra.Command, args []string) error {
 		session.StartOpts{TrackRelays: true, Mode: "run"},
 	)
 	if err != nil {
+		sigs.Stop()
 		return handleSessionError(err)
 	}
 
+	if sig := sigs.StartupSignal(); sig != 0 {
+		sigs.Stop()
+		endRunSession(sess)
+		exitFromSignal(sig)
+	}
+
 	if sess.GetExitIP() == "" {
+		sigs.Stop()
 		sess.Stop()
 		display.Error("Connection failed — no working upstream. Try again.")
 		return fmt.Errorf("no exit IP")
@@ -140,6 +152,7 @@ func runRun(cmd *cobra.Command, args []string) error {
 	childCmd.Env = buildProxyEnv(os.Environ(), sess.ProxyURL())
 
 	if err := childCmd.Start(); err != nil {
+		sigs.Stop()
 		sess.Stop()
 		return fmt.Errorf("failed to start command: %w", err)
 	}
@@ -169,33 +182,40 @@ func runRun(cmd *cobra.Command, args []string) error {
 		fmt.Fprintln(os.Stderr, "\n  Connection lost during execution. Command killed.")
 	}
 
+	sigs.Attach(-childCmd.Process.Pid, &childRunning, killCancel)
+
 	childErr := childCmd.Wait()
 	childRunning.Store(false)
 	close(killCancel)
+	sigs.Stop()
 
-	// Tear down session
-	resp, stopErr := sess.Stop()
-	if !runNoStat {
-		if childErr != nil {
-			display.Error("Command failed: %s", args[0])
-		}
-		if stopErr != nil {
-			display.Error("Failed to end session. Try again.")
-		} else {
-			display.SessionSummary("shellroute session ended.", resp.DurationSec, resp.BytesTotal, resp.CostUSD, resp.BalanceUSD)
-			if resp.BalanceUSD <= 0.001 {
-				display.Warn("Balance depleted. Top up at https://console.shellroute.com, then reconnect.")
-			}
-		}
+	// A signal death is reported by exiting the same way, not as a failure.
+	exitErr, _ := childErr.(*exec.ExitError)
+	if childErr != nil && !runNoStat && (exitErr == nil || childSignal(exitErr) == 0) {
+		display.Error("Command failed: %s", args[0])
 	}
+	endRunSession(sess)
 
-	if childErr != nil {
-		if exitErr, ok := childErr.(*exec.ExitError); ok {
-			os.Exit(exitErr.ExitCode())
-		}
-		return childErr
+	if exitErr != nil {
+		exitAsChild(exitErr)
+	}
+	return childErr
+}
+
+// endRunSession ends the session and prints the summary unless --no-stat.
+func endRunSession(sess *session.Session) {
+	resp, err := sess.Stop()
+	if runNoStat {
+		return
+	}
+	if err != nil {
+		display.Error("Failed to end session. Try again.")
+		return
+	}
+	display.SessionSummary("shellroute session ended.", resp.DurationSec, resp.BytesTotal, resp.CostUSD, resp.BalanceUSD)
+	if resp.BalanceUSD <= 0.001 {
+		display.Warn("Balance depleted. Top up at https://console.shellroute.com, then reconnect.")
 	}
-	return nil
 }
 
 const defaultNoProxy = "localhost,127.0.0.1,::1"

internal/cli/run_exit_test.go

@@ -0,0 +1,142 @@
+//go:build !windows
+
+package cli
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"os/exec"
+	"syscall"
+	"testing"
+	"time"
+)
+
+// exitAsChild and exitFromSignal end the process, so each case runs in a
+// re-executed test binary driven by TestHelperProcess.
+
+func TestHelperProcess(t *testing.T) {
+	switch os.Getenv("SR_TEST_HELPER") {
+	case "":
+		return
+	case "exit-as-child":
+		var cmd *exec.Cmd
+		if code := os.Getenv("SR_TEST_CHILD_EXIT"); code != "" {
+			cmd = exec.Command("bash", "-c", "exit "+code)
+		} else {
+			cmd = exec.Command("sleep", "30")
+		}
+		cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+		if err := cmd.Start(); err != nil {
+			fmt.Println(err)
+			os.Exit(99)
+		}
+		if name := os.Getenv("SR_TEST_CHILD_SIG"); name != "" {
+			time.Sleep(200 * time.Millisecond)
+			syscall.Kill(-cmd.Process.Pid, signalByName(name))
+		}
+		err := cmd.Wait()
+		var ee *exec.ExitError
+		if errors.As(err, &ee) {
+			exitAsChild(ee)
+		}
+		if err != nil {
+			fmt.Println(err)
+			os.Exit(99)
+		}
+		os.Exit(0)
+	case "exit-from-signal":
+		exitFromSignal(signalByName(os.Getenv("SR_TEST_SIG")))
+	case "ignored-hup":
+		helperIgnoredHUP()
+	}
+}
+
+func signalByName(name string) syscall.Signal {
+	switch name {
+	case "INT":
+		return syscall.SIGINT
+	case "TERM":
+		return syscall.SIGTERM
+	case "HUP":
+		return syscall.SIGHUP
+	case "KILL":
+		return syscall.SIGKILL
+	case "USR1":
+		return syscall.SIGUSR1
+	}
+	panic("unknown signal " + name)
+}
+
+type helperResult struct {
+	code int // exit status, or -1 when killed by a signal
+	sig  syscall.Signal
+	out  string
+}
+
+// runHelper re-executes the test binary running only TestHelperProcess.
+// With hupIgnored, SIGHUP is ignored on entry, as under nohup.
+func runHelper(t *testing.T, hupIgnored bool, env ...string) helperResult {
+	t.Helper()
+	var cmd *exec.Cmd
+	if hupIgnored {
+		cmd = exec.Command("bash", "-c", `trap '' HUP; exec "$0" -test.run='^TestHelperProcess$'`, os.Args[0])
+	} else {
+		cmd = exec.Command(os.Args[0], "-test.run=^TestHelperProcess$")
+	}
+	cmd.Env = append(os.Environ(), env...)
+	out, err := cmd.CombinedOutput()
+	res := helperResult{out: string(out)}
+	if err == nil {
+		return res
+	}
+	var ee *exec.ExitError
+	if !errors.As(err, &ee) {
+		t.Fatalf("helper: %v", err)
+	}
+	ws := ee.Sys().(syscall.WaitStatus)
+	if ws.Signaled() {
+		res.code, res.sig = -1, ws.Signal()
+	} else {
+		res.code = ws.ExitStatus()
+	}
+	return res
+}
+
+func TestExitAsChild(t *testing.T) {
+	cases := []struct {
+		name     string
+		env      []string
+		wantCode int
+		wantSig  syscall.Signal
+	}{
+		{"exit code kept", []string{"SR_TEST_CHILD_EXIT=7"}, 7, 0},
+		{"SIGTERM re-raised", []string{"SR_TEST_CHILD_SIG=TERM"}, -1, syscall.SIGTERM},
+		{"SIGINT re-raised", []string{"SR_TEST_CHILD_SIG=INT"}, -1, syscall.SIGINT},
+		{"SIGHUP re-raised", []string{"SR_TEST_CHILD_SIG=HUP"}, -1, syscall.SIGHUP},
+		{"SIGKILL re-raised", []string{"SR_TEST_CHILD_SIG=KILL"}, -1, syscall.SIGKILL},
+		{"other signal uses 128+n", []string{"SR_TEST_CHILD_SIG=USR1"}, 128 + int(syscall.SIGUSR1), 0},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			res := runHelper(t, false, append(c.env, "SR_TEST_HELPER=exit-as-child")...)
+			if res.code != c.wantCode || res.sig != c.wantSig {
+				t.Errorf("got code=%d sig=%v, want code=%d sig=%v\n%s", res.code, res.sig, c.wantCode, c.wantSig, res.out)
+			}
+		})
+	}
+}
+
+// exitFromSignal must not re-enable a signal that is ignored on entry: with
+// SIGHUP ignored (nohup), it falls back to exit status 129.
+func TestExitFromSignal_IgnoredFallsBack(t *testing.T) {
+	res := runHelper(t, true, "SR_TEST_HELPER=exit-from-signal", "SR_TEST_SIG=HUP")
+	if res.code != 129 || res.sig != 0 {
+		t.Errorf("got code=%d sig=%v, want code=129\n%s", res.code, res.sig, res.out)
+	}
+
+	res = runHelper(t, false, "SR_TEST_HELPER=exit-from-signal", "SR_TEST_SIG=HUP")
+	if res.sig != syscall.SIGHUP {
+		t.Errorf("got code=%d sig=%v, want killed by SIGHUP\n%s", res.code, res.sig, res.out)
+	}
+}