From 332bc909cee5336c6fe4184667445c8d1b5d3cd9 Mon Sep 17 00:00:00 2001 From: sneak Date: Mon, 21 Sep 2026 13:30:25 +0000 Subject: [PATCH] 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 | 4 +- internal/cli/cli.go | 16 ++++- internal/cli/ssh/to.go | 8 +++ internal/cli/ssh_test.go | 132 +++++++++++++++++++++++++++++++++++++-- 4 files changed, 154 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 975ecd7..5f4e523 100644 --- a/README.md +++ b/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=` 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` diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 57b4da1..2dd443e 100644 --- a/internal/cli/cli.go +++ b/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 } diff --git a/internal/cli/ssh/to.go b/internal/cli/ssh/to.go index 1afd7a9..6569c5c 100644 --- a/internal/cli/ssh/to.go +++ b/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 diff --git a/internal/cli/ssh_test.go b/internal/cli/ssh_test.go index 775d1e9..1bf8193 100644 --- a/internal/cli/ssh_test.go +++ b/internal/cli/ssh_test.go @@ -3,11 +3,14 @@ package cli_test import ( "bytes" "os" + "os/exec" "path/filepath" "slices" "strconv" "strings" + "syscall" "testing" + "time" "git.eeqj.de/sneak/keyfunc/internal/cli" "git.eeqj.de/sneak/keyfunc/internal/cli/ssh" @@ -15,6 +18,23 @@ import ( "github.com/stretchr/testify/require" ) +// runAsTool, set in the environment of a re-executed test binary, tells +// TestMain to run the tool through Main rather than the suite, so the +// signal test can drive the real signal path in a process it can send a +// signal to. +const runAsTool = "KEYFUNC_TEST_RUN_AS_TOOL" + +// TestMain re-executes the test binary as the tool when runAsTool is +// set, and otherwise runs the suite. The signal test starts the tool +// this way, as a subprocess it can signal and watch clean up. +func TestMain(m *testing.M) { + if os.Getenv(runAsTool) == "1" { + os.Exit(cli.Main()) + } + + os.Exit(m.Run()) +} + // The modes the host is supposed to end up with, and the mode the // stand-ins need so that they can be run at all. const ( @@ -47,6 +67,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 +189,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 +456,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 +465,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 +483,100 @@ 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()) + + // The three signals the tool handles, checked one after another. + signals := []struct { + name string + signal os.Signal + }{ + {"SIGTERM", syscall.SIGTERM}, + {"SIGINT", syscall.SIGINT}, + {"SIGHUP", syscall.SIGHUP}, + } + + for _, ending := range signals { + signalEndsTheTool(t, ending.name, ending.signal) + } +} + +// signalEndsTheTool runs the tool as a subprocess against a stand-in +// ssh that blocks, waits until the agent is up and ssh is running +// against it, sends the tool the signal, and requires the agent socket +// and its directory to be gone once the tool has ended. The subprocess +// goes through Main and its signal handling, so with that handling +// removed the signal kills the tool outright, no deferred cleanup runs, +// the directory is left behind, and the check fails. +func signalEndsTheTool(t *testing.T, name string, signal os.Signal) { + t.Helper() + + noted := filepath.Join(t.TempDir(), "socket") + t.Setenv("KEYFUNC_TEST_SOCKET", noted) + standIn(t, "ssh", sleeper) + + //nolint:gosec // the binary is this test's own, re-run as the tool + command := exec.CommandContext( + t.Context(), os.Args[0], subcommand, "to", host, remoteCommand, + ) + + command.Env = append(os.Environ(), runAsTool+"=1") + require.NoError(t, command.Start()) + + // The stand-in notes the socket only once the agent is up and ssh + // is running against it, so this is where the signal lands. + socket := waitForSocket(t, noted) + + require.NoError(t, command.Process.Signal(signal)) + waitForTool(t, name, command) + + // The signal ended the tool, and its deferred cleanup still ran: + // the agent socket and its directory are gone. + require.NoDirExists(t, filepath.Dir(socket), name) +} + +// waitForTool waits for the subprocess to end, and fails the test if it +// does not end in time. +func waitForTool(t *testing.T, name string, command *exec.Cmd) { + t.Helper() + + done := make(chan error, 1) + go func() { done <- command.Wait() }() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatalf("the tool did not end after %s", name) + } +} + +// 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 +598,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