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)
	}

	// Reclaim what earlier processes left behind before any worker starts
	// filling the directory again. A scan that is interrupted by a restart
	// never runs its own cleanup, so without this every restart leaks the
	// in-flight layout and extraction permanently.
	sweepTmpDir(wp.cfg.Vuln.TmpDir, wp.cfg.Vuln.SweepMaxAge)

	// 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")
			}
		}()
	}

	// Say so at boot rather than leaving it to be inferred from duplicate
	// scans months later. The scanner cannot read the hold's configuration, so
	// this compares against the value the hold shipped with; a hold running a
	// different budget makes this advisory, which is why it is a warning and
	// not a refusal to start.
	switch d := wp.jobTimeout(); {
	case d <= 0:
		slog.Warn("No per-job scan deadline configured; one wedged job will "+
			"hold its worker until the process restarts",
			"config", "scanner.job_timeout")
	case d >= holdScanningTimeout:
		slog.Warn("scanner.job_timeout is at or above the hold's scanning "+
			"timeout, so the hold will give up on a job before this scanner "+
			"does and re-dispatch work that is still running",
			"job_timeout", d, "hold_scanning_timeout", holdScanningTimeout)
	}

	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,
		"job_timeout", wp.jobTimeout())
}

// 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)

		// The job's own clock starts here, at the same instant the hold's
		// does. See holdScanningTimeout for why the two must not be equal.
		jobCtx, cancelJob := wp.jobContext(ctx)
		result, err := wp.processJob(jobCtx, job)

		// Ask the context, not the error. Only the context knows whether the
		// budget ran out; several stages wrap or replace the cause on the way
		// back, and one of them (stereoscope's extraction) does not take a
		// context at all, so the error it returns after an overrun says
		// nothing about the deadline.
		//
		// Only when the job actually failed. A scan that finished inside an
		// uninterruptible stage after the deadline passed still holds a real
		// verdict, and the hold has not reclaimed the row: reporting a timeout
		// instead would throw away work that is already done.
		timedOut := err != nil && errors.Is(jobCtx.Err(), context.DeadlineExceeded)
		cancelJob()

		if err != nil {
			var skipErr *SkipError
			switch {
			case timedOut:
				// "error" and not "skipped", deliberately. A timeout is a
				// statement about this host at this moment — a contended CPU,
				// a slow S3, a co-tenant hold running garbage collection — not
				// about the image, and the same bytes may well scan on the
				// next pass. The hold re-offers errors and never re-offers
				// skips, so a skip here would retire an image permanently on
				// one bad afternoon.
				//
				// The cost of being wrong in this direction is a rescan at
				// full price on the stale-scan schedule, and vuln.max_image_size
				// already refuses the pathological images before a byte moves.
				// The cost of being wrong in the other direction is an image
				// that is never scanned again and no way to tell from the
				// record why.
				slog.Error("Scan job timed out",
					"worker_id", id,
					"repository", job.Repository,
					"timeout", wp.jobTimeout(),
					"error", err)
				wp.client.SendError(job.Seq, fmt.Sprintf(
					"scan timed out after %s: %v", wp.jobTimeout(), err))
			case 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)
			default:
				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):
		}
	}
}

// holdScanningTimeout mirrors scanningTimeout in
// pkg/hold/pds/scan_broadcaster.go. It is duplicated rather than imported
// because the scanner is a separate module and a hold may be running a
// different version than this scanner; it is a reference point for a startup
// warning, never a bound this process enforces.
//
// The ordering it exists to protect: the hold stamps started_at when it
// receives the "started" message a worker sends on dequeue, and fails the row
// this long after. The scanner starts its own clock at the same instant, so
// the two budgets measure the same interval and whichever is shorter decides
// what happens. Shorter here means the scanner stops, sends a terminal message,
// and the hold records the real reason. Shorter there means the hold marks the
// row failed and hands it to the next scanner with a worker free, while this
// one is still scanning: the same image scanned twice, and the verdict that
// arrives late lands on a row somebody else now owns.
const holdScanningTimeout = 10 * time.Minute

// jobTimeout is the configured per-job budget, zero or less meaning none.
func (wp *WorkerPool) jobTimeout() time.Duration {
	return wp.cfg.Scanner.JobTimeout
}

// jobContext derives the context one job runs under.
//
// It always derives from the pool context, deadline or not, so shutdown
// interrupts a job whether or not an operator has kept the deadline.
func (wp *WorkerPool) jobContext(ctx context.Context) (context.Context, context.CancelFunc) {
	if d := wp.jobTimeout(); d > 0 {
		return context.WithTimeout(ctx, d)
	}
	return context.WithCancel(ctx)
}

// 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(ctx, 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, job.ManifestDigest)
	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)
}