vaultik/internal/vaultik/restore_interrupt_test.go

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

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

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

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

	return matches
}