at-container-registry/scanner/internal/mockhold/blobsource.go

// Package mockhold provides an in-process stand-in for the hold service so
// scanner tests can drive the real scanner binary through failure scenarios
// that are impractical to provoke against a live hold: mid-scan disconnects,
// truncated blobs, stalled downloads, and artifact shapes the real hold's
// dispatch guards currently refuse to send.
//
// The scanner touches exactly three hold endpoints, and this package serves
// all three:
//
//	GET /xrpc/io.atcr.hold.subscribeScanJobs   WebSocket, jobs out / acks in
//	GET /xrpc/com.atproto.sync.getBlob         returns {"url": "..."}
//	GET /blobs/{hex}                           the bytes that URL points at
//
// Blob bytes come from a BlobSource, so the same scenario can run against a
// real image pulled with skopeo or against synthetic bytes built in the test.
package mockhold

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// ErrBlobNotFound is returned by a BlobSource that has no bytes for a digest.
// The mock hold turns this into a 404 from getBlob, which is what a real hold
// returns for a blob that was garbage collected out from under a queued job.
var ErrBlobNotFound = errors.New("mockhold: blob not found")

// BlobSource supplies blob bytes by digest. Implementations must be safe for
// concurrent use: the scanner downloads a config and every layer in sequence,
// but several workers may be scanning at once.
//
// Size is returned separately from the reader so the mock can set
// Content-Length. A source that does not know the size ahead of time (a
// generated stream, say) may return -1, and the mock will use chunked
// encoding.
type BlobSource interface {
	Open(digest string) (io.ReadCloser, int64, error)
}

// DigestHex extracts the hex portion of a digest string, mirroring the
// scanner's own digestHex so a source keyed either way resolves the same.
func DigestHex(digest string) string {
	if _, hex, ok := strings.Cut(digest, ":"); ok {
		return hex
	}
	return digest
}

// OCILayout serves blobs out of an OCI image layout directory, which is what
// `skopeo copy docker://<ref> oci:<dir>:<tag>` writes. The layout stores blobs
// at blobs/sha256/<hex>, keyed by exactly the digest the scan job references,
// so a layout pulled from a real registry can back a scan with no rewriting.
type OCILayout struct {
	Dir string
}

// NewOCILayout returns a source reading from an OCI layout directory.
func NewOCILayout(dir string) *OCILayout { return &OCILayout{Dir: dir} }

// Open implements BlobSource.
func (o *OCILayout) Open(digest string) (io.ReadCloser, int64, error) {
	path := filepath.Join(o.Dir, "blobs", "sha256", DigestHex(digest))

	// Refuse a digest that escapes the layout. Digests are attacker-controlled
	// in the sense that a test may deliberately feed a malformed one, and a
	// traversal would read outside the fixture rather than failing the way a
	// real hold would.
	if !strings.HasPrefix(filepath.Clean(path), filepath.Clean(o.Dir)+string(os.PathSeparator)) {
		return nil, 0, fmt.Errorf("%w: %q escapes the layout", ErrBlobNotFound, digest)
	}

	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, 0, fmt.Errorf("%w: %s", ErrBlobNotFound, digest)
		}
		return nil, 0, err
	}

	fi, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, 0, err
	}

	return f, fi.Size(), nil
}

// Memory serves blobs from an in-memory map keyed by digest. Use it for
// synthetic content: a hand-built tar layer, a gzip bomb, bytes that do not
// hash to the digest claiming them.
type Memory struct {
	blobs map[string][]byte
}

// NewMemory returns an empty in-memory source.
func NewMemory() *Memory { return &Memory{blobs: make(map[string][]byte)} }

// Add stores bytes under a digest. The digest is stored by hex, so callers may
// pass either "sha256:abc" or "abc".
func (m *Memory) Add(digest string, data []byte) *Memory {
	m.blobs[DigestHex(digest)] = data
	return m
}

// Open implements BlobSource.
func (m *Memory) Open(digest string) (io.ReadCloser, int64, error) {
	data, ok := m.blobs[DigestHex(digest)]
	if !ok {
		return nil, 0, fmt.Errorf("%w: %s", ErrBlobNotFound, digest)
	}
	return io.NopCloser(bytes.NewReader(data)), int64(len(data)), nil
}

// Multi tries each source in order and returns the first hit. It lets one
// scenario span several skopeo-pulled layouts, or overlay synthetic bytes on
// top of a real image to corrupt exactly one layer.
type Multi struct {
	sources []BlobSource
}

// NewMulti returns a source that consults each of sources in order.
func NewMulti(sources ...BlobSource) *Multi { return &Multi{sources: sources} }

// Open implements BlobSource.
func (m *Multi) Open(digest string) (io.ReadCloser, int64, error) {
	for _, s := range m.sources {
		rc, n, err := s.Open(digest)
		if err == nil {
			return rc, n, nil
		}
		// A source that failed for a reason other than "I don't have it" is a
		// real fault and must not be masked by a later source happening to.
		if !errors.Is(err, ErrBlobNotFound) {
			return nil, 0, err
		}
	}
	return nil, 0, fmt.Errorf("%w: %s", ErrBlobNotFound, digest)
}