diff --git a/internal/cli/app.go b/internal/cli/app.go index 4a366cf..dc305b3 100644 --- a/internal/cli/app.go +++ b/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 }, diff --git a/internal/cli/runapp_interrupt_test.go b/internal/cli/runapp_interrupt_test.go new file mode 100644 index 0000000..cfd12d7 --- /dev/null +++ b/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") +} 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 != ""