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
-// 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 {
+// stallingStorer wraps a Storer so that each download of a key chosen by
+// stalls hands over the first half of the object and then stalls,
+// ignoring the context, as a file:// read on a hung USB stick or network
+// share does. Only closing the returned reader ends the stall.
+type stallingStorer struct {
 	storage.Storer
 
+	stalls  func(key string) bool
 	once    sync.Once
 	stalled chan struct{}
 }
 
-func newStallingBlobStorer(inner storage.Storer) *stallingBlobStorer {
-	return &stallingBlobStorer{Storer: inner, stalled: make(chan struct{})}
+func newStallingStorer(
+	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
-// 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).
+// TestInterruptDuringStalledReadCleansTempDir stops a restore and a deep
+// verify while a read of a blob or of the snapshot database 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 (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,8 +289,9 @@ func TestInterruptDuringStalledBlobReadCleansTempDir(t *testing.T) {
 		},
 	}
 
-	for name, op := range operations {
-		t.Run(name, func(t *testing.T) {
+	for readName, stalls := range stalledReads {
+		for opName, op := range operations {
+			t.Run(readName+" read during "+opName, func(t *testing.T) {
 				fs := afero.NewOsFs()
 				root := t.TempDir()
 
@@ -287,7 +304,7 @@ func TestInterruptDuringStalledBlobReadCleansTempDir(t *testing.T) {
 				require.NoError(t, fs.MkdirAll(scratch, 0o755))
 				t.Setenv("TMPDIR", scratch)
 
-			stalling := newStallingBlobStorer(storer)
+				stalling := newStallingStorer(storer, stalls)
 
 				ctx, cancel := context.WithCancel(context.Background())
 				defer cancel()
@@ -304,10 +321,14 @@ func TestInterruptDuringStalledBlobReadCleansTempDir(t *testing.T) {
 				select {
 				case <-stalling.stalled:
 				case <-time.After(30 * time.Second):
-				t.Fatal("the operation never reached a blob read")
+					t.Fatal("the operation never reached the stalled read")
 				}
 
-			require.NotEmpty(t, scratchEntries(t, scratch),
+				// 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(
@@ -315,10 +336,11 @@ func TestInterruptDuringStalledBlobReadCleansTempDir(t *testing.T) {
 				defer stopCancel()
 
 				require.True(t, stop(stopCtx),
-				"stop timed out; the stalled blob read was not interrupted")
+					"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")
 			})
 		}
 	}
+}

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)