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