at-container-registry/scanner/internal/scan/worker.go

// Package scan implements the vulnerability scanning pipeline:
// extract layers → generate SBOM → scan vulnerabilities → send result.
package scan

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"runtime"
	"sync"
	"time"

	scanner "atcr.io/scanner"
	"atcr.io/scanner/internal/client"
	"atcr.io/scanner/internal/config"
	"atcr.io/scanner/internal/queue"
)

// SkipError is returned by processJob when the scanner intentionally bypasses
// an artifact type it can't analyze (helm charts, in-toto attestations, DSSE).
// The worker dispatches these to hold via SendSkipped so the hold can mark
// the scan record "skipped" instead of "failed". Skipped records are never
// retried by the stale-scan loop; failures are.
type SkipError struct {
	Reason string
}

func (e *SkipError) Error() string { return "skipped: " + e.Reason }

// WorkerPool manages a pool of scan workers
type WorkerPool struct {
	cfg    *config.Config
	queue  *queue.JobQueue
	client *client.HoldClient
	wg     sync.WaitGroup
}

// NewWorkerPool creates a new worker pool
func NewWorkerPool(cfg *config.Config, q *queue.JobQueue, c *client.HoldClient) *WorkerPool {
	return &WorkerPool{
		cfg:    cfg,
		queue:  q,
		client: c,
	}
}

// Start launches worker goroutines
func (wp *WorkerPool) Start(ctx context.Context) {
	// Point TMPDIR at the configured tmp dir so Grype's DB download
	// (go-getter zstd decompression can be 1 GB+) and stereoscope's layer
	// extraction both land on the same partition as the scanner volume —
	// NOT on /tmp, which is typically tmpfs with ~400 MB and would silently
	// fail mid-extract. This must be set before any scanner/grype goroutine
	// starts and must never be restored to a smaller default mid-process.
	if wp.cfg.Vuln.TmpDir != "" {
		if err := os.MkdirAll(wp.cfg.Vuln.TmpDir, 0o755); err != nil {
			slog.Warn("Failed to create scanner tmp dir", "path", wp.cfg.Vuln.TmpDir, "error", err)
		}
		os.Setenv("TMPDIR", wp.cfg.Vuln.TmpDir)
	}

	// Initialize vuln database on startup if enabled
	if wp.cfg.Vuln.Enabled {
		go func() {
			if err := initializeVulnDatabase(wp.cfg.Vuln.DBPath); err != nil {
				slog.Error("Failed to initialize vulnerability database", "error", err)
				slog.Warn("Vulnerability scanning will be disabled until database is available")
			}
		}()
	}

	for i := 0; i < wp.cfg.Scanner.Workers; i++ {
		wp.wg.Add(1)
		go wp.worker(ctx, i)
	}

	slog.Info("Scanner worker pool started", "workers", wp.cfg.Scanner.Workers)
}

// Wait blocks until all workers finish
func (wp *WorkerPool) Wait() {
	wp.wg.Wait()
}

func (wp *WorkerPool) worker(ctx context.Context, id int) {
	defer wp.wg.Done()

	slog.Info("Scanner worker started", "worker_id", id)

	for {
		job := wp.queue.Dequeue()
		if job == nil {
			slog.Info("Scanner worker shutting down", "worker_id", id)
			return
		}

		slog.Info("Processing scan job",
			"worker_id", id,
			"repository", job.Repository,
			"tag", job.Tag,
			"digest", job.ManifestDigest,
			"tier", job.Tier)

		// Tell the hold the scan is actually starting. The ack it already has
		// was sent on receipt, before this job joined the queue, so it cannot
		// tell queueing from scanning without this.
		wp.client.SendStarted(job.Seq)

		result, err := wp.processJob(ctx, job)
		if err != nil {
			var skipErr *SkipError
			if errors.As(err, &skipErr) {
				slog.Info("Scan job skipped",
					"worker_id", id,
					"repository", job.Repository,
					"reason", skipErr.Reason)
				wp.client.SendSkipped(job.Seq, skipErr.Reason)
			} else {
				slog.Error("Scan job failed",
					"worker_id", id,
					"repository", job.Repository,
					"error", err)
				wp.client.SendError(job.Seq, err.Error())
			}
		} else {
			wp.client.SendResult(job.Seq, result)

			// A nil Summary means Grype never ran (vuln.enabled=false), which
			// is not the same as "scanned, found nothing". Log the completion
			// without a count rather than printing a zero the scan never
			// established.
			if result.Summary != nil {
				slog.Info("Scan job completed",
					"worker_id", id,
					"repository", job.Repository,
					"vulnerabilities", result.Summary.Total)
			} else {
				slog.Info("Scan job completed",
					"worker_id", id,
					"repository", job.Repository,
					"vulnerabilities", "not scanned")
			}
		}

		// Free large scan artifacts and trigger GC before the cooldown
		// so memory is reclaimed between jobs. Syft/Grype allocate heavily
		// and Go's GC needs idle time to catch up under sustained load.
		result = nil
		runtime.GC()

		// Cooldown between scans to reduce sustained memory pressure
		select {
		case <-ctx.Done():
			return
		case <-time.After(JobCooldown):
		}
	}
}

// JobCooldown is the pause a worker takes after each job so Go's GC can
// reclaim what Syft and Grype allocated before the next scan starts.
//
// It is a variable, and exported, solely so tests in other packages can
// shorten it: at the production value a scenario that runs a handful of jobs
// through a single worker spends nearly all its runtime asleep. Production
// code must not change it.
var JobCooldown = 10 * time.Second

// unscannable config media types — these are OCI artifacts that aren't
// container images so Syft/Grype can't analyze their layers.
var unscannableConfigTypes = map[string]bool{
	"application/vnd.cncf.helm.config.v1+json": true, // Helm charts
	"application/vnd.in-toto+json":             true, // In-toto attestations
	"application/vnd.dsse.envelope.v1+json":    true, // DSSE envelopes (SLSA)
}

// skipReason reports why a job cannot be scanned, or "" when it can be.
//
// Everything decided here is a permanent property of the manifest record, so
// everything here is a skip rather than an error: the hold records a skip once
// and re-offers a failure on every stale pass, forever.
func skipReason(job *scanner.ScanJob) string {
	if unscannableConfigTypes[job.Config.MediaType] {
		return fmt.Sprintf("unscannable artifact type %s", job.Config.MediaType)
	}

	// A buildx attestation manifest carries an ordinary image config with a
	// single in-toto or DSSE payload as its layer, so the config media type
	// alone does not identify it. buildOCILayout drops every non-tar layer,
	// which would hand Syft an image with nothing in it.
	if len(job.Layers) > 0 && !hasScannableLayer(job) {
		return fmt.Sprintf("no scannable layers (%s)", job.Layers[0].MediaType)
	}

	// Every digest the scan will use has to be a digest. Each one names a blob
	// to ask the hold for and a file to write into the layout, and a string
	// that is neither cannot start being one on a later attempt. Checking here
	// rejects the job before a single request goes out; buildOCILayout parses
	// the same digests again because it is what turns them into paths, and
	// that boundary must hold on its own.
	for _, ref := range referencedBlobs(job) {
		if _, err := scanner.ParseDigest(ref.Descriptor.Digest); err != nil {
			return fmt.Sprintf("%s: %v", ref.what(), err)
		}
	}

	return ""
}

// hasScannableLayer reports whether any layer survives the filter
// buildOCILayout applies. It asks referencedBlobs rather than repeating the
// media-type rule, so this answer and the layout can never disagree.
func hasScannableLayer(job *scanner.ScanJob) bool {
	for _, ref := range referencedBlobs(job) {
		if !ref.isConfig() {
			return true
		}
	}
	return false
}

func (wp *WorkerPool) processJob(ctx context.Context, job *scanner.ScanJob) (*scanner.ScanResult, error) {
	startTime := time.Now()

	// Skip non-container OCI artifacts (Helm charts, in-toto, DSSE, etc.).
	// Returning *SkipError tells the worker dispatch loop to send a "skipped"
	// message rather than an "error" — the hold marks these records as
	// permanently skipped and won't retry them on the rescan interval.
	if reason := skipReason(job); reason != "" {
		return nil, &SkipError{Reason: reason}
	}

	// Ensure tmp dir exists
	if err := ensureDir(wp.cfg.Vuln.TmpDir); err != nil {
		return nil, fmt.Errorf("failed to create tmp dir: %w", err)
	}

	// The cheap guard: refuse an image that admits to being over the ceiling
	// before a byte moves. It is only a pre-check, because the sizes it adds up
	// come from the same user-writable record as the digests; buildOCILayout
	// enforces the same ceiling against the bytes that actually arrive.
	//
	// Either way the verdict is permanent — an image does not shrink — so it is
	// a skip, not a failure the stale loop will offer back forever.
	if wp.cfg.Vuln.MaxImageSize > 0 {
		var totalSize int64
		for _, layer := range job.Layers {
			totalSize += layer.Size
		}
		totalSize += job.Config.Size
		if totalSize > wp.cfg.Vuln.MaxImageSize {
			return nil, &SkipError{Reason: fmt.Sprintf(
				"image too large: %d bytes compressed (limit %d bytes)", totalSize, wp.cfg.Vuln.MaxImageSize)}
		}
	}

	// Step 1: Build OCI image layout from hold via presigned URLs
	slog.Info("Building OCI layout", "repository", job.Repository)
	ociLayoutDir, cleanup, err := buildOCILayout(job, wp.cfg.Vuln.TmpDir, wp.cfg.Hold.Secret, wp.cfg.Vuln.MaxImageSize)
	if err != nil {
		return nil, fmt.Errorf("failed to build OCI layout: %w", err)
	}
	defer cleanup()

	// Step 2: Generate SBOM with Syft
	slog.Info("Generating SBOM", "repository", job.Repository)
	sbomResult, sbomJSON, sbomDigest, err := generateSBOM(ctx, ociLayoutDir)
	if err != nil {
		return nil, fmt.Errorf("failed to generate SBOM: %w", err)
	}

	result := &scanner.ScanResult{
		ManifestDigest: job.ManifestDigest,
		SBOM:           sbomJSON,
		SBOMDigest:     sbomDigest,
	}

	// Step 3: Scan SBOM with Grype (if enabled)
	if wp.cfg.Vuln.Enabled {
		slog.Info("Scanning for vulnerabilities", "repository", job.Repository, "handle", job.UserHandle)
		vulnJSON, vulnDigest, summary, err := scanVulnerabilities(ctx, sbomResult, wp.cfg.Vuln.DBPath)
		if err != nil {
			return nil, fmt.Errorf("failed to scan vulnerabilities: %w", err)
		}
		result.VulnReport = vulnJSON
		result.VulnDigest = vulnDigest
		result.Summary = &summary
	}
	sbomResult = nil // release SBOM catalog for GC

	duration := time.Since(startTime)
	slog.Info("Scan pipeline completed",
		"repository", job.Repository,
		"duration", duration)

	return result, nil
}

func ensureDir(path string) error {
	return os.MkdirAll(path, 0755)
}