vaultik/internal/vaultik/restore.go

package vaultik

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"math"
	"os"
	"path/filepath"
	"strings"
	"time"

	"filippo.io/age"
	"github.com/spf13/afero"
	"sneak.berlin/go/vaultik/internal/blobgen"
	"sneak.berlin/go/vaultik/internal/database"
	"sneak.berlin/go/vaultik/internal/log"
	"sneak.berlin/go/vaultik/internal/snapshot"
	"sneak.berlin/go/vaultik/internal/types"
)

// Sentinel errors for restore failures.
var (
	errFilesFailedRestore    = errors.New("file(s) failed to restore")
	errFilesFailedVerify     = errors.New("files failed verification")
	errDecryptionKeyRequired = errors.New(
		"decryption key required for restore\n\n" +
			"Set the VAULTIK_AGE_SECRET_KEY environment variable to your " +
			"age private key:\n" +
			"  export VAULTIK_AGE_SECRET_KEY='AGE-SECRET-KEY-...'")
	errBlobMissingFromIndex     = errors.New("blob hash missing from blob index")
	errChunkNotInAnyBlob        = errors.New("chunk not found in any blob")
	errBlobIDNotInHashIndex     = errors.New("blob id missing from hash index")
	errShortChunkRead           = errors.New("short read")
	errRestorePathEscapesTarget = errors.New(
		"refusing to restore path outside the target directory")
	errTrailingRestoreData = errors.New(
		"restored file has trailing data after its last chunk")
	errRestoreIncomplete = errors.New(
		"restore loop ended with files still pending")
	errSnapshotDBMismatch = errors.New(
		"decrypted database is not the requested snapshot")
)

// snapshotDBFilename is the name the decrypted snapshot database is
// written under inside its private temp directory.
const snapshotDBFilename = "snapshot.db"

// restoreDirMode is the permission mode for directories created while
// restoring (parent directories and the target root; restored
// directories themselves get their stored mode).
const restoreDirMode = 0o755

// restoreFileMode is the restrictive mode a regular file is created with
// during restore. Content is written while the file holds this mode; the
// stored mode is applied only after the file is fully written and closed,
// so a file whose stored mode is restrictive is never briefly readable by
// other local users while its content is being written.
const restoreFileMode = 0o600

// sweepIntervalDivisor sets the sweeper threshold to one N-th of the
// configured blob size limit.
const sweepIntervalDivisor = 100

// restoreStatusInterval is how often periodic progress lines are
// printed during restore and verify.
const restoreStatusInterval = 15 * time.Second

// RestoreOptions contains options for the restore operation
type RestoreOptions struct {
	SnapshotID string
	TargetDir  string
	Paths      []string // Optional paths to restore (empty = all)
	Verify     bool     // Verify restored files by checking chunk hashes
	SkipErrors bool     // Continue past file-restore errors instead of aborting
}

// RestoreResult contains statistics from a restore operation
type RestoreResult struct {
	FilesRestored   int
	BytesRestored   int64
	BlobsDownloaded int
	BytesDownloaded int64
	Duration        time.Duration
	// Verification results (only populated if Verify option is set)
	FilesVerified int
	BytesVerified int64
	FilesFailed   int
	FailedFiles   []string // Paths of files that failed verification
}

// Restore restores files from a snapshot to the target directory
func (v *Vaultik) Restore(opts *RestoreOptions) error {
	startTime := time.Now()

	identity, err := v.prepareRestoreIdentity()
	if err != nil {
		return err
	}

	log.Info("Starting restore operation",
		"snapshot_id", opts.SnapshotID,
		"target_dir", opts.TargetDir,
		"paths", opts.Paths,
	)

	// Step 1: Download and decrypt the snapshot metadata database
	log.Info("Downloading snapshot metadata...")

	tempDB, tempDir, err := v.downloadSnapshotDB(opts.SnapshotID, identity)
	if err != nil {
		return fmt.Errorf("downloading snapshot database: %w", err)
	}

	defer func() {
		err := tempDB.Close()
		if err != nil {
			log.Debug("Failed to close temp database", "error", err)
		}
		// Remove the whole private directory, so the decrypted database
		// and any SQLite side files it produced are gone on every path.
		err = v.Fs.RemoveAll(tempDir)
		if err != nil {
			log.Debug("Failed to remove temp database directory", "error", err)
		}
	}()

	repos := database.NewRepositories(tempDB)

	// Step 2: Get list of files to restore
	files, err := v.getFilesToRestore(v.ctx, repos, opts.Paths)
	if err != nil {
		return fmt.Errorf("getting files to restore: %w", err)
	}

	if len(files) == 0 {
		log.Warn("No files found to restore")
		v.UI.Warningf("No files found to restore.")

		return nil
	}

	log.Info("Found files to restore", "count", len(files))
	v.UI.Infof("Found %s files to restore.", v.UI.Count(len(files)))

	// Step 3: Create target directory
	err = v.Fs.MkdirAll(opts.TargetDir, restoreDirMode)
	if err != nil {
		return fmt.Errorf("creating target directory: %w", err)
	}

	// Step 4: Build a map of chunks to blobs for efficient restoration
	chunkToBlobMap, err := v.buildChunkToBlobMap(v.ctx, repos)
	if err != nil {
		return fmt.Errorf("building chunk-to-blob map: %w", err)
	}

	// Step 5: Restore files
	result, err := v.restoreAllFiles(files, repos, opts, identity, chunkToBlobMap)
	if err != nil {
		return err
	}

	result.Duration = time.Since(startTime)

	log.Info("Restore complete",
		"files_restored", result.FilesRestored,
		"bytes_restored", ubytes(result.BytesRestored),
		"blobs_downloaded", result.BlobsDownloaded,
		"bytes_downloaded", ubytes(result.BytesDownloaded),
		"duration", result.Duration,
	)

	v.UI.Completef("Restored %s files (%s) in %s.",
		v.UI.Count(result.FilesRestored),
		v.UI.Size(result.BytesRestored),
		v.UI.Duration(result.Duration),
	)

	return v.finishRestore(repos, files, opts, result)
}

// finishRestore emits the post-restore warnings, runs optional
// verification, and converts any failed-file count into an error.
func (v *Vaultik) finishRestore(
	repos *database.Repositories,
	files []*database.File,
	opts *RestoreOptions,
	result *RestoreResult,
) error {
	if os.Geteuid() != 0 {
		v.UI.Warningf("Restore did not preserve file ownership: chown(2) " +
			"requires root. Re-run as root (e.g. with sudo) if you need " +
			"original UID/GID preserved.")
	}

	if result.FilesFailed > 0 {
		v.UI.Warningf("%d file(s) failed to restore:", result.FilesFailed)

		for _, path := range result.FailedFiles {
			v.UI.Detailf("%s", v.UI.Path(path))
		}
	}

	// Run verification if requested
	if opts.Verify {
		err := v.handleRestoreVerification(repos, files, opts, result)
		if err != nil {
			return err
		}
	}

	if result.FilesFailed > 0 {
		return fmt.Errorf("%d %w", result.FilesFailed, errFilesFailedRestore)
	}

	return nil
}

// prepareRestoreIdentity validates that an age secret key is configured
// and parses it.
//
//nolint:ireturn // age.Identity is the decryption abstraction by design
func (v *Vaultik) prepareRestoreIdentity() (age.Identity, error) {
	if v.Config.AgeSecretKey == "" {
		return nil, errDecryptionKeyRequired
	}

	identity, err := age.ParseX25519Identity(v.Config.AgeSecretKey)
	if err != nil {
		return nil, fmt.Errorf("parsing age secret key: %w", err)
	}

	return identity, nil
}

// restoreAllFiles processes files in blob-locality order: drain every
// file whose blob set is on disk, download the missing blobs for the
// pending file with the smallest uncached count, repeat. This keeps
// peak cache occupancy near 1 even on snapshots whose path order
// interleaves blobs, and lets the sweeper free each blob the moment
// its file set is exhausted.
func (v *Vaultik) restoreAllFiles(
	files []*database.File,
	repos *database.Repositories,
	opts *RestoreOptions,
	identity age.Identity,
	chunkToBlobMap map[string]*database.BlobChunk,
) (*RestoreResult, error) {
	result := &RestoreResult{}

	// The restore-side blob cache is unbounded — restores may read any
	// blob many times across deduplicated files and we want to avoid
	// re-downloading until we can prove a blob is no longer needed.
	// Cleanup is driven by the sweeper below, not by LRU.
	blobCache, err := newBlobDiskCache(math.MaxInt64)
	if err != nil {
		return nil, fmt.Errorf("creating blob cache: %w", err)
	}

	if v.restoreCacheObserver != nil {
		v.restoreCacheObserver(blobCache)
	}

	defer func() {
		if v.restoreCacheObserver != nil {
			v.restoreCacheObserver(blobCache)
		}

		_ = blobCache.Close()
	}()

	// Per-restore sweep state: every blob_size_limit/100 bytes written,
	// scan the cache and delete any blob whose remaining file references
	// are all already restored.
	sweeper := newRestoreSweeper(v.ctx, repos, blobCache,
		v.Config.BlobSizeLimit.Int64()/sweepIntervalDivisor)

	blobByHash, blobIDToHash, err := v.buildBlobIndexes(repos)
	if err != nil {
		return nil, err
	}

	plan, err := newRestorePlan(v.ctx, repos, files, chunkToBlobMap, blobIDToHash)
	if err != nil {
		return nil, fmt.Errorf("building restore plan: %w", err)
	}

	filesByID, totalBytesExpected := indexRestoreFiles(files)

	v.UI.Beginf("Restoring %s files (%s) to %s.",
		v.UI.Count(len(files)),
		v.UI.Size(totalBytesExpected),
		v.UI.Path(opts.TargetDir))

	session := &restoreSession{
		v:              v,
		ctx:            v.ctx,
		repos:          repos,
		opts:           opts,
		identity:       identity,
		chunkToBlobMap: chunkToBlobMap,
		blobByHash:     blobByHash,
		blobIDToHash:   blobIDToHash,
		blobCache:      blobCache,
		sweeper:        sweeper,
		result:         result,
		runningAsRoot:  os.Geteuid() == 0,
	}

	err = v.runRestoreLoop(session, plan, filesByID, totalBytesExpected)
	if err != nil {
		return nil, err
	}

	return result, nil
}

// runRestoreLoop drains the restore plan: restore files as their blobs
// become available, download the next blob set when nothing is ready,
// and emit periodic progress.
func (v *Vaultik) runRestoreLoop(
	session *restoreSession, plan *restorePlan,
	filesByID map[types.FileID]*database.File, totalBytesExpected int64,
) error {
	// Periodic progress output, matching the snapshot create cadence.
	startTime := time.Now()
	lastStatusTime := startTime
	processed := 0
	totalFiles := len(filesByID)

	for plan.hasPending() {
		if v.ctx.Err() != nil {
			return v.ctx.Err()
		}

		fileID, ready := plan.popReady()
		if !ready {
			downloaded, err := session.downloadNextBlobSet(plan)
			if err != nil {
				return err
			}

			if !downloaded {
				break
			}

			continue
		}

		file := filesByID[fileID]

		err := session.restoreFile(file)
		if err != nil {
			err = v.handleRestoreFileError(
				plan, session.opts, session.result, file, fileID, err)
			if err != nil {
				return err
			}

			continue
		}

		// Record the file as restored so the sweeper can free blobs
		// once all referencing files are done, and drop it from the
		// plan's indexes so future picks ignore it.
		session.sweeper.fileRestored(fileID.String())
		plan.finishFile(fileID)

		processed++

		v.restoreProgressTick(processed, totalFiles,
			session.result.BytesRestored,
			totalBytesExpected, startTime, &lastStatusTime)
	}

	// The loop above stops as soon as nothing is ready and nothing more
	// can be downloaded. If files still remain, they were abandoned
	// rather than restored; fail loudly instead of reporting success.
	if plan.hasPending() {
		return errRestoreIncomplete
	}

	return nil
}

// downloadNextBlobSet is invoked when no file is fully cache-served.
// It first frees any blobs whose file sets are exhausted — without
// this, the blob whose last file we just finished would still be
// cached when we Put the next one, briefly pushing peak occupancy from
// 1 to 2. It then picks the pending file with the smallest uncached
// blob set and downloads its blobs; after each blob lands, the plan
// moves any pending file whose set just emptied onto the ready queue.
// Returns false when nothing is pending download (the caller stops).
func (s *restoreSession) downloadNextBlobSet(plan *restorePlan) (bool, error) {
	s.sweeper.sweep()

	next, ok := plan.pickNextDownload()
	if !ok {
		return false, nil
	}

	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])
		}

		err := s.downloadBlobToCache(hash, blob.CompressedSize)
		if err != nil {
			return false, fmt.Errorf("downloading blob %s: %w", hash[:16], err)
		}

		s.result.BlobsDownloaded++
		s.result.BytesDownloaded += blob.CompressedSize

		plan.markBlobCached(hash)
	}

	return true, nil
}

// indexRestoreFiles indexes files by ID for plan lookups and sums the
// expected byte total for percentage / ETA arithmetic.
func indexRestoreFiles(
	files []*database.File,
) (map[types.FileID]*database.File, int64) {
	filesByID := make(map[types.FileID]*database.File, len(files))

	var totalBytesExpected int64

	for _, f := range files {
		filesByID[f.ID] = f
		totalBytesExpected += f.Size
	}

	return filesByID, totalBytesExpected
}

// buildBlobIndexes pre-fetches every blob row once so chunk extraction
// can map a blob_id to its hash without a DB round-trip per chunk.
func (v *Vaultik) buildBlobIndexes(
	repos *database.Repositories,
) (map[string]*database.Blob, map[string]string, error) {
	blobsByID, err := repos.Blobs.GetAll(v.ctx)
	if err != nil {
		return nil, nil, fmt.Errorf("fetching blob index: %w", err)
	}

	blobIDToHash := make(map[string]string, len(blobsByID))

	blobByHash := make(map[string]*database.Blob, len(blobsByID))
	for id, blob := range blobsByID {
		hash := blob.Hash.String()
		blobIDToHash[id] = hash
		blobByHash[hash] = blob
	}

	return blobByHash, blobIDToHash, nil
}

// restoreProgressTick emits the periodic UI status line and structured
// progress log during the restore loop.
func (v *Vaultik) restoreProgressTick(
	processed, totalFiles int, bytesRestored, totalBytesExpected int64,
	startTime time.Time, lastStatusTime *time.Time,
) {
	if time.Since(*lastStatusTime) >= restoreStatusInterval {
		v.printRestoreProgress(
			processed, totalFiles, bytesRestored,
			totalBytesExpected, startTime)

		*lastStatusTime = time.Now()
	}

	// Structured progress log for --verbose / JSON consumers.
	if processed%progressLogEvery == 0 || processed == totalFiles {
		log.Info("Restore progress",
			"files", fmt.Sprintf("%d/%d", processed, totalFiles),
			"bytes", ubytes(bytesRestored),
		)
	}
}

// handleRestoreFileError records a per-file restore failure: fatal unless
// --skip-errors is set, in which case the file is counted as failed and
// dropped from the plan.
func (v *Vaultik) handleRestoreFileError(
	plan *restorePlan, opts *RestoreOptions, result *RestoreResult,
	file *database.File, fileID types.FileID, err error,
) error {
	log.Error("Failed to restore file", "path", file.Path, "error", err)

	if !opts.SkipErrors {
		return fmt.Errorf(
			"restoring %s: %w (pass --skip-errors to continue past "+
				"restore failures)", file.Path, err)
	}

	v.UI.Errorf("Failed to restore %s: %v. Skipping (--skip-errors).",
		v.UI.Path(file.Path.String()), err)

	result.FilesFailed++
	result.FailedFiles = append(result.FailedFiles, file.Path.String())

	plan.finishFile(fileID)

	return nil
}

// printRestoreProgress emits a periodic restore-phase status line via
// the UI writer, mirroring scanner.printProcessingProgress so the two
// long-running commands have the same on-screen rhythm.
func (v *Vaultik) printRestoreProgress(
	filesDone, totalFiles int, bytesDone, totalBytes int64, startTime time.Time,
) {
	v.printPhaseProgress("Restore", "restore",
		filesDone, totalFiles, bytesDone, totalBytes, startTime)
}

// printPhaseProgress emits a periodic status line for a long-running
// phase (restore or verify) so user-facing pacing is uniform.
func (v *Vaultik) printPhaseProgress(
	title, phase string,
	filesDone, totalFiles int, bytesDone, totalBytes int64, startTime time.Time,
) {
	elapsed := time.Since(startTime)
	pct := float64(bytesDone) / float64(totalBytes) * percentScale
	byteRate := float64(bytesDone) / elapsed.Seconds()
	fileRate := float64(filesDone) / elapsed.Seconds()

	remainingBytes := totalBytes - bytesDone

	var eta time.Duration
	if byteRate > 0 && remainingBytes > 0 {
		eta = time.Duration(float64(remainingBytes)/byteRate) * time.Second
	}

	if eta > 0 {
		v.UI.Progressf("%s: %s/%s files (%s), %s/%s, %s, %.0f files/sec, "+
			"%s elapsed: %s, %s ETA: %s (est remain %s).",
			title,
			v.UI.Count(filesDone),
			v.UI.Count(totalFiles),
			v.UI.Percent(pct),
			v.UI.Size(bytesDone),
			v.UI.Size(totalBytes),
			v.UI.Speed(byteRate),
			fileRate,
			phase,
			v.UI.Duration(elapsed),
			phase,
			v.UI.Time(time.Now().Add(eta)),
			v.UI.Duration(eta))

		return
	}

	v.UI.Progressf("%s: %s/%s files (%s), %s/%s, %s, %.0f files/sec, "+
		"%s elapsed: %s.",
		title,
		v.UI.Count(filesDone),
		v.UI.Count(totalFiles),
		v.UI.Percent(pct),
		v.UI.Size(bytesDone),
		v.UI.Size(totalBytes),
		v.UI.Speed(byteRate),
		fileRate,
		phase,
		v.UI.Duration(elapsed))
}

// handleRestoreVerification runs post-restore verification if requested
func (v *Vaultik) handleRestoreVerification(
	repos *database.Repositories,
	files []*database.File,
	opts *RestoreOptions,
	result *RestoreResult,
) error {
	err := v.verifyRestoredFiles(v.ctx, repos, files, opts.TargetDir, result)
	if err != nil {
		return fmt.Errorf("verification failed: %w", err)
	}

	if result.FilesFailed > 0 {
		v.UI.Errorf("Verification failed: %s files did not match expected checksums.",
			v.UI.Count(result.FilesFailed))

		for _, path := range result.FailedFiles {
			v.UI.Detailf("%s", v.UI.Path(path))
		}

		return fmt.Errorf("%d %w", result.FilesFailed, errFilesFailedVerify)
	}

	v.UI.Completef("Verified %s files (%s).",
		v.UI.Count(result.FilesVerified),
		v.UI.Size(result.BytesVerified))

	return nil
}

// downloadSnapshotDB downloads and decrypts the snapshot metadata
// database. The identifier is resolved to the snapshot's remote key: a
// human ID is hashed, and a remote key (or its abbreviation, as printed
// for a remote-only snapshot) is used as-is, so a host with no local
// index can restore the snapshots it can only see on the store.
func (v *Vaultik) downloadSnapshotDB(
	snapshotID string, identity age.Identity,
) (*database.DB, string, error) {
	remoteKey, err := v.resolveSnapshotRemoteKey(snapshotID)
	if err != nil {
		return nil, "", err
	}

	// Download encrypted database from storage
	dbKey := fmt.Sprintf("metadata/%s/db.zst.age", remoteKey)

	reader, err := v.Storage.Get(v.ctx, dbKey)
	if err != nil {
		return nil, "", fmt.Errorf("downloading %s: %w", dbKey, err)
	}

	defer func() { _ = reader.Close() }()

	// Read all data
	encryptedData, err := io.ReadAll(reader)
	if err != nil {
		return nil, "", fmt.Errorf("reading encrypted data: %w", err)
	}

	log.Debug("Downloaded encrypted database",
		"size", ubytes(int64(len(encryptedData))))

	// Decrypt and decompress using blobgen.Reader
	blobReader, err := blobgen.NewReader(bytes.NewReader(encryptedData), identity)
	if err != nil {
		return nil, "", fmt.Errorf("creating decryption reader: %w", err)
	}

	defer func() { _ = blobReader.Close() }()

	// Read the binary SQLite database
	dbData, err := io.ReadAll(blobReader)
	if err != nil {
		return nil, "", fmt.Errorf("decrypting and decompressing: %w", err)
	}

	log.Debug("Decrypted database", "size", ubytes(int64(len(dbData))))

	db, tempDir, err := v.materializeSnapshotDB(dbData)
	if err != nil {
		return nil, "", err
	}

	// Confirm the decrypted database really is the snapshot named by
	// remoteKey before any files are read from it. On mismatch, close the
	// database and remove its private directory so nothing is left behind.
	err = v.verifySnapshotDBIdentity(db, snapshotID, remoteKey)
	if err != nil {
		_ = db.Close()
		_ = v.Fs.RemoveAll(tempDir)

		return nil, "", err
	}

	return db, tempDir, nil
}

// verifySnapshotDBIdentity confirms the decrypted metadata database really
// is the snapshot named by remoteKey. age decryption proves the database
// is readable, not that the object served at
// metadata/<remoteKey>/db.zst.age is the snapshot that was requested: an
// attacker who swaps in another valid db.zst.age (which needs no key
// material) would otherwise redirect restore and deep verify to a
// different snapshot's contents. The exported per-snapshot database holds
// exactly one snapshot row, and a snapshot's remote key is derived from
// that row's ID, so the database is the requested one exactly when its
// sole snapshot hashes back to remoteKey. Comparing the requested
// identifier directly would not do: it may be a remote-key prefix a
// recovery host uses in place of a human snapshot ID it cannot know.
func (v *Vaultik) verifySnapshotDBIdentity(
	db *database.DB, requested, remoteKey string,
) error {
	repos := database.NewRepositories(db)

	snap, err := repos.Snapshots.GetOnlySnapshot(v.ctx)
	if err != nil {
		return fmt.Errorf("checking identity of database for %s: %w", requested, err)
	}

	if snapshot.RemoteSnapshotKey(snap.ID.String()) != remoteKey {
		return fmt.Errorf("%w: requested %s but the database is snapshot %s",
			errSnapshotDBMismatch, requested, snap.ID)
	}

	return nil
}

// materializeSnapshotDB writes the decrypted snapshot database bytes into
// a fresh private (0700) temp directory and opens the file read-only. On
// any failure it removes the directory before returning, so no decrypted
// metadata is left on disk when the open is interrupted or the payload is
// damaged. On success the returned directory is the caller's to remove.
func (v *Vaultik) materializeSnapshotDB(
	dbData []byte,
) (*database.DB, string, error) {
	tempDir, err := afero.TempDir(v.Fs, "", "vaultik-restore-")
	if err != nil {
		return nil, "", fmt.Errorf("creating temp directory: %w", err)
	}

	success := false

	defer func() {
		if !success {
			_ = v.Fs.RemoveAll(tempDir)
		}
	}()

	dbPath := filepath.Join(tempDir, snapshotDBFilename)

	err = afero.WriteFile(v.Fs, dbPath, dbData, restoreFileMode)
	if err != nil {
		return nil, "", fmt.Errorf("writing database file: %w", err)
	}

	log.Debug("Created restore database", "path", dbPath)

	db, err := database.OpenReadOnly(v.ctx, dbPath)
	if err != nil {
		return nil, "", fmt.Errorf("opening restore database: %w", err)
	}

	success = true

	return db, tempDir, nil
}

// getFilesToRestore returns the list of files to restore based on path filters
func (v *Vaultik) getFilesToRestore(
	ctx context.Context, repos *database.Repositories, pathFilters []string,
) ([]*database.File, error) {
	// If no filters, get all files
	if len(pathFilters) == 0 {
		return repos.Files.ListAll(ctx)
	}

	// Get files matching the path filters
	var result []*database.File

	seen := make(map[string]bool)

	for _, filter := range pathFilters {
		// Normalize the filter path
		filter = filepath.Clean(filter)

		// Get files with this prefix
		files, err := repos.Files.ListByPrefix(ctx, filter)
		if err != nil {
			return nil, fmt.Errorf("listing files with prefix %s: %w", filter, err)
		}

		for _, file := range files {
			if !seen[file.ID.String()] {
				seen[file.ID.String()] = true
				result = append(result, file)
			}
		}
	}

	return result, nil
}

// buildChunkToBlobMap creates a mapping from chunk hash to blob information
func (v *Vaultik) buildChunkToBlobMap(
	ctx context.Context, repos *database.Repositories,
) (map[string]*database.BlobChunk, error) {
	// Query all blob_chunks
	query := `SELECT blob_id, chunk_hash, offset, length FROM blob_chunks`

	rows, err := repos.DB().Conn().QueryContext(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("querying blob_chunks: %w", err)
	}

	defer func() { _ = rows.Close() }()

	result := make(map[string]*database.BlobChunk)

	for rows.Next() {
		var (
			bc                      database.BlobChunk
			blobIDStr, chunkHashStr string
		)

		err = rows.Scan(&blobIDStr, &chunkHashStr, &bc.Offset, &bc.Length)
		if err != nil {
			return nil, fmt.Errorf("scanning blob_chunk: %w", err)
		}

		blobID, err := types.ParseBlobID(blobIDStr)
		if err != nil {
			return nil, fmt.Errorf("parsing blob ID: %w", err)
		}

		bc.BlobID = blobID
		bc.ChunkHash = types.ChunkHash(chunkHashStr)
		result[chunkHashStr] = &bc
	}

	return result, rows.Err()
}

// restoreSession holds every piece of per-restore state shared by the
// restore-time methods. Each restore builds one of these from the
// snapshot's metadata and then drives the file loop through methods on
// it. Keeping this state on the struct rather than threading it
// through every function signature keeps the inner-loop call sites
// readable: restoreFile(file) instead of a ten-argument helper.
type restoreSession struct {
	v              *Vaultik
	ctx            context.Context //nolint:containedctx // per-restore state by design
	repos          *database.Repositories
	opts           *RestoreOptions
	identity       age.Identity
	chunkToBlobMap map[string]*database.BlobChunk
	blobByHash     map[string]*database.Blob
	blobIDToHash   map[string]string
	blobCache      *blobDiskCache
	sweeper        *restoreSweeper
	result         *RestoreResult
	// runningAsRoot gates chown(2). On every Unix-ish kernel, only
	// root can chown a file to an arbitrary UID/GID — non-root chown
	// always fails with EPERM. Attempting it anyway produces N
	// guaranteed-failed syscalls + N noisy debug lines, so we skip
	// the call entirely as non-root and emit one warning at the end
	// of the restore explaining that ownership was not preserved.
	runningAsRoot bool
}

// containedRestorePath resolves rel — a path read from the snapshot
// database — to its location under targetDir and confirms the write will
// stay inside the target.
//
// age decryption proves a snapshot is readable, not that it is honest, so
// every stored path is treated as hostile. rel is rejected unless
// filepath.IsLocal accepts it once the leading separator is stripped:
// stored paths are absolute and the join to targetDir drops that
// separator, so "/etc/passwd" is judged as the relative "etc/passwd" it
// becomes on disk. This bars "..", absolute, and empty paths.
//
// A stored symlink whose target points outside the tree is still honest
// (and restored verbatim), but a later entry must not be written through
// it. Each existing ancestor directory below the target is therefore
// Lstat'ed and a symlink among them is refused. The leaf itself is not
// traversed: honest snapshots restore symlinks at leaf positions, and the
// unique-path constraint keeps a leaf from being both a symlink and a
// regular file. The target directory itself may be a symlink; only
// components below it are checked.
func containedRestorePath(fs afero.Fs, targetDir, rel string) (string, error) {
	local := strings.TrimPrefix(rel, string(filepath.Separator))
	if !filepath.IsLocal(local) {
		return "", fmt.Errorf("%w: %s", errRestorePathEscapesTarget, rel)
	}

	local = filepath.Clean(local)
	targetPath := filepath.Join(targetDir, local)

	relDir := filepath.Dir(local)
	if relDir == "." {
		return targetPath, nil
	}

	current := targetDir
	for component := range strings.SplitSeq(relDir, string(filepath.Separator)) {
		current = filepath.Join(current, component)

		info, err := lstatIfPossible(fs, current)
		if err != nil {
			if os.IsNotExist(err) {
				continue
			}

			return "", fmt.Errorf("checking restore path %s: %w", current, err)
		}

		if info.Mode()&os.ModeSymlink != 0 {
			return "", fmt.Errorf("%w: %s descends through symlink %s",
				errRestorePathEscapesTarget, rel, current)
		}
	}

	return targetPath, nil
}

// lstatIfPossible performs a symlink-aware stat when the filesystem
// supports it. afero.OsFs does; MemMapFs, which has no symlinks, reports
// that Lstat was not used and its result never carries ModeSymlink.
func lstatIfPossible(fs afero.Fs, name string) (os.FileInfo, error) {
	if lstater, ok := fs.(afero.Lstater); ok {
		info, _, err := lstater.LstatIfPossible(name)

		return info, err
	}

	return fs.Stat(name)
}

// restoreFile dispatches to the right per-kind restorer.
func (s *restoreSession) restoreFile(file *database.File) error {
	targetPath, err := containedRestorePath(
		s.v.Fs, s.opts.TargetDir, file.Path.String())
	if err != nil {
		return err
	}

	parentDir := filepath.Dir(targetPath)

	err = s.v.Fs.MkdirAll(parentDir, restoreDirMode)
	if err != nil {
		return fmt.Errorf("creating parent directory: %w", err)
	}

	if file.IsSymlink() {
		return s.restoreSymlink(file, targetPath)
	}

	if file.Mode&uint32(os.ModeDir) != 0 {
		return s.restoreDirectory(file, targetPath)
	}

	return s.restoreRegularFile(file, targetPath)
}

// restoreSymlink restores a symbolic link.
func (s *restoreSession) restoreSymlink(file *database.File, targetPath string) error {
	_ = s.v.Fs.Remove(targetPath)
	// afero.MemMapFs doesn't support symlinks, so route real-FS
	// symlinks through os.
	if _, ok := s.v.Fs.(*afero.OsFs); ok {
		err := os.Symlink(file.LinkTarget.String(), targetPath)
		if err != nil {
			return fmt.Errorf("creating symlink: %w", err)
		}
	} else {
		log.Debug("Symlink creation not supported on this filesystem",
			"path", file.Path, "target", file.LinkTarget)
	}

	s.result.FilesRestored++

	log.Debug("Restored symlink", "path", file.Path, "target", file.LinkTarget)

	return nil
}

// restoreDirectory restores a directory with its permissions, mtime,
// and (on real filesystems, with sufficient privileges) ownership.
func (s *restoreSession) restoreDirectory(
	file *database.File, targetPath string,
) error {
	err := s.v.Fs.MkdirAll(targetPath, os.FileMode(file.Mode))
	if err != nil {
		return fmt.Errorf("creating directory: %w", err)
	}

	// MkdirAll applies the process umask, so chmod to the exact stored
	// mode. A failure here is non-fatal.
	err = s.v.Fs.Chmod(targetPath, os.FileMode(file.Mode))
	if err != nil {
		log.Debug("Failed to set permissions", "path", targetPath, "error", err)
	}

	s.applyFileMetadata(file, targetPath)

	s.result.FilesRestored++

	return nil
}

// applyFileMetadata applies ownership (when running as root on a real
// filesystem) and mtime to a restored path. Permission mode is applied
// separately by each caller, with different failure handling, so it is
// not touched here. Failures are logged at debug level and do not abort
// the restore.
func (s *restoreSession) applyFileMetadata(file *database.File, targetPath string) {
	if s.runningAsRoot {
		if _, ok := s.v.Fs.(*afero.OsFs); ok {
			err := os.Chown(targetPath, int(file.UID), int(file.GID))
			if err != nil {
				log.Debug("Failed to set ownership", "path", targetPath, "error", err)
			}
		}
	}

	err := s.v.Fs.Chtimes(targetPath, file.MTime, file.MTime)
	if err != nil {
		log.Debug("Failed to set mtime", "path", targetPath, "error", err)
	}
}

// chunkWriteTimings accumulates per-phase durations while writing a
// file's chunks out of the blob cache. Debug instrumentation only.
type chunkWriteTimings struct {
	readAt  time.Duration
	write   time.Duration
	sweeper time.Duration
}

// restoreRegularFile reconstructs a regular file by reading chunks
// directly out of cached blobs via ReadAt. The expectation when this
// method runs is that every blob this file needs is already in the
// disk cache — the planner guarantees that by only marking files
// "ready" once their full blob set is on disk.
func (s *restoreSession) restoreRegularFile(
	file *database.File, targetPath string,
) error {
	fileStart := time.Now()

	t0 := time.Now()
	fileChunks, err := s.repos.FileChunks.GetByFileID(s.ctx, file.ID)
	fileChunksQueryDur := time.Since(t0)

	if err != nil {
		return fmt.Errorf("getting file chunks: %w", err)
	}

	t0 = time.Now()

	// Remove any existing entry, then create the file with a restrictive
	// mode via O_EXCL. The stored mode is applied only after the content
	// is written and the file closed, so a file whose stored mode is
	// restrictive is never briefly readable by other local users while
	// its content is written. Removing first (rather than failing on a
	// leftover file) matches the documented behaviour that re-running
	// restore overwrites partial output.
	_ = s.v.Fs.Remove(targetPath)

	outFile, err := s.v.Fs.OpenFile(
		targetPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, restoreFileMode)
	createDur := time.Since(t0)

	if err != nil {
		return fmt.Errorf("creating output file: %w", err)
	}

	bytesWritten, timings, err := s.writeFileChunks(outFile, fileChunks)
	if err != nil {
		// Do not leave a partial file behind.
		_ = outFile.Close()

		s.removePartialRestore(targetPath)

		return err
	}

	log.Debug("Restored regular file (timings)",
		"path", file.Path,
		"chunks", len(fileChunks),
		"bytes_written", bytesWritten,
		"ms_total", time.Since(fileStart).Milliseconds(),
		"ms_file_chunks_query", fileChunksQueryDur.Milliseconds(),
		"ms_create", createDur.Milliseconds(),
		"ms_readat", timings.readAt.Milliseconds(),
		"ms_writes", timings.write.Milliseconds(),
		"ms_sweeper", timings.sweeper.Milliseconds(),
	)

	err = outFile.Close()
	if err != nil {
		s.removePartialRestore(targetPath)

		return fmt.Errorf("closing output file: %w", err)
	}

	s.applyRestoredFileMode(file, targetPath)
	s.applyFileMetadata(file, targetPath)

	s.result.FilesRestored++
	s.result.BytesRestored += bytesWritten

	log.Debug("Restored file", "path", file.Path, "size", ubytes(bytesWritten))

	return nil
}

// applyRestoredFileMode applies the stored permission bits to a
// just-written regular file (created with restoreFileMode). A failure is
// a user-visible warning, not a fatal error: the file's content is
// intact and it remains at the restrictive create-time mode, so the
// restore is not aborted or discarded over it.
func (s *restoreSession) applyRestoredFileMode(
	file *database.File, targetPath string,
) {
	err := s.v.Fs.Chmod(targetPath, os.FileMode(file.Mode))
	if err != nil {
		s.v.UI.Warningf("Failed to set mode %s on %s: %v",
			os.FileMode(file.Mode).Perm(), s.v.UI.Path(targetPath), err)
	}
}

// removePartialRestore deletes a restore output file whose write did not
// complete, so a failed restore never leaves a partial file behind.
func (s *restoreSession) removePartialRestore(targetPath string) {
	err := s.v.Fs.Remove(targetPath)
	if err != nil {
		log.Debug("Failed to remove partial restore file",
			"path", targetPath, "error", err)
	}
}

// writeFileChunks streams each of the file's chunks from the blob disk
// cache into outFile, crediting restored bytes to the sweeper as it
// goes. Returns the bytes written plus per-phase timing accumulators.
func (s *restoreSession) writeFileChunks(
	outFile afero.File, fileChunks []*database.FileChunk,
) (int64, chunkWriteTimings, error) {
	var (
		timings      chunkWriteTimings
		bytesWritten int64
	)

	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]
		if !ok {
			return bytesWritten, timings, fmt.Errorf(
				"%w: %s", errChunkNotInAnyBlob, chunkHashStr[:16])
		}

		blobHash, ok := s.blobIDToHash[blobChunk.BlobID.String()]
		if !ok {
			return bytesWritten, timings, fmt.Errorf(
				"%w: %s", errBlobIDNotInHashIndex, blobChunk.BlobID)
		}

		t0 := time.Now()
		chunkData, err := s.blobCache.ReadAt(
			blobHash, blobChunk.Offset, blobChunk.Length)
		timings.readAt += time.Since(t0)

		if err != nil {
			return bytesWritten, timings, fmt.Errorf(
				"reading chunk %s from cached blob %s: %w",
				fc.ChunkHash[:16], blobHash[:16], err)
		}

		t0 = time.Now()
		n, err := outFile.Write(chunkData)
		timings.write += time.Since(t0)

		if err != nil {
			return bytesWritten, timings, fmt.Errorf("writing chunk: %w", err)
		}

		bytesWritten += int64(n)

		t0 = time.Now()

		s.sweeper.chunkRestored(int64(n))

		timings.sweeper += time.Since(t0)
	}

	return bytesWritten, timings, nil
}

// downloadBlobToCache streams a blob from remote storage straight into
// the disk cache, decrypting and decompressing on the fly. The
// plaintext never lives fully in memory — io.Copy through
// blobDiskCache.PutFromReader uses a 32 KiB buffer regardless of blob
// size, which is what makes multi-GB blobs tractable on machines with
// less RAM than the blob.
func (s *restoreSession) downloadBlobToCache(
	blobHash string, expectedSize int64,
) error {
	start := time.Now()

	t0 := time.Now()
	rc, err := s.v.FetchAndDecryptBlob(s.ctx, blobHash, expectedSize, s.identity)
	fetchSetupDur := time.Since(t0)

	if err != nil {
		return err
	}

	t0 = time.Now()
	written, copyErr := s.blobCache.PutFromReader(blobHash, rc)
	streamDur := time.Since(t0)
	closeErr := rc.Close()

	// closeErr carries the blob's hash-verification result (a mismatch,
	// or the stream not being fully read). On any failure, drop the
	// cache entry so a blob that failed verification is never read back
	// as if it were valid.
	if copyErr != nil {
		s.blobCache.Delete(blobHash)

		return copyErr
	}

	if closeErr != nil {
		s.blobCache.Delete(blobHash)

		return closeErr
	}

	log.Debug("Streamed blob into disk cache",
		"hash", blobHash[:16],
		"compressed_bytes", expectedSize,
		"plaintext_bytes", written,
		"ms_total", time.Since(start).Milliseconds(),
		"ms_fetch_setup", fetchSetupDur.Milliseconds(),
		"ms_stream_decrypt_decompress", streamDur.Milliseconds(),
	)

	return nil
}

// verifyRestoredFiles verifies that all restored files match their
// expected chunk hashes.
func (v *Vaultik) verifyRestoredFiles(
	ctx context.Context,
	repos *database.Repositories,
	files []*database.File,
	targetDir string,
	result *RestoreResult,
) error {
	// Calculate total bytes to verify for progress bar
	var totalBytes int64

	regularFiles := make([]*database.File, 0, len(files))
	for _, file := range files {
		// Skip symlinks and directories - only verify regular files
		if file.IsSymlink() || file.Mode&uint32(os.ModeDir) != 0 {
			continue
		}

		regularFiles = append(regularFiles, file)
		totalBytes += file.Size
	}

	if len(regularFiles) == 0 {
		log.Info("No regular files to verify")

		return nil
	}

	log.Info("Verifying restored files",
		"files", len(regularFiles),
		"bytes", ubytes(totalBytes),
	)
	v.UI.Beginf("Verifying %s files (%s).",
		v.UI.Count(len(regularFiles)),
		v.UI.Size(totalBytes))

	startTime := time.Now()
	lastStatusTime := startTime

	var bytesProcessed int64

	for i, file := range regularFiles {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		targetPath, err := containedRestorePath(v.Fs, targetDir, file.Path.String())
		if err == nil {
			var bytesVerified int64

			bytesVerified, err = v.verifyFile(ctx, repos, file, targetPath)
			if err == nil {
				result.FilesVerified++
				result.BytesVerified += bytesVerified
			}
		}

		if err != nil {
			log.Error("File verification failed", "path", file.Path, "error", err)

			result.FilesFailed++
			result.FailedFiles = append(result.FailedFiles, file.Path.String())
		}

		bytesProcessed += file.Size

		if time.Since(lastStatusTime) >= restoreStatusInterval {
			v.printVerifyProgress(
				i+1, len(regularFiles), bytesProcessed, totalBytes, startTime)

			lastStatusTime = time.Now()
		}
	}

	log.Info("Verification complete",
		"files_verified", result.FilesVerified,
		"bytes_verified", ubytes(result.BytesVerified),
		"files_failed", result.FilesFailed,
	)

	return nil
}

// printVerifyProgress emits a periodic verify-phase status line. Same
// shape as the restore progress line so user-facing pacing is uniform
// across the two phases.
func (v *Vaultik) printVerifyProgress(
	filesDone, totalFiles int, bytesDone, totalBytes int64, startTime time.Time,
) {
	v.printPhaseProgress("Verify", "verify",
		filesDone, totalFiles, bytesDone, totalBytes, startTime)
}

// verifyFile verifies a single restored file by checking its chunk hashes
func (v *Vaultik) verifyFile(
	ctx context.Context,
	repos *database.Repositories,
	file *database.File,
	targetPath string,
) (int64, error) {
	// Get file chunks in order
	fileChunks, err := repos.FileChunks.GetByFileID(ctx, file.ID)
	if err != nil {
		return 0, fmt.Errorf("getting file chunks: %w", err)
	}

	// Open the restored file
	f, err := v.Fs.Open(targetPath)
	if err != nil {
		return 0, fmt.Errorf("opening file: %w", err)
	}

	defer func() { _ = f.Close() }()

	// Verify each chunk
	var bytesVerified int64

	for _, fc := range fileChunks {
		// Get chunk size from database
		chunk, err := repos.Chunks.GetByHash(ctx, fc.ChunkHash.String())
		if err != nil {
			return bytesVerified, fmt.Errorf("getting chunk %s: %w",
				fc.ChunkHash.String()[:16], err)
		}

		// Read chunk data from file
		chunkData := make([]byte, chunk.Size)

		n, err := io.ReadFull(f, chunkData)
		if err != nil {
			return bytesVerified, fmt.Errorf("reading chunk data: %w", err)
		}

		if int64(n) != chunk.Size {
			return bytesVerified, fmt.Errorf("%w: expected %d bytes, got %d",
				errShortChunkRead, chunk.Size, n)
		}

		// Calculate hash and compare
		hash := sha256.Sum256(chunkData)
		actualHash := hex.EncodeToString(hash[:])
		expectedHash := fc.ChunkHash.String()

		if actualHash != expectedHash {
			return bytesVerified, fmt.Errorf("%w: chunk %d: expected %s, got %s",
				errChunkHashMismatch, fc.Idx, expectedHash[:16], actualHash[:16])
		}

		bytesVerified += int64(n)
	}

	// The stored chunks account for the whole file, so the reader must
	// be at EOF now. Trailing bytes past the last chunk are corruption
	// the per-chunk loop cannot see.
	extra := make([]byte, 1)

	n, err := f.Read(extra)
	if n != 0 || !errors.Is(err, io.EOF) {
		return bytesVerified, fmt.Errorf("%w: file longer than its %d chunk(s)",
			errTrailingRestoreData, len(fileChunks))
	}

	log.Debug("File verified",
		"path", file.Path, "bytes", bytesVerified, "chunks", len(fileChunks))

	return bytesVerified, nil
}