vaultik/internal/vaultik/restore_interrupt_test.go

package vaultik //nolint:testpackage // sets ctx/cancel and inspects scratch files

import (
	"bytes"
	"context"
	"io"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/spf13/afero"
	"github.com/stretchr/testify/assert"
	"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 (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-verify-*",
	} {
		found, err := filepath.Glob(filepath.Join(dir, pattern))
		require.NoError(t, err)

		matches = append(matches, found...)
	}

	return matches
}

// TestRestoreSkipErrorsCancelDuringBlobDownload cancels a SkipErrors
// restore while a blob download is in progress. The download fails only
// because of the cancel, so Restore must return context.Canceled without
// reporting the file that needs the blob as failed.
func TestRestoreSkipErrorsCancelDuringBlobDownload(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	fs := afero.NewOsFs()
	tempDir := t.TempDir()

	cfg, storer, snapshotID, srcPath := backupOneFile(context.Background(),
		t, fs, tempDir, "a.txt", []byte("hello vaultik"), 0o644)

	gate := newBlockingBlobStorer(storer)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var out bytes.Buffer

	v := newRestoreVaultik(ctx, cfg, gate, fs)
	v.UI = ui.NewWithColor(&out, false)

	restoreErr := make(chan error, 1)

	go func() {
		restoreErr <- v.Restore(&RestoreOptions{
			SnapshotID: snapshotID,
			TargetDir:  filepath.Join(tempDir, "restored"),
			SkipErrors: true,
		})
	}()

	select {
	case <-gate.entered:
	case <-time.After(30 * time.Second):
		t.Fatal("restore never reached the blob-download phase")
	}

	cancel()

	require.ErrorIs(t, <-restoreErr, context.Canceled)
	assert.NotContains(t, out.String(), srcPath,
		"the cancel was reported as a failed file")
}

// 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 newStallingStorer(
	inner storage.Storer, stalls func(key string) bool,
) *stallingStorer {
	return &stallingStorer{
		Storer:  inner,
		stalls:  stalls,
		stalled: make(chan struct{}),
	}
}

func (s *stallingStorer) Get(
	ctx context.Context, key string,
) (io.ReadCloser, error) {
	rc, err := s.Storer.Get(ctx, key)
	if err != nil || !s.stalls(key) {
		return rc, err
	}

	object, 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(object[:len(object)/2])

		s.once.Do(func() { close(s.stalled) })
	}()

	return pr, nil
}

// 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 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{
				SnapshotID: snapshotID,
				TargetDir:  targetDir,
			})
		},
		"deep verify": func(v *Vaultik, snapshotID, _ string) error {
			return v.RunDeepVerify(snapshotID, &VerifyOptions{Deep: true})
		},
	}

	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()

				// 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 := newStallingStorer(storer, stalls)

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