1 Commits
Author SHA1 Message Date
sneak 332bc909ce Clean up the agent socket and working files when a signal ends the tool (closes #17)
check / check (push) Successful in 39s
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
2026-09-21 14:07:40 +00:00
+68 -18
View File
@@ -2,12 +2,13 @@ package cli_test
import (
"bytes"
"context"
"os"
"os/exec"
"path/filepath"
"slices"
"strconv"
"strings"
"syscall"
"testing"
"time"
@@ -17,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 (
@@ -473,36 +491,68 @@ func TestTheToolEndsWithTheStatusSSHEndedWith(t *testing.T) {
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)
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
//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,
)
root := cli.Root()
root.SetOut(&bytes.Buffer{})
root.SetErr(&bytes.Buffer{})
root.SetArgs([]string{subcommand, "to", host, remoteCommand})
command.Env = append(os.Environ(), runAsTool+"=1")
require.NoError(t, command.Start())
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.
// 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)
cancel()
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.Fatal("the tool did not end after the context was cancelled")
t.Fatalf("the tool did not end after %s", name)
}
// 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