diff --git a/internal/vaultik/blob_fetch.go b/internal/vaultik/blob_fetch.go index 714c7d1..aa00243 100644 --- a/internal/vaultik/blob_fetch.go +++ b/internal/vaultik/blob_fetch.go @@ -103,6 +103,8 @@ func (v *Vaultik) FetchAndDecryptBlob( } // FetchBlob downloads a blob and returns a reader for the encrypted data. +// The reader is closed as soon as ctx is cancelled, which ends a read +// blocked inside it. func (v *Vaultik) FetchBlob( ctx context.Context, blobHash string, ) (io.ReadCloser, error) { @@ -120,5 +122,47 @@ func (v *Vaultik) FetchBlob( return nil, fmt.Errorf("downloading blob %s: %w", shortHash(blobHash), err) } - return rc, nil + return newCloseOnCancelReader(ctx, rc), nil +} + +// closeOnCancelReader closes the reader it wraps when ctx is cancelled. +// Not every backend's reader stops on cancel by itself (the file +// backend's does not), and a read stalled on a slow or hung source would +// otherwise outlast the shutdown timeout, so the process would exit +// before restore or deep verify removed its decrypted temporary files. +type closeOnCancelReader struct { + io.ReadCloser + + ctx context.Context //nolint:containedctx // Read reports its error + stop func() bool +} + +func newCloseOnCancelReader( + ctx context.Context, rc io.ReadCloser, +) *closeOnCancelReader { + return &closeOnCancelReader{ + ReadCloser: rc, + ctx: ctx, + stop: context.AfterFunc(ctx, func() { _ = rc.Close() }), + } +} + +// Read returns the context's error, not the closed reader's, once ctx +// has been cancelled. +func (r *closeOnCancelReader) Read(p []byte) (int, error) { + n, err := r.ReadCloser.Read(p) + if err != nil && r.ctx.Err() != nil { + return n, r.ctx.Err() + } + + return n, err +} + +// Close closes the wrapped reader, unless the cancel already closed it. +func (r *closeOnCancelReader) Close() error { + if !r.stop() { + return nil + } + + return r.ReadCloser.Close() } diff --git a/internal/vaultik/restore_interrupt_test.go b/internal/vaultik/restore_interrupt_test.go index 0625ccb..2024a0d 100644 --- a/internal/vaultik/restore_interrupt_test.go +++ b/internal/vaultik/restore_interrupt_test.go @@ -142,14 +142,14 @@ func TestRestoreCleansTempDirOnInterrupt(t *testing.T) { } // scratchEntries returns the vaultik blob-cache and snapshot-database -// scratch entries currently present in dir. +// scratch entries (restore's and deep verify's) 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-*", + "vaultik-blobcache-*", "vaultik-restore-*", "vaultik-verify-*", } { found, err := filepath.Glob(filepath.Join(dir, pattern)) require.NoError(t, err) @@ -206,3 +206,119 @@ func TestRestoreSkipErrorsCancelDuringBlobDownload(t *testing.T) { assert.NotContains(t, out.String(), srcPath, "the cancel was reported as a failed file") } + +// stallingBlobStorer wraps a Storer so that every blob download hands +// over the first half of the blob and then stalls, ignoring the context, +// as a file:// blob on a hung USB stick or network share does. Only +// closing the returned reader ends the stall. +type stallingBlobStorer struct { + storage.Storer + + once sync.Once + stalled chan struct{} +} + +func newStallingBlobStorer(inner storage.Storer) *stallingBlobStorer { + return &stallingBlobStorer{Storer: inner, stalled: make(chan struct{})} +} + +func (s *stallingBlobStorer) Get( + ctx context.Context, key string, +) (io.ReadCloser, error) { + rc, err := s.Storer.Get(ctx, key) + if err != nil || !strings.HasPrefix(key, "blobs/") { + return rc, err + } + + blob, err := io.ReadAll(rc) + _ = rc.Close() + + if err != nil { + return nil, err + } + + // The writer is never closed, so once the first half has been read + // the next Read blocks until the reader is closed. + pr, pw := io.Pipe() + + go func() { + _, _ = pw.Write(blob[:len(blob)/2]) + + s.once.Do(func() { close(s.stalled) }) + }() + + return pr, nil +} + +// TestInterruptDuringStalledBlobReadCleansTempDir stops a restore and a +// deep verify while a blob read is stalled on a source that ignores the +// context. Stopping must end that read, so the operation unwinds and +// removes its decrypted scratch files well within the shutdown timeout +// (issue #325). +// +// Not parallel: it points TMPDIR at a private directory (via t.Setenv) +// so it can assert on exactly the scratch files the operation created. +func TestInterruptDuringStalledBlobReadCleansTempDir(t *testing.T) { + log.Initialize(log.Config{}) + + operations := map[string]func(v *Vaultik, snapshotID, targetDir string) error{ + "restore": func(v *Vaultik, snapshotID, targetDir string) error { + return v.Restore(&RestoreOptions{ + SnapshotID: snapshotID, + TargetDir: targetDir, + }) + }, + "deep verify": func(v *Vaultik, snapshotID, _ string) error { + return v.RunDeepVerify(snapshotID, &VerifyOptions{Deep: true}) + }, + } + + for name, op := range operations { + t.Run(name, func(t *testing.T) { + fs := afero.NewOsFs() + root := t.TempDir() + + // Random content, so the blob stays large enough to stall + // part-way through its compressed, encrypted body. + cfg, storer, snapshotID, _ := backupOneFile(context.Background(), + t, fs, root, "a.bin", randomBytes(t, 256*1024), 0o644) + + scratch := filepath.Join(root, "scratch") + require.NoError(t, fs.MkdirAll(scratch, 0o755)) + t.Setenv("TMPDIR", scratch) + + stalling := newStallingBlobStorer(storer) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + v := newRestoreVaultik(ctx, cfg, stalling, fs) + v.cancel = cancel + + var opErr error + + stop := v.StartOperation(func() { + opErr = op(v, snapshotID, filepath.Join(root, "restored")) + }) + + select { + case <-stalling.stalled: + case <-time.After(30 * time.Second): + t.Fatal("the operation never reached a blob read") + } + + require.NotEmpty(t, scratchEntries(t, scratch), + "expected decrypted scratch files to exist mid-read") + + stopCtx, stopCancel := context.WithTimeout( + context.Background(), 10*time.Second) + defer stopCancel() + + require.True(t, stop(stopCtx), + "stop timed out; the stalled blob read was not interrupted") + require.ErrorIs(t, opErr, context.Canceled) + require.Empty(t, scratchEntries(t, scratch), + "decrypted scratch files remained after the interrupt") + }) + } +}