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)
+	shutdownCtx, cancel := context.WithTimeout(
+		context.WithoutCancel(ctx), shutdownTimeout)
+	defer cancel()
 
-		<-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(
-			context.WithoutCancel(ctx), shutdownTimeout)
-		defer shutdownCancel()
-
-		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
+	err = app.Stop(shutdownCtx)
+	if err != nil {
+		log.Error("Error during shutdown", "error", err)
 	}
+
+	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

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