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

@@ -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
 				},

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

@@ -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 != ""