vaultik/internal/snapshot/backup_test.go

package snapshot_test

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"testing"
	"testing/fstest"
	"time"

	"sneak.berlin/go/vaultik/internal/database"
	"sneak.berlin/go/vaultik/internal/types"
)

// errBlobNotFound is returned by the mock S3 client for unknown hashes.
var errBlobNotFound = errors.New("blob not found")

// testFile1Name is the shared fixture filename used across backup tests.
const testFile1Name = "file1.txt"

// MockS3Client is a mock implementation of S3 operations for testing
type MockS3Client struct {
	storage map[string][]byte
}

func NewMockS3Client() *MockS3Client {
	return &MockS3Client{
		storage: make(map[string][]byte),
	}
}

func (m *MockS3Client) PutBlob(_ context.Context, hash string, data []byte) error {
	m.storage[hash] = data

	return nil
}

func (m *MockS3Client) GetBlob(_ context.Context, hash string) ([]byte, error) {
	data, ok := m.storage[hash]
	if !ok {
		return nil, fmt.Errorf("%w: %s", errBlobNotFound, hash)
	}

	return data, nil
}

func (m *MockS3Client) BlobExists(_ context.Context, hash string) (bool, error) {
	_, ok := m.storage[hash]

	return ok, nil
}

func (m *MockS3Client) CreateBucket(_ context.Context, _ string) error {
	return nil
}

// verifyBackupFiles checks the file records created by a backup against
// the fixture filesystem.
func verifyBackupFiles(
	ctx context.Context,
	t *testing.T,
	repos *database.Repositories,
	testFS fstest.MapFS,
) {
	t.Helper()

	files, err := repos.Files.ListByPrefix(ctx, "")
	if err != nil {
		t.Fatalf("Failed to list files: %v", err)
	}

	expectedFiles := map[string]bool{
		testFile1Name:           true,
		"dir1/file2.txt":        true,
		"dir1/subdir/file3.txt": true,
		"largefile.bin":         true,
	}

	if len(files) != len(expectedFiles) {
		t.Errorf("Expected %d files, got %d", len(expectedFiles), len(files))
	}

	for _, file := range files {
		if !expectedFiles[file.Path.String()] {
			t.Errorf("Unexpected file in database: %s", file.Path)
		}

		delete(expectedFiles, file.Path.String())

		// Verify file metadata
		fsFile := testFS[file.Path.String()]
		if fsFile == nil {
			t.Errorf("File %s not found in test filesystem", file.Path)

			continue
		}

		if file.Size != int64(len(fsFile.Data)) {
			t.Errorf("File %s: expected size %d, got %d",
				file.Path, len(fsFile.Data), file.Size)
		}

		if file.Mode != uint32(fsFile.Mode) {
			t.Errorf("File %s: expected mode %o, got %o",
				file.Path, fsFile.Mode, file.Mode)
		}
	}

	if len(expectedFiles) > 0 {
		t.Errorf("Files not found in database: %v", expectedFiles)
	}
}

// verifyBackupChunksAndBlobs checks that chunking produced the expected
// records and every referenced blob exists in the mock S3 store.
func verifyBackupChunksAndBlobs(
	ctx context.Context,
	t *testing.T,
	repos *database.Repositories,
	s3Client *MockS3Client,
	snapshotID string,
) {
	t.Helper()

	chunks, err := repos.Chunks.List(ctx)
	if err != nil {
		t.Fatalf("Failed to list chunks: %v", err)
	}

	if len(chunks) == 0 {
		t.Error("No chunks found in database")
	}

	// The large file should create 10 chunks (10MB / 1MB chunk size)
	// Plus the small files
	minExpectedChunks := 10 + 3
	if len(chunks) < minExpectedChunks {
		t.Errorf("Expected at least %d chunks, got %d", minExpectedChunks, len(chunks))
	}

	// Verify at least one blob was created and uploaded
	// We can't list blobs directly, but we can check via snapshot blobs
	blobHashes, err := repos.Snapshots.GetBlobHashes(ctx, snapshotID)
	if err != nil {
		t.Fatalf("Failed to get blob hashes: %v", err)
	}

	if len(blobHashes) == 0 {
		t.Error("Expected at least one blob to be created")
	}

	for _, blobHash := range blobHashes {
		// Check blob exists in mock S3
		exists, err := s3Client.BlobExists(ctx, blobHash)
		if err != nil {
			t.Errorf("Failed to check blob %s: %v", blobHash, err)
		}

		if !exists {
			t.Errorf("Blob %s not found in S3", blobHash)
		}
	}
}

func TestBackupWithInMemoryFS(t *testing.T) {
	t.Parallel()

	// Create a temporary directory for the database
	tempDir := t.TempDir()
	dbPath := filepath.Join(tempDir, "test.db")

	// Create test filesystem
	testFS := fstest.MapFS{
		testFile1Name: &fstest.MapFile{
			Data:    []byte("Hello, World!"),
			Mode:    0644,
			ModTime: time.Now(),
		},
		"dir1/file2.txt": &fstest.MapFile{
			Data:    []byte("This is a test file with some content."),
			Mode:    0755,
			ModTime: time.Now(),
		},
		"dir1/subdir/file3.txt": &fstest.MapFile{
			Data:    []byte("Another file in a subdirectory."),
			Mode:    0600,
			ModTime: time.Now(),
		},
		"largefile.bin": &fstest.MapFile{
			Data:    generateLargeFileContent(10 * 1024 * 1024), // 10MB file with varied content
			Mode:    0644,
			ModTime: time.Now(),
		},
	}

	// Initialize the database
	ctx := context.Background()

	db, err := database.New(ctx, dbPath)
	if err != nil {
		t.Fatalf("Failed to create database: %v", err)
	}

	defer func() {
		err := db.Close()
		if err != nil {
			t.Logf("Failed to close database: %v", err)
		}
	}()

	repos := database.NewRepositories(db)

	// Create mock S3 client
	s3Client := NewMockS3Client()

	// Run backup
	backupEngine := &BackupEngine{
		repos:    repos,
		s3Client: s3Client,
	}

	snapshotID, err := backupEngine.Backup(ctx, testFS, ".")
	if err != nil {
		t.Fatalf("Backup failed: %v", err)
	}

	// Verify snapshot was created
	snapshot, err := repos.Snapshots.GetByID(ctx, snapshotID)
	if err != nil {
		t.Fatalf("Failed to get snapshot: %v", err)
	}

	if snapshot == nil {
		t.Fatal("Snapshot not found")
	}

	if snapshot.FileCount == 0 {
		t.Error("Expected snapshot to have files")
	}

	// Verify files, chunks, and blob records
	verifyBackupFiles(ctx, t, repos, testFS)
	verifyBackupChunksAndBlobs(ctx, t, repos, s3Client, snapshotID)
}

func TestBackupDeduplication(t *testing.T) {
	t.Parallel()

	// Create a temporary directory for the database
	tempDir := t.TempDir()
	dbPath := filepath.Join(tempDir, "test.db")

	// Create test filesystem with duplicate content
	testFS := fstest.MapFS{
		testFile1Name: &fstest.MapFile{
			Data:    []byte("Duplicate content"),
			Mode:    0644,
			ModTime: time.Now(),
		},
		"file2.txt": &fstest.MapFile{
			Data:    []byte("Duplicate content"),
			Mode:    0644,
			ModTime: time.Now(),
		},
		"file3.txt": &fstest.MapFile{
			Data:    []byte("Unique content"),
			Mode:    0644,
			ModTime: time.Now(),
		},
	}

	// Initialize the database
	ctx := context.Background()

	db, err := database.New(ctx, dbPath)
	if err != nil {
		t.Fatalf("Failed to create database: %v", err)
	}

	defer func() {
		err := db.Close()
		if err != nil {
			t.Logf("Failed to close database: %v", err)
		}
	}()

	repos := database.NewRepositories(db)

	// Create mock S3 client
	s3Client := NewMockS3Client()

	// Run backup
	backupEngine := &BackupEngine{
		repos:    repos,
		s3Client: s3Client,
	}

	_, err = backupEngine.Backup(ctx, testFS, ".")
	if err != nil {
		t.Fatalf("Backup failed: %v", err)
	}

	// Verify deduplication
	chunks, err := repos.Chunks.List(ctx)
	if err != nil {
		t.Fatalf("Failed to list chunks: %v", err)
	}

	// Should have only 2 unique chunks (duplicate content + unique content)
	if len(chunks) != 2 {
		t.Errorf("Expected 2 unique chunks, got %d", len(chunks))
	}

	// Verify chunk references
	for _, chunk := range chunks {
		files, err := repos.ChunkFiles.GetByChunkHash(ctx, chunk.ChunkHash)
		if err != nil {
			t.Errorf("Failed to get files for chunk %s: %v", chunk.ChunkHash, err)
		}

		// The duplicate content chunk should be referenced by 2 files
		if chunk.Size == int64(len("Duplicate content")) && len(files) != 2 {
			t.Errorf("Expected duplicate chunk to be referenced by 2 files, got %d",
				len(files))
		}
	}
}

// BackupEngine performs backup operations
type BackupEngine struct {
	repos    *database.Repositories
	s3Client interface {
		PutBlob(ctx context.Context, hash string, data []byte) error
		BlobExists(ctx context.Context, hash string) (bool, error)
	}
}

// backupCounters accumulates statistics across a test backup run.
type backupCounters struct {
	fileCount  int64
	chunkCount int64
	blobCount  int64
	totalSize  int64
	blobSize   int64
}

// Backup performs a backup of the given filesystem
func (b *BackupEngine) Backup(
	ctx context.Context, fsys fs.FS, root string,
) (string, error) {
	// Create a new snapshot
	hostname, _ := os.Hostname()
	snapshotID := time.Now().Format(time.RFC3339)
	snapshot := &database.Snapshot{
		ID:             types.SnapshotID(snapshotID),
		Hostname:       types.Hostname(hostname),
		VaultikVersion: "test",
		StartedAt:      time.Now(),
		CompletedAt:    nil,
	}

	// Create initial snapshot record
	err := b.repos.WithTx(ctx, func(ctx context.Context, tx *sql.Tx) error {
		return b.repos.Snapshots.Create(ctx, tx, snapshot)
	})
	if err != nil {
		return "", err
	}

	counters := &backupCounters{}

	// Track which chunks we've seen to handle deduplication
	processedChunks := make(map[string]bool)

	// Scan the filesystem and process files
	err = fs.WalkDir(fsys, root, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}

		// Skip directories
		if d.IsDir() {
			return nil
		}

		// Get file info
		info, err := d.Info()
		if err != nil {
			return err
		}

		// Handle symlinks
		if info.Mode()&fs.ModeSymlink != 0 {
			// For testing, we'll skip symlinks since fstest doesn't support them well
			return nil
		}

		return b.backupOneFile(ctx, fsys, path, info, processedChunks, counters)
	})
	if err != nil {
		return "", err
	}

	// After all files are processed, create blobs for new chunks
	err = b.createBlobsForChunks(ctx, snapshotID, processedChunks, counters)
	if err != nil {
		return "", err
	}

	// Update snapshot with final counts
	err = b.repos.WithTx(ctx, func(ctx context.Context, tx *sql.Tx) error {
		return b.repos.Snapshots.UpdateCounts(ctx, tx, snapshotID,
			counters.fileCount, counters.chunkCount, counters.blobCount,
			counters.totalSize, counters.blobSize)
	})
	if err != nil {
		return "", err
	}

	return snapshotID, nil
}

// backupOneFile records a single regular file and its chunks.
func (b *BackupEngine) backupOneFile(
	ctx context.Context,
	fsys fs.FS,
	path string,
	info fs.FileInfo,
	processedChunks map[string]bool,
	counters *backupCounters,
) error {
	// Create file record in a short transaction
	file := &database.File{
		Path:  types.FilePath(path),
		Size:  info.Size(),
		Mode:  uint32(info.Mode()),
		MTime: info.ModTime(),
		UID:   1000, // Default UID for test
		GID:   1000, // Default GID for test
	}

	err := b.repos.WithTx(ctx, func(ctx context.Context, tx *sql.Tx) error {
		return b.repos.Files.Create(ctx, tx, file)
	})
	if err != nil {
		return err
	}

	counters.fileCount++
	counters.totalSize += info.Size()

	// Read and process file in chunks
	f, err := fsys.Open(path)
	if err != nil {
		return err
	}
	defer func() {
		err := f.Close()
		if err != nil {
			// Log but don't fail since we're already in an error path potentially
			fmt.Fprintf(os.Stderr, "Failed to close file: %v\n", err)
		}
	}()

	// Process file in chunks
	chunkIndex := 0
	buffer := make([]byte, defaultChunkSize)

	for {
		n, err := f.Read(buffer)
		if err != nil && !errors.Is(err, io.EOF) {
			return err
		}

		if n == 0 {
			break
		}

		err = b.recordChunk(ctx, file, chunkIndex, buffer[:n], processedChunks)
		if err != nil {
			return err
		}

		chunkIndex++
	}

	return nil
}

// recordChunk creates the chunk record (if new) and its file associations.
func (b *BackupEngine) recordChunk(
	ctx context.Context,
	file *database.File,
	chunkIndex int,
	chunkData []byte,
	processedChunks map[string]bool,
) error {
	chunkHash := calculateHash(chunkData)

	// Check if chunk already exists (outside of transaction)
	existingChunk, _ := b.repos.Chunks.GetByHash(ctx, chunkHash)
	if existingChunk == nil {
		// Create new chunk in a short transaction
		err := b.repos.WithTx(ctx, func(ctx context.Context, tx *sql.Tx) error {
			chunk := &database.Chunk{
				ChunkHash: types.ChunkHash(chunkHash),
				Size:      int64(len(chunkData)),
			}

			return b.repos.Chunks.Create(ctx, tx, chunk)
		})
		if err != nil {
			return err
		}

		processedChunks[chunkHash] = true
	}

	// Create file-chunk mapping in a short transaction
	err := b.repos.WithTx(ctx, func(ctx context.Context, tx *sql.Tx) error {
		fileChunk := &database.FileChunk{
			FileID:    file.ID,
			Idx:       chunkIndex,
			ChunkHash: types.ChunkHash(chunkHash),
		}

		return b.repos.FileChunks.Create(ctx, tx, fileChunk)
	})
	if err != nil {
		return err
	}

	// Create chunk-file mapping in a short transaction
	return b.repos.WithTx(ctx, func(ctx context.Context, tx *sql.Tx) error {
		chunkFile := &database.ChunkFile{
			ChunkHash:  types.ChunkHash(chunkHash),
			FileID:     file.ID,
			FileOffset: int64(chunkIndex * defaultChunkSize),
			Length:     int64(len(chunkData)),
		}

		return b.repos.ChunkFiles.Create(ctx, tx, chunkFile)
	})
}

// createBlobsForChunks uploads one blob per new chunk and records the blob
// metadata and snapshot association.
func (b *BackupEngine) createBlobsForChunks(
	ctx context.Context,
	snapshotID string,
	processedChunks map[string]bool,
	counters *backupCounters,
) error {
	for chunkHash := range processedChunks {
		// Get chunk data (outside of transaction)
		chunk, err := b.repos.Chunks.GetByHash(ctx, chunkHash)
		if err != nil {
			return err
		}

		counters.chunkCount++

		// In a real system, blobs would contain multiple chunks and be encrypted
		// For testing, we'll create a blob with a "blob-" prefix to differentiate
		blobHash := "blob-" + chunkHash

		// For the test, we'll create dummy data since we don't have the original
		dummyData := []byte(chunkHash)

		// Upload to S3 as a blob
		err = b.s3Client.PutBlob(ctx, blobHash, dummyData)
		if err != nil {
			return err
		}

		// Create blob entry in a short transaction
		blobID := types.NewBlobID()

		err = b.repos.WithTx(ctx, func(ctx context.Context, tx *sql.Tx) error {
			blob := &database.Blob{
				ID:        blobID,
				Hash:      types.BlobHash(blobHash),
				CreatedTS: time.Now(),
			}

			return b.repos.Blobs.Create(ctx, tx, blob)
		})
		if err != nil {
			return err
		}

		counters.blobCount++
		counters.blobSize += chunk.Size

		// Create blob-chunk mapping in a short transaction
		err = b.repos.WithTx(ctx, func(ctx context.Context, tx *sql.Tx) error {
			blobChunk := &database.BlobChunk{
				BlobID:    blobID,
				ChunkHash: types.ChunkHash(chunkHash),
				Offset:    0,
				Length:    chunk.Size,
			}

			return b.repos.BlobChunks.Create(ctx, tx, blobChunk)
		})
		if err != nil {
			return err
		}

		// Add blob to snapshot in a short transaction
		err = b.repos.WithTx(ctx, func(ctx context.Context, tx *sql.Tx) error {
			return b.repos.Snapshots.AddBlob(ctx, tx, snapshotID, blobID,
				types.BlobHash(blobHash))
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func calculateHash(data []byte) string {
	h := sha256.New()
	h.Write(data)

	return hex.EncodeToString(h.Sum(nil))
}

func generateLargeFileContent(size int) []byte {
	data := make([]byte, size)
	// Fill with pattern that changes every chunk to avoid deduplication
	for i := range size {
		chunkNum := i / defaultChunkSize
		data[i] = byte((i + chunkNum) % 256)
	}

	return data
}

const defaultChunkSize = 1024 * 1024 // 1MB chunks