vaultik/internal/blobgen/writer.go

package blobgen

import (
	"crypto/sha256"
	"errors"
	"fmt"
	"hash"
	"io"
	"runtime"

	"filippo.io/age"
	"github.com/klauspost/compress/zstd"
)

// Zstd compression level bounds accepted by NewWriter.
const (
	minCompressionLevel = 1
	maxCompressionLevel = 19
)

// reservedCompressionCPUs is how many CPUs are left free of zstd
// compression work for I/O and hashing.
const reservedCompressionCPUs = 2

// ErrInvalidCompressionLevel is returned when the zstd compression level
// is outside the accepted 1-19 range.
var ErrInvalidCompressionLevel = errors.New(
	"invalid compression level: must be between 1 and 19")

// Writer wraps compression and encryption with SHA256 hashing.
// Data flows: input -> tee(hasher, compressor -> encryptor -> destination)
// The hash is computed on the uncompressed input for deterministic content-addressing.
type Writer struct {
	teeWriter        io.Writer      // Tee to hasher and compressor
	compressor       *zstd.Encoder  // Compression layer
	encryptor        io.WriteCloser // Encryption layer
	hasher           hash.Hash      // SHA256 hasher (on uncompressed input)
	compressionLevel int
	bytesWritten     int64
}

// NewWriter creates a new Writer that compresses, encrypts, and hashes
// data. The hash is computed on the uncompressed input for deterministic
// content-addressing.
func NewWriter(
	w io.Writer, compressionLevel int, recipients []string,
) (*Writer, error) {
	// Validate compression level
	err := validateCompressionLevel(compressionLevel)
	if err != nil {
		return nil, err
	}

	// Create SHA256 hasher for the uncompressed input
	hasher := sha256.New()

	// Parse recipients
	var ageRecipients []age.Recipient

	for _, recipient := range recipients {
		r, err := age.ParseX25519Recipient(recipient)
		if err != nil {
			return nil, fmt.Errorf("parsing recipient %s: %w", recipient, err)
		}

		ageRecipients = append(ageRecipients, r)
	}

	// Create encryption writer that outputs to destination
	encWriter, err := age.Encrypt(w, ageRecipients...)
	if err != nil {
		return nil, fmt.Errorf("creating encryption writer: %w", err)
	}

	// Calculate compression concurrency: CPUs - 2, minimum 1
	concurrency := max(runtime.NumCPU()-reservedCompressionCPUs, 1)

	// Create compression writer with encryption as destination
	compressor, err := zstd.NewWriter(encWriter,
		zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(compressionLevel)),
		zstd.WithEncoderConcurrency(concurrency),
	)
	if err != nil {
		_ = encWriter.Close()

		return nil, fmt.Errorf("creating compression writer: %w", err)
	}

	// Create tee writer: input goes to both hasher and compressor
	teeWriter := io.MultiWriter(hasher, compressor)

	return &Writer{
		teeWriter:        teeWriter,
		compressor:       compressor,
		encryptor:        encWriter,
		hasher:           hasher,
		compressionLevel: compressionLevel,
	}, nil
}

// Write implements io.Writer
func (w *Writer) Write(p []byte) (int, error) {
	n, err := w.teeWriter.Write(p)
	w.bytesWritten += int64(n)

	return n, err
}

// Close closes all layers and returns any errors
func (w *Writer) Close() error {
	// Close compressor first
	err := w.compressor.Close()
	if err != nil {
		return fmt.Errorf("closing compressor: %w", err)
	}

	// Then close encryptor
	err = w.encryptor.Close()
	if err != nil {
		return fmt.Errorf("closing encryptor: %w", err)
	}

	return nil
}

// Sum256 returns the double SHA256 hash of the uncompressed input data.
// Double hashing (SHA256(SHA256(data))) prevents information leakage about
// the plaintext - an attacker cannot confirm existence of known content
// by computing its hash and checking for a matching blob filename.
func (w *Writer) Sum256() []byte {
	// First hash: SHA256(plaintext)
	firstHash := w.hasher.Sum(nil)
	// Second hash: SHA256(firstHash) - this is the blob ID
	secondHash := sha256.Sum256(firstHash)

	return secondHash[:]
}

// BytesWritten returns the number of uncompressed bytes written
func (w *Writer) BytesWritten() int64 {
	return w.bytesWritten
}

func validateCompressionLevel(level int) error {
	// Zstd compression levels: 1-19 (default is 3)
	// SpeedFastest = 1, SpeedDefault = 3, SpeedBetterCompression = 7,
	// SpeedBestCompression = 11
	if level < minCompressionLevel || level > maxCompressionLevel {
		return fmt.Errorf("%w: got %d", ErrInvalidCompressionLevel, level)
	}

	return nil
}