vaultik/internal/vaultik/blob_fetch.go

package vaultik

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"time"

	"filippo.io/age"
	"sneak.berlin/go/vaultik/internal/blobgen"
	"sneak.berlin/go/vaultik/internal/log"
)

// errBlobHashMismatch is returned when a fetched blob's content hash does
// not match the expected double-SHA-256 hash.
var errBlobHashMismatch = errors.New("blob hash mismatch")

// hashVerifyReader wraps a blobgen.Reader and verifies the double-SHA-256 hash
// of decrypted plaintext when Close is called. It reuses the hash that
// blobgen.Reader already computes internally via its TeeReader, avoiding
// redundant SHA-256 computation.
type hashVerifyReader struct {
	reader   *blobgen.Reader // underlying decrypted blob reader (has internal hasher)
	fetcher  io.ReadCloser   // raw fetched stream (closed on Close)
	blobHash string          // expected double-SHA-256 hex
	done     bool            // EOF reached
}

func (h *hashVerifyReader) Read(p []byte) (int, error) {
	n, err := h.reader.Read(p)
	if errors.Is(err, io.EOF) {
		h.done = true
	}

	return n, err
}

// Close verifies the hash (if the stream was fully read) and closes underlying readers.
func (h *hashVerifyReader) Close() error {
	readerErr := h.reader.Close()
	fetcherErr := h.fetcher.Close()

	if h.done {
		firstHash := h.reader.Sum256()
		secondHasher := sha256.New()
		secondHasher.Write(firstHash)

		actualHashHex := hex.EncodeToString(secondHasher.Sum(nil))
		if actualHashHex != h.blobHash {
			return fmt.Errorf("%w: expected %s, got %s",
				errBlobHashMismatch, h.blobHash[:16], actualHashHex[:16])
		}
	}

	if readerErr != nil {
		return readerErr
	}

	return fetcherErr
}

// FetchAndDecryptBlob downloads a blob, decrypts and decompresses it, and
// returns a streaming reader that computes the double-SHA-256 hash on the fly.
// The hash is verified when the returned reader is closed (after fully reading).
// This avoids buffering the entire blob in memory.
func (v *Vaultik) FetchAndDecryptBlob(
	ctx context.Context, blobHash string, expectedSize int64, identity age.Identity,
) (io.ReadCloser, error) {
	rc, _, err := v.FetchBlob(ctx, blobHash, expectedSize)
	if err != nil {
		return nil, err
	}

	reader, err := blobgen.NewReader(rc, identity)
	if err != nil {
		_ = rc.Close()

		return nil, fmt.Errorf("creating blob reader: %w", err)
	}

	return &hashVerifyReader{
		reader:   reader,
		fetcher:  rc,
		blobHash: blobHash,
	}, nil
}

// FetchBlob downloads a blob and returns a reader for the encrypted data.
// Times the Storage.Get and Storage.Stat round-trips separately at
// debug level so we can see whether the size-only Stat (which is an
// extra request on every fetch) is hurting throughput.
func (v *Vaultik) FetchBlob(
	ctx context.Context, blobHash string, expectedSize int64,
) (io.ReadCloser, int64, error) {
	blobPath := fmt.Sprintf("blobs/%s/%s/%s", blobHash[:2], blobHash[2:4], blobHash)

	t0 := time.Now()
	rc, err := v.Storage.Get(ctx, blobPath)
	getDur := time.Since(t0)

	if err != nil {
		return nil, 0, fmt.Errorf("downloading blob %s: %w", blobHash[:16], err)
	}

	t0 = time.Now()
	info, err := v.Storage.Stat(ctx, blobPath)
	statDur := time.Since(t0)

	if err != nil {
		_ = rc.Close()

		return nil, 0, fmt.Errorf("stat blob %s: %w", blobHash[:16], err)
	}

	log.Debug("FetchBlob round-trips",
		"hash", blobHash[:16],
		"ms_storage_get", getDur.Milliseconds(),
		"ms_storage_stat", statDur.Milliseconds(),
		"expected_size", expectedSize,
		"stat_size", info.Size,
	)

	return rc, info.Size, nil
}