Clean up the agent socket and working files when a signal ends the tool (closes #17)

Main ran the command tree on a background context, so SIGINT, SIGTERM or
SIGHUP killed the process before the deferred cleanup ran: the agent's
temporary directory and socket, and the install working directory, were
left behind.

Main now runs the tree on a signal.NotifyContext for those three signals.
A signal cancels the context, which ends the child ssh or sftp started
with exec.CommandContext, and the deferred cleanup then runs. The exit
status after a signal stays 1 unless ssh reported one of its own.

For "ssh to" the child is cancelled with SIGTERM rather than the default
kill, so ssh restores the terminal before it goes.

Model: opus-4-8

README.md

@@ -146,7 +146,9 @@ Derives the key, serves it from an SSH agent that runs inside the tool on a unix
 socket in a new private `0700` temporary directory, then runs the system `ssh`
 with `-o IdentityAgent=<that socket>` followed by the host and all remaining
 arguments unchanged. The tool exits with `ssh`'s exit status and removes the
-socket and directory on the way out. The private key is never written to disk.
+socket and directory on the way out. The private key is never written to disk. A
+SIGINT, SIGTERM or SIGHUP ends `ssh` and still removes the socket and directory,
+and the tool then exits with status 1 unless `ssh` reported one of its own.
 
 ## age identities: `keyfunc age`
 

internal/cli/cli.go

@@ -2,10 +2,13 @@
 package cli
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
+	"os/signal"
 	"runtime/debug"
+	"syscall"
 
 	"git.eeqj.de/sneak/keyfunc/internal/cli/age"
 	"git.eeqj.de/sneak/keyfunc/internal/cli/mnemonic"
@@ -66,8 +69,19 @@ func Root() *cobra.Command {
 // status of its own, which "ssh to" uses to hand on the status ssh
 // ended with. ssh has already said whatever it had to say in that
 // case, so nothing more is printed.
+//
+// SIGINT, SIGTERM and SIGHUP cancel the command's context instead of
+// killing the process outright, so the child ssh or sftp ends and the
+// deferred cleanup that removes the agent socket and the install
+// working directory still runs.
 func Main() int {
-	err := Root().Execute()
+	ctx, stop := signal.NotifyContext(
+		context.Background(),
+		syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP,
+	)
+	defer stop()
+
+	err := Root().ExecuteContext(ctx)
 	if err == nil {
 		return 0
 	}

internal/cli/ssh/to.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/exec"
 	"slices"
+	"syscall"
 
 	"github.com/spf13/cobra"
 )
@@ -75,6 +76,13 @@ func connect(ctx context.Context, argv []string) error {
 	command.Stdout = os.Stdout
 	command.Stderr = os.Stderr
 
+	// A cancelled context means a signal ended the tool. Send ssh a
+	// SIGTERM rather than the default kill, so it puts the terminal
+	// back the way it found it before it goes.
+	command.Cancel = func() error {
+		return command.Process.Signal(syscall.SIGTERM)
+	}
+
 	err := command.Run()
 	if err == nil {
 		return nil

internal/cli/ssh_test.go

@@ -2,12 +2,14 @@ package cli_test
 
 import (
 	"bytes"
+	"context"
 	"os"
 	"path/filepath"
 	"slices"
 	"strconv"
 	"strings"
 	"testing"
+	"time"
 
 	"git.eeqj.de/sneak/keyfunc/internal/cli"
 	"git.eeqj.de/sneak/keyfunc/internal/cli/ssh"
@@ -47,6 +49,9 @@ const (
 	keptIn    = "authorized_keys"
 )
 
+// remoteCommand is the command the "to" tests hand ssh after the host.
+const remoteCommand = "uptime"
+
 // The tool's own name, as it stands in the arguments a test hands to
 // Main, the ssh subcommand both commands the tests here drive live
 // under, and the one of those two these tests name most.
@@ -166,6 +171,18 @@ fi
 exit "$KEYFUNC_TEST_STATUS"
 `
 
+// sleeper is a stand-in for the system ssh that notes the agent socket
+// and then blocks, so a test can cancel the context while it is running
+// and watch the tool take the agent down. The wait ends on its own only
+// as a backstop, well after the test has cancelled and looked.
+const sleeper = `
+socket=${2#IdentityAgent=}
+if [ -S "$socket" ]; then
+	printf '%s\n' "$socket" > "$KEYFUNC_TEST_SOCKET"
+fi
+sleep 5
+`
+
 // pretended is where a stand-in writes down what it was asked to do.
 type pretended struct {
 	// home stands in for the home directory on the host.
@@ -421,7 +438,7 @@ func TestSSHIsPointedAtTheAgentAndItsStatusIsHandedOn(t *testing.T) {
 
 	arguments, noted := pretendCall(t)
 
-	_, err := execute(t, subcommand, "to", host, "uptime")
+	_, err := execute(t, subcommand, "to", host, remoteCommand)
 
 	var passed ssh.StatusError
 
@@ -430,7 +447,7 @@ func TestSSHIsPointedAtTheAgentAndItsStatusIsHandedOn(t *testing.T) {
 
 	given := recorded(t, arguments)
 	require.Equal(t, "-o", given[0])
-	require.Equal(t, []string{host, "uptime"}, given[2:])
+	require.Equal(t, []string{host, remoteCommand}, given[2:])
 
 	// The stand-in wrote the path down only because there really was
 	// a socket there while it ran.
@@ -448,11 +465,68 @@ func TestTheToolEndsWithTheStatusSSHEndedWith(t *testing.T) {
 
 	t.Cleanup(func() { os.Args = given })
 
-	os.Args = []string{tool, subcommand, "to", host, "uptime"}
+	os.Args = []string{tool, subcommand, "to", host, remoteCommand}
 
 	require.Equal(t, failingStatus, cli.Main())
 }
 
+func TestASignalTakesTheAgentDirectoryDown(t *testing.T) {
+	t.Setenv(mnemonic.Variable, example())
+
+	noted := filepath.Join(t.TempDir(), "socket")
+	t.Setenv("KEYFUNC_TEST_SOCKET", noted)
+	standIn(t, "ssh", sleeper)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	root := cli.Root()
+	root.SetOut(&bytes.Buffer{})
+	root.SetErr(&bytes.Buffer{})
+	root.SetArgs([]string{subcommand, "to", host, remoteCommand})
+
+	done := make(chan error, 1)
+	go func() { done <- root.ExecuteContext(ctx) }()
+
+	// The stand-in notes the socket only once it is up and ssh is
+	// running against it, so this is where a signal would land.
+	socket := waitForSocket(t, noted)
+
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(10 * time.Second):
+		t.Fatal("the tool did not end after the context was cancelled")
+	}
+
+	// The deferred cleanup ran even though a cancellation, not a clean
+	// exit, ended ssh: the agent socket and its directory are gone.
+	require.NoDirExists(t, filepath.Dir(socket))
+}
+
+// waitForSocket waits for the stand-in to write down the agent socket
+// and gives back the path, which means the agent is up and ssh is
+// running against it.
+func waitForSocket(t *testing.T, noted string) string {
+	t.Helper()
+
+	var socket string
+
+	require.Eventually(t, func() bool {
+		content, err := os.ReadFile(noted) //nolint:gosec // test path
+		if err != nil {
+			return false
+		}
+
+		socket = strings.TrimSpace(string(content))
+
+		return socket != ""
+	}, 5*time.Second, 5*time.Millisecond)
+
+	return socket
+}
+
 func TestTheMnemonicIsNotHandedToSFTP(t *testing.T) {
 	t.Setenv(mnemonic.CommandVariable, "echo "+example())
 	t.Setenv(mnemonic.Variable, example())
@@ -474,7 +548,7 @@ func TestTheMnemonicIsNotHandedToSSH(t *testing.T) {
 	pretendCall(t)
 	environment := recordEnvironment(t)
 
-	_, err := execute(t, subcommand, "to", host, "uptime")
+	_, err := execute(t, subcommand, "to", host, remoteCommand)
 
 	var passed ssh.StatusError