vaultik/internal/blob/packer_test.go

package blob_test

import (
	"bytes"
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"io"
	"testing"

	"filippo.io/age"
	"github.com/klauspost/compress/zstd"
	"github.com/spf13/afero"
	"sneak.berlin/go/vaultik/internal/blob"
	"sneak.berlin/go/vaultik/internal/database"
	"sneak.berlin/go/vaultik/internal/log"
	"sneak.berlin/go/vaultik/internal/types"
)

const (
	// Test key from test/insecure-integration-test.key
	testPrivateKey = "AGE-SECRET-KEY-19CR5YSFW59HM4TLD6GXVEDMZFTVVF7PPHKUT68TXSFPK7A" +
		"PHXA2QS2NJA5"
	testPublicKey = "age1ezrjmfpwsc95svdg0y54mums3zevgzu0x0ecq2f7tp8a05gl0sjq9q9wjg"

	defaultMaxBlobSize = 10 * 1024 * 1024 // 10MB
	testChunkSize      = 1000
	testChunkCount     = 10
)

// parseTestIdentity parses the fixed test age identity.
func parseTestIdentity(t *testing.T) *age.X25519Identity {
	t.Helper()

	identity, err := age.ParseX25519Identity(testPrivateKey)
	if err != nil {
		t.Fatalf("failed to parse test identity: %v", err)
	}

	return identity
}

// newTestPacker creates a test database and a Packer backed by it.
func newTestPacker(
	t *testing.T, maxBlobSize int64,
) (*database.Repositories, *blob.Packer) {
	t.Helper()

	db, err := database.NewTestDB()
	if err != nil {
		t.Fatalf("failed to create test db: %v", err)
	}

	t.Cleanup(func() { _ = db.Close() })

	repos := database.NewRepositories(db)

	packer, err := blob.NewPacker(blob.PackerConfig{
		MaxBlobSize:      maxBlobSize,
		CompressionLevel: 3,
		Recipients:       []string{testPublicKey},
		Repositories:     repos,
		Fs:               afero.NewMemMapFs(),
	})
	if err != nil {
		t.Fatalf("failed to create packer: %v", err)
	}

	return repos, packer
}

// makeChunk creates a ChunkRef for data and registers the chunk in the
// database.
func makeChunk(
	t *testing.T, repos *database.Repositories, data []byte,
) *blob.ChunkRef {
	t.Helper()

	hash := sha256.Sum256(data)
	hashStr := hex.EncodeToString(hash[:])

	dbChunk := &database.Chunk{
		ChunkHash: types.ChunkHash(hashStr),
		Size:      int64(len(data)),
	}

	err := repos.WithTx(
		context.Background(),
		func(ctx context.Context, tx *sql.Tx) error {
			return repos.Chunks.Create(ctx, tx, dbChunk)
		})
	if err != nil {
		t.Fatalf("failed to create chunk in db: %v", err)
	}

	return &blob.ChunkRef{
		Hash: hashStr,
		Data: data,
	}
}

// decryptAndDecompress reverses the blob pipeline: age decrypt, then zstd
// decompress.
func decryptAndDecompress(
	t *testing.T, blobData []byte, identity *age.X25519Identity,
) []byte {
	t.Helper()

	decrypted, err := age.Decrypt(bytes.NewReader(blobData), identity)
	if err != nil {
		t.Fatalf("failed to decrypt blob: %v", err)
	}

	reader, err := zstd.NewReader(decrypted)
	if err != nil {
		t.Fatalf("failed to create decompressor: %v", err)
	}
	defer reader.Close()

	var decompressed bytes.Buffer

	_, err = io.Copy(&decompressed, reader)
	if err != nil {
		t.Fatalf("failed to decompress: %v", err)
	}

	return decompressed.Bytes()
}

func TestPackerSingleChunk(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	identity := parseTestIdentity(t)
	repos, packer := newTestPacker(t, defaultMaxBlobSize)
	ctx := context.Background()

	data := []byte("Hello, World!")
	chunk := makeChunk(t, repos, data)

	err := packer.AddChunk(ctx, chunk)
	if err != nil {
		t.Fatalf("failed to add chunk: %v", err)
	}

	err = packer.Flush(ctx)
	if err != nil {
		t.Fatalf("failed to flush: %v", err)
	}

	blobs := packer.GetFinishedBlobs()
	if len(blobs) != 1 {
		t.Fatalf("expected 1 blob, got %d", len(blobs))
	}

	finished := blobs[0]
	if len(finished.Chunks) != 1 {
		t.Errorf("expected 1 chunk in blob, got %d", len(finished.Chunks))
	}

	// Note: Very small data may not compress well
	t.Logf("Compression: %d -> %d bytes",
		finished.Uncompressed, finished.Compressed)

	decompressed := decryptAndDecompress(t, finished.Data, identity)
	if !bytes.Equal(decompressed, data) {
		t.Error("decompressed data doesn't match original")
	}
}

func TestPackerMultipleChunks(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	repos, packer := newTestPacker(t, defaultMaxBlobSize)
	ctx := context.Background()

	chunks := make([]*blob.ChunkRef, testChunkCount)
	for i := range testChunkCount {
		data := bytes.Repeat([]byte{byte(i)}, testChunkSize)
		chunks[i] = makeChunk(t, repos, data)
	}

	for _, chunk := range chunks {
		err := packer.AddChunk(ctx, chunk)
		if err != nil {
			t.Fatalf("failed to add chunk: %v", err)
		}
	}

	err := packer.Flush(ctx)
	if err != nil {
		t.Fatalf("failed to flush: %v", err)
	}

	blobs := packer.GetFinishedBlobs()
	if len(blobs) != 1 {
		t.Fatalf("expected 1 blob, got %d", len(blobs))
	}

	if len(blobs[0].Chunks) != testChunkCount {
		t.Errorf("expected %d chunks in blob, got %d",
			testChunkCount, len(blobs[0].Chunks))
	}

	// Verify offsets are correct
	expectedOffset := int64(0)

	for i, chunkRef := range blobs[0].Chunks {
		if chunkRef.Offset != expectedOffset {
			t.Errorf("chunk %d: expected offset %d, got %d",
				i, expectedOffset, chunkRef.Offset)
		}

		if chunkRef.Length != testChunkSize {
			t.Errorf("chunk %d: expected length %d, got %d",
				i, testChunkSize, chunkRef.Length)
		}

		expectedOffset += chunkRef.Length
	}
}

func TestPackerSizeLimit(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	const (
		maxBlobSize        = 5000 // 5KB max, forces multiple blobs
		maxBlobawoOverhead = 6000 // allow some overhead over the limit
	)

	repos, packer := newTestPacker(t, maxBlobSize)
	ctx := context.Background()

	chunks := make([]*blob.ChunkRef, testChunkCount)
	for i := range testChunkCount {
		data := bytes.Repeat([]byte{byte(i)}, testChunkSize) // 1KB each
		chunks[i] = makeChunk(t, repos, data)
	}

	blobCount := 0

	// Add chunks and handle size limit errors
	for _, chunk := range chunks {
		err := packer.AddChunk(ctx, chunk)
		if errors.Is(err, blob.ErrBlobSizeLimitExceeded) {
			// Finalize current blob
			err = packer.FinalizeBlob(ctx)
			if err != nil {
				t.Fatalf("failed to finalize blob: %v", err)
			}

			blobCount++

			// Retry adding the chunk
			err = packer.AddChunk(ctx, chunk)
			if err != nil {
				t.Fatalf("failed to add chunk after finalize: %v", err)
			}
		} else if err != nil {
			t.Fatalf("failed to add chunk: %v", err)
		}
	}

	err := packer.Flush(ctx)
	if err != nil {
		t.Fatalf("failed to flush: %v", err)
	}

	blobs := packer.GetFinishedBlobs()

	totalBlobs := blobCount + len(blobs)
	if totalBlobs < 2 {
		t.Errorf("expected multiple blobs due to size limit, got %d", totalBlobs)
	}

	// Verify each blob respects size limit (approximately)
	for _, finished := range blobs {
		if finished.Compressed > maxBlobawoOverhead {
			t.Errorf("blob size %d exceeds limit", finished.Compressed)
		}
	}
}

func TestPackerEncryption(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	identity := parseTestIdentity(t)
	repos, packer := newTestPacker(t, defaultMaxBlobSize)
	ctx := context.Background()

	data := bytes.Repeat([]byte("Test data for encryption!"), 100)
	chunk := makeChunk(t, repos, data)

	err := packer.AddChunk(ctx, chunk)
	if err != nil {
		t.Fatalf("failed to add chunk: %v", err)
	}

	err = packer.Flush(ctx)
	if err != nil {
		t.Fatalf("failed to flush: %v", err)
	}

	blobs := packer.GetFinishedBlobs()
	if len(blobs) != 1 {
		t.Fatalf("expected 1 blob, got %d", len(blobs))
	}

	decompressed := decryptAndDecompress(t, blobs[0].Data, identity)
	if !bytes.Equal(decompressed, data) {
		t.Error("decrypted and decompressed data doesn't match original")
	}
}