internal/vaultik/restore.go

@@ -691,6 +691,8 @@ func (v *Vaultik) downloadSnapshotDB(
 		return nil, "", fmt.Errorf("downloading %s: %w", dbKey, err)
 	}
 
+	reader = newCloseOnCancelReader(v.ctx, reader)
+
 	defer func() { _ = reader.Close() }()
 
 	// Decrypt and decompress straight from the storage stream, then stream

internal/vaultik/restore_interrupt_test.go

@@ -207,30 +207,37 @@ func TestRestoreSkipErrorsCancelDuringBlobDownload(t *testing.T) {
 		"the cancel was reported as a failed file")
 }
 
-// stallingBlobStorer wraps a Storer so that every blob download hands
+// stallingStorer wraps a Storer so that each download of a key chosen by
-// over the first half of the blob and then stalls, ignoring the context,
+// stalls hands over the first half of the object and then stalls,
-// as a file:// blob on a hung USB stick or network share does. Only
+// ignoring the context, as a file:// read on a hung USB stick or network
-// closing the returned reader ends the stall.
+// share does. Only closing the returned reader ends the stall.
-type stallingBlobStorer struct {
+type stallingStorer struct {
 	storage.Storer
 
+	stalls  func(key string) bool
 	once    sync.Once
 	stalled chan struct{}
 }
 
-func newStallingBlobStorer(inner storage.Storer) *stallingBlobStorer {
+func newStallingStorer(
-	return &stallingBlobStorer{Storer: inner, stalled: make(chan struct{})}
+	inner storage.Storer, stalls func(key string) bool,
+) *stallingStorer {
+	return &stallingStorer{
+		Storer:  inner,
+		stalls:  stalls,
+		stalled: make(chan struct{}),
+	}
 }
 
-func (s *stallingBlobStorer) Get(
+func (s *stallingStorer) Get(
 	ctx context.Context, key string,
 ) (io.ReadCloser, error) {
 	rc, err := s.Storer.Get(ctx, key)
-	if err != nil || !strings.HasPrefix(key, "blobs/") {
+	if err != nil || !s.stalls(key) {
 		return rc, err
 	}
 
-	blob, err := io.ReadAll(rc)
+	object, err := io.ReadAll(rc)
 	_ = rc.Close()
 
 	if err != nil {
@@ -242,7 +249,7 @@ func (s *stallingBlobStorer) Get(
 	pr, pw := io.Pipe()
 
 	go func() {
-		_, _ = pw.Write(blob[:len(blob)/2])
+		_, _ = pw.Write(object[:len(object)/2])
 
 		s.once.Do(func() { close(s.stalled) })
 	}()
@@ -250,17 +257,26 @@ func (s *stallingBlobStorer) Get(
 	return pr, nil
 }
 
-// TestInterruptDuringStalledBlobReadCleansTempDir stops a restore and a
+// TestInterruptDuringStalledReadCleansTempDir stops a restore and a deep
-// deep verify while a blob read is stalled on a source that ignores the
+// verify while a read of a blob or of the snapshot database is stalled
-// context. Stopping must end that read, so the operation unwinds and
+// on a source that ignores the context. Stopping must end that read, so
-// removes its decrypted scratch files well within the shutdown timeout
+// the operation unwinds and removes its decrypted scratch files well
-// (issue #325).
+// within the shutdown timeout (issues #325 and #356).
 //
 // 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) {
+func TestInterruptDuringStalledReadCleansTempDir(t *testing.T) {
 	log.Initialize(log.Config{})
 
+	stalledReads := map[string]func(key string) bool{
+		"blob": func(key string) bool {
+			return strings.HasPrefix(key, "blobs/")
+		},
+		"snapshot database": func(key string) bool {
+			return strings.HasSuffix(key, "/db.zst.age")
+		},
+	}
+
 	operations := map[string]func(v *Vaultik, snapshotID, targetDir string) error{
 		"restore": func(v *Vaultik, snapshotID, targetDir string) error {
 			return v.Restore(&RestoreOptions{
@@ -273,52 +289,58 @@ func TestInterruptDuringStalledBlobReadCleansTempDir(t *testing.T) {
 		},
 	}
 
-	for name, op := range operations {
+	for readName, stalls := range stalledReads {
-		t.Run(name, func(t *testing.T) {
+		for opName, op := range operations {
-			fs := afero.NewOsFs()
+			t.Run(readName+" read during "+opName, func(t *testing.T) {
-			root := t.TempDir()
+				fs := afero.NewOsFs()
+				root := t.TempDir()
 
-			// Random content, so the blob stays large enough to stall
+				// Random content, so the blob stays large enough to stall
-			// part-way through its compressed, encrypted body.
+				// part-way through its compressed, encrypted body.
-			cfg, storer, snapshotID, _ := backupOneFile(context.Background(),
+				cfg, storer, snapshotID, _ := backupOneFile(context.Background(),
-				t, fs, root, "a.bin", randomBytes(t, 256*1024), 0o644)
+					t, fs, root, "a.bin", randomBytes(t, 256*1024), 0o644)
 
-			scratch := filepath.Join(root, "scratch")
+				scratch := filepath.Join(root, "scratch")
-			require.NoError(t, fs.MkdirAll(scratch, 0o755))
+				require.NoError(t, fs.MkdirAll(scratch, 0o755))
-			t.Setenv("TMPDIR", scratch)
+				t.Setenv("TMPDIR", scratch)
 
-			stalling := newStallingBlobStorer(storer)
+				stalling := newStallingStorer(storer, stalls)
 
-			ctx, cancel := context.WithCancel(context.Background())
+				ctx, cancel := context.WithCancel(context.Background())
-			defer cancel()
+				defer cancel()
 
-			v := newRestoreVaultik(ctx, cfg, stalling, fs)
+				v := newRestoreVaultik(ctx, cfg, stalling, fs)
-			v.cancel = cancel
+				v.cancel = cancel
 
-			var opErr error
+				var opErr error
 
-			stop := v.StartOperation(func() {
+				stop := v.StartOperation(func() {
-				opErr = op(v, snapshotID, filepath.Join(root, "restored"))
+					opErr = op(v, snapshotID, filepath.Join(root, "restored"))
+				})
+
+				select {
+				case <-stalling.stalled:
+				case <-time.After(30 * time.Second):
+					t.Fatal("the operation never reached the stalled read")
+				}
+
+				// The snapshot database's temp directory is created just
+				// after its read starts, so it can appear after the stall.
+				require.Eventually(t, func() bool {
+					return len(scratchEntries(t, scratch)) > 0
+				}, 10*time.Second, 10*time.Millisecond,
+					"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 read was not interrupted")
+				require.ErrorIs(t, opErr, context.Canceled)
+				require.Empty(t, scratchEntries(t, scratch),
+					"decrypted scratch files remained after the interrupt")
 			})
-
+		}
-			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")
-		})
 	}
 }

internal/vaultik/verify.go

@@ -246,6 +246,8 @@ func (v *Vaultik) downloadVerifiedSnapshotDB(
 			fmt.Errorf("failed to download database: %w", err))
 	}
 
+	dbReader = newCloseOnCancelReader(v.ctx, dbReader)
+
 	defer func() { _ = dbReader.Close() }()
 
 	tdb, err := v.decryptAndLoadDatabase(dbReader, identities)