Wait for the interrupted operation to clean up before exit (closes #159)

On SIGINT/SIGTERM the process could exit before the interrupted command cleanup defers ran, leaving decrypted data in the temp directory (the blob cache and the decrypted snapshot database).

RunApp now mirrors fx run sequence: start, block on app.Wait(), then app.Stop(), returning only after Stop completes. fx delivers both an OS interrupt and the finished operation Shutdowner.Shutdown() on one channel. Stop runs the OnStop hooks; the operation hook cancels the command and waits for its goroutine to return (bounded by shutdownTimeout) before exit. The old code returned as soon as app.Done fired, without Stop, so a real interrupt unwound to os.Exit while cleanup still ran. Restore loops check the context between chunks and blobs so the wait ends promptly. A cli test drives RunApp through the OnStop hook.

Model: opus-4-8

internal/cli/app.go

@@ -7,12 +7,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"os"
-	"os/signal"
 	"path/filepath"
 	"strings"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/adrg/xdg"
@@ -161,64 +158,45 @@ func cleanStartupError(err error) error {
 	return &startupError{msg: msg}
 }
 
-// RunApp starts and stops the fx application within the given context.
-// It handles graceful shutdown on interrupt signals (SIGINT, SIGTERM) and
-// ensures the application stops cleanly. The function blocks until the
-// application completes or is interrupted. Returns an error if startup fails.
+// RunApp starts the fx application, blocks until it is asked to stop, and
+// then stops it. The app is asked to stop either by an OS interrupt
+// (SIGINT/SIGTERM — fx installs its own handler when app.Wait is called) or,
+// on normal completion, by the finished operation calling
+// Shutdowner.Shutdown(); both arrive on the app.Wait channel.
+//
+// Stopping runs the fx OnStop hooks, and RunApp does not return until Stop
+// returns. On an interrupt the operation's OnStop hook cancels the running
+// command and waits for it to unwind — removing its decrypted scratch files —
+// so the process cannot proceed to exit mid-cleanup (issue #159). Waiting for
+// Stop before returning is what makes that hook effective: routing the
+// interrupt through app.Stop and not returning until it completes is required,
+// because fx also fires the app.Wait channel on the signal, and an earlier
+// version returned on that alone — unwinding to os.Exit while the concurrent
+// cleanup still ran. The stop is bounded by shutdownTimeout. Returns an error
+// if startup fails.
 func RunApp(ctx context.Context, app *fx.App) error {
-	// Set up signal handling for graceful shutdown
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
-
-	// Create a context that will be cancelled on signal
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
-	// Start the app
 	err := app.Start(ctx)
 	if err != nil {
 		return cleanStartupError(err)
 	}
 
-	// Handle shutdown
-	shutdownComplete := make(chan struct{})
+	// Block until an interrupt or the finished operation's
+	// Shutdowner.Shutdown() arrives, then stop the app in this goroutine so we
+	// return only after its OnStop hooks — including the operation's cleanup
+	// wait — have run. Detach the stop from ctx's cancellation but keep its
+	// values, and bound it by shutdownTimeout.
+	<-app.Wait()
 
-	go func() {
-		defer close(shutdownComplete)
-
-		<-sigChan
-		log.Notice("Received interrupt signal, shutting down gracefully...")
-
-		// Create a timeout context for shutdown. The parent ctx is being
-		// cancelled, so detach from its cancellation but keep its values.
-		shutdownCtx, shutdownCancel := context.WithTimeout(
+	shutdownCtx, cancel := context.WithTimeout(
 		context.WithoutCancel(ctx), shutdownTimeout)
-		defer shutdownCancel()
+	defer cancel()
 
-		err := app.Stop(shutdownCtx)
+	err = app.Stop(shutdownCtx)
 	if err != nil {
 		log.Error("Error during shutdown", "error", err)
 	}
-	}()
 
-	// Wait for the signal handler to complete shutdown or the app to
-	// request shutdown.
-	select {
-	case <-shutdownComplete:
-		// Shutdown completed via signal
 	return nil
-	case <-ctx.Done():
-		// Context cancelled (shouldn't happen in normal operation)
-		err := app.Stop(context.WithoutCancel(ctx))
-		if err != nil {
-			log.Error("Error stopping app", "error", err)
-		}
-
-		return ctx.Err()
-	case <-app.Done():
-		// App finished running (e.g., backup completed)
-		return nil
-	}
 }
 
 // errReported marks a failure the operation has already shown the user
@@ -238,7 +216,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 +235,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 +253,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
 				},

internal/cli/runapp_interrupt_test.go

@@ -0,0 +1,97 @@
+package cli_test
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/fx"
+	"sneak.berlin/go/vaultik/internal/cli"
+)
+
+// TestRunAppWaitsForOperationCleanupOnShutdown drives RunApp with an fx app
+// wired the way RunOperation wires a command: a single lifecycle hook whose
+// OnStart launches the operation in its own goroutine and whose OnStop cancels
+// it and blocks until that goroutine returns. The operation stands in for a
+// restore blocked mid-download — it holds a decrypted "scratch" file and only
+// removes it as it unwinds on cancellation.
+//
+// The app is asked to stop once the operation is running (standing in for an
+// OS interrupt; fx delivers a real signal and Shutdowner.Shutdown() on the
+// same app.Wait channel, so both drive the identical shutdown path). RunApp
+// must not return until app.Stop has run the OnStop hook, so the scratch file
+// must be gone by the time RunApp returns. Before the fix RunApp returned as
+// soon as the app.Wait/Done channel fired, without running app.Stop, so the
+// cleanup never ran and this file would still be on disk (issue #159).
+func TestRunAppWaitsForOperationCleanupOnShutdown(t *testing.T) {
+	t.Parallel()
+
+	scratch := filepath.Join(t.TempDir(), "decrypted-scratch")
+	require.NoError(t, os.WriteFile(scratch, []byte("secret"), 0o600))
+
+	// Cancel and reap the operation even if RunApp returns without doing so
+	// (the buggy path), so the goroutine cannot leak past the test.
+	opCtx, opCancel := context.WithCancel(context.Background())
+	t.Cleanup(opCancel)
+
+	var stop func(context.Context) bool
+
+	app := fx.New(
+		fx.NopLogger,
+		fx.Invoke(func(lc fx.Lifecycle, sh fx.Shutdowner) {
+			lc.Append(fx.Hook{
+				OnStart: func(_ context.Context) error {
+					done := make(chan struct{})
+
+					go func() {
+						defer close(done)
+
+						// Blocked mid-operation until cancelled, then run the
+						// cleanup an interrupted restore would run.
+						<-opCtx.Done()
+
+						_ = os.Remove(scratch)
+					}()
+
+					stop = func(ctx context.Context) bool {
+						opCancel()
+
+						select {
+						case <-done:
+							return true
+						case <-ctx.Done():
+							return false
+						}
+					}
+
+					// Ask the app to stop now that the operation is running.
+					go func() { _ = sh.Shutdown() }()
+
+					return nil
+				},
+				OnStop: func(ctx context.Context) error {
+					stop(ctx)
+
+					return nil
+				},
+			})
+		}),
+	)
+
+	done := make(chan error, 1)
+	go func() { done <- cli.RunApp(context.Background(), app) }()
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(30 * time.Second):
+		t.Fatal("RunApp did not return after shutdown was requested")
+	}
+
+	_, err := os.Stat(scratch)
+	require.True(t, os.IsNotExist(err),
+		"RunApp returned before the operation removed its decrypted scratch file")
+}

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]

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
+}

internal/vaultik/vaultik.go

@@ -128,6 +128,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 != ""