From 1e9aa5dedab77f66210e0b8563f01b5112496f85 Mon Sep 17 00:00:00 2001 From: sneak Date: Tue, 22 Sep 2026 11:07:53 +0000 Subject: [PATCH] Wait for the interrupted operation to clean up before exit (closes #159) On SIGINT/SIGTERM, RunOperation's OnStop hook only cancelled the running command and returned; nothing waited for its goroutine, so the process could exit before restore's cleanup defers ran. That left decrypted data in the temp directory: the blob cache (vaultik-blobcache-*) and the decrypted snapshot-database directory (vaultik-restore-*). StartOperation now returns a stop function that cancels the context and waits for the goroutine to return, bounded by the existing shutdownTimeout; OnStop calls it and warns if it times out. The fix is in the shared runner, so it covers every command, not only restore. Restore's chunk-write and blob-download loops also check the context between steps so the wait ends promptly. Model: opus-4-8 --- internal/cli/app.go | 23 ++- internal/vaultik/restore.go | 12 ++ internal/vaultik/restore_interrupt_test.go | 159 +++++++++++++++++++++ internal/vaultik/vaultik.go | 34 +++++ 4 files changed, 223 insertions(+), 5 deletions(-) create mode 100644 internal/vaultik/restore_interrupt_test.go diff --git a/internal/cli/app.go b/internal/cli/app.go index 4a366cf..3f4e2c8 100644 --- a/internal/cli/app.go +++ b/internal/cli/app.go @@ -238,7 +238,10 @@ var errReported = errors.New("operation failed") // // op runs in a goroutine so OnStart returns promptly and an interrupt // can still cancel through OnStop; when it finishes, success or failure, -// it triggers shutdown, which is what lets RunWithApp return. report is +// it triggers shutdown, which is what lets RunWithApp return. On an +// interrupt OnStop cancels op and waits for the goroutine to return, so +// op's cleanup (removing decrypted scratch files) runs before the +// process exits; the wait is bounded by shutdownTimeout. report is // called with a non-canceled failure so the caller can log it (and // suppress it under --json) before it becomes errReported. A context // cancellation is the interrupt path, not a failure: it is neither @@ -254,9 +257,11 @@ func RunOperation( opts.Invokes = append(opts.Invokes, fx.Invoke(func(v *vaultik.Vaultik, lc fx.Lifecycle) { + var stop func(context.Context) bool + lc.Append(fx.Hook{ OnStart: func(_ context.Context) error { - go func() { + stop = v.StartOperation(func() { err := op(v) if err != nil && !errors.Is(err, context.Canceled) { report(err) @@ -270,12 +275,20 @@ func RunOperation( if stopErr != nil { log.Error("Failed to shutdown", "error", stopErr) } - }() + }) return nil }, - OnStop: func(_ context.Context) error { - v.Cancel() + // On an interrupt, cancel the operation and wait for it to + // unwind so its cleanup defers (which remove decrypted + // scratch files from the temp directory) run before the + // process exits. The wait is bounded by ctx, the existing + // shutdownTimeout. + OnStop: func(ctx context.Context) error { + if !stop(ctx) { + log.Warn("Shutdown timed out before the operation " + + "finished; decrypted temporary files may remain") + } return nil }, diff --git a/internal/vaultik/restore.go b/internal/vaultik/restore.go index 1244943..32d51df 100644 --- a/internal/vaultik/restore.go +++ b/internal/vaultik/restore.go @@ -402,6 +402,12 @@ func (s *restoreSession) downloadNextBlobSet(plan *restorePlan) (bool, error) { } for _, hash := range plan.blobsNeeded(next) { + // Stop between blobs on cancel so an interrupt ends the download + // phase promptly rather than fetching the rest of the set. + if s.ctx.Err() != nil { + return false, s.ctx.Err() + } + blob, ok := s.blobByHash[hash] if !ok { return false, fmt.Errorf("%w: %s", errBlobMissingFromIndex, hash[:16]) @@ -1077,6 +1083,12 @@ func (s *restoreSession) writeFileChunks( ) for _, fc := range fileChunks { + // Stop between chunks on cancel so an interrupt does not keep + // writing a large file after the operation has been told to stop. + if s.ctx.Err() != nil { + return bytesWritten, timings, s.ctx.Err() + } + chunkHashStr := fc.ChunkHash.String() blobChunk, ok := s.chunkToBlobMap[chunkHashStr] diff --git a/internal/vaultik/restore_interrupt_test.go b/internal/vaultik/restore_interrupt_test.go new file mode 100644 index 0000000..ccf07a5 --- /dev/null +++ b/internal/vaultik/restore_interrupt_test.go @@ -0,0 +1,159 @@ +package vaultik //nolint:testpackage // sets ctx/cancel and inspects scratch files + +import ( + "context" + "io" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/spf13/afero" + "github.com/stretchr/testify/require" + "sneak.berlin/go/vaultik/internal/log" + "sneak.berlin/go/vaultik/internal/storage" + "sneak.berlin/go/vaultik/internal/ui" +) + +// blockingBlobStorer wraps a Storer and blocks the first blob download +// until its context is cancelled, so a test can catch a restore while it +// is mid-download. Metadata reads pass straight through, so the restore +// reaches the blob-download phase — having already written its decrypted +// scratch files — before it blocks. +type blockingBlobStorer struct { + storage.Storer + + once sync.Once + entered chan struct{} +} + +func newBlockingBlobStorer(inner storage.Storer) *blockingBlobStorer { + return &blockingBlobStorer{Storer: inner, entered: make(chan struct{})} +} + +func (b *blockingBlobStorer) Get( + ctx context.Context, key string, +) (io.ReadCloser, error) { + if strings.HasPrefix(key, "blobs/") { + b.once.Do(func() { close(b.entered) }) + <-ctx.Done() + + return nil, ctx.Err() + } + + return b.Storer.Get(ctx, key) +} + +// TestRestoreCleansTempDirOnInterrupt drives a restore through the stop +// path (v.StartOperation, which is what the fx OnStop hook uses) instead +// of calling Restore directly, catches it mid-download, and asserts that +// stopping waits for the operation to unwind and removes its decrypted +// scratch files — the blob cache and the temporary snapshot database — +// from the temp directory. Without the wait a SIGINT exits the process +// before those defers run, leaving decrypted data on disk (issue #159). +// +// Not parallel: it points TMPDIR at a private directory (via t.Setenv) +// so it can assert on exactly the scratch files this restore created. +func TestRestoreCleansTempDirOnInterrupt(t *testing.T) { + log.Initialize(log.Config{}) + + fs := afero.NewOsFs() + root := t.TempDir() + + dataDir := filepath.Join(root, "source") + storeDir := filepath.Join(root, "remote") + restoreDir := filepath.Join(root, "restored") + dbPath := filepath.Join(root, "index.sqlite") + + require.NoError(t, fs.MkdirAll(dataDir, 0o755)) + + buildLocalityFixture(t, fs, dataDir) + + cfg, storer, snapshotID := setupLocalityBackup( + context.Background(), t, fs, dataDir, storeDir, dbPath) + + // Point the "" temp paths (the blob cache directory and the + // snapshot-database directory) at a private directory so the test can + // assert on exactly the scratch this restore creates. + scratch := filepath.Join(root, "scratch") + require.NoError(t, fs.MkdirAll(scratch, 0o755)) + t.Setenv("TMPDIR", scratch) + + gate := newBlockingBlobStorer(storer) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + v := &Vaultik{ + Config: cfg, + Storage: gate, + Fs: fs, + Stdout: io.Discard, + Stderr: io.Discard, + UI: ui.NewWithColor(io.Discard, false), + ctx: ctx, + cancel: cancel, + } + + var ( + opReturned atomic.Bool + restoreErr error + ) + + stop := v.StartOperation(func() { + defer opReturned.Store(true) + + restoreErr = v.Restore(&RestoreOptions{ + SnapshotID: snapshotID, + TargetDir: restoreDir, + }) + }) + + // Wait until the restore is blocked mid-download; its decrypted + // scratch files exist by now. + select { + case <-gate.entered: + case <-time.After(30 * time.Second): + t.Fatal("restore never reached the blob-download phase") + } + + require.NotEmpty(t, scratchEntries(t, scratch), + "expected decrypted scratch files to exist mid-restore") + + // Stop the operation the way the fx OnStop hook does. + stopCtx, stopCancel := context.WithTimeout( + context.Background(), 30*time.Second) + defer stopCancel() + + require.True(t, stop(stopCtx), + "stop timed out; the operation goroutine did not return") + + // stop returns only once the operation goroutine has returned, so its + // cleanup defers have run by the time we read these. + require.True(t, opReturned.Load(), + "stop returned before the operation goroutine finished") + require.ErrorIs(t, restoreErr, context.Canceled) + require.Empty(t, scratchEntries(t, scratch), + "decrypted scratch files remained after the interrupt") +} + +// scratchEntries returns the vaultik blob-cache and snapshot-database +// scratch entries currently present in dir. +func scratchEntries(t *testing.T, dir string) []string { + t.Helper() + + var matches []string + + for _, pattern := range []string{ + "vaultik-blobcache-*", "vaultik-restore-*", + } { + found, err := filepath.Glob(filepath.Join(dir, pattern)) + require.NoError(t, err) + + matches = append(matches, found...) + } + + return matches +} diff --git a/internal/vaultik/vaultik.go b/internal/vaultik/vaultik.go index 21783eb..9511965 100644 --- a/internal/vaultik/vaultik.go +++ b/internal/vaultik/vaultik.go @@ -136,6 +136,40 @@ func (v *Vaultik) Cancel() { v.cancel() } +// StartOperation runs fn in its own goroutine and returns a stop +// function. fn is the command being run (a restore, verify, prune, and +// so on); it observes cancellation through the Vaultik context and +// removes its decrypted scratch files (the blob cache and the temporary +// snapshot database) from the temp directory as it unwinds. +// +// Calling stop cancels the Vaultik context and then blocks until fn has +// returned — so that unwinding, and the cleanup it does, completes +// before the caller proceeds — or until the passed context is done, +// whichever comes first. It reports whether fn returned before that +// deadline. A signal-driven shutdown must call stop before the process +// exits; otherwise the process can exit mid-operation and leave +// decrypted data behind. +func (v *Vaultik) StartOperation(fn func()) func(context.Context) bool { + done := make(chan struct{}) + + go func() { + defer close(done) + + fn() + }() + + return func(ctx context.Context) bool { + v.Cancel() + + select { + case <-done: + return true + case <-ctx.Done(): + return false + } + } +} + // CanDecrypt returns true if this Vaultik instance has decryption capabilities func (v *Vaultik) CanDecrypt() bool { return v.Config.AgeSecretKey != ""