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"
	"fmt"
	"log/slog"
	"os"
	"sync"
	"time"

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

// 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) {
	// Initialize vuln database on startup if enabled
	if wp.cfg.VulnEnabled {
		go func() {
			if err := initializeVulnDatabase(wp.cfg.VulnDBPath, wp.cfg.TmpDir); 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.Workers; i++ {
		wp.wg.Add(1)
		go wp.worker(ctx, i)
	}

	slog.Info("Scanner worker pool started", "workers", wp.cfg.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)

		result, err := wp.processJob(ctx, job)
		if err != nil {
			slog.Error("Scan job failed",
				"worker_id", id,
				"repository", job.Repository,
				"error", err)
			wp.client.SendError(job.Seq, err.Error())
			continue
		}

		wp.client.SendResult(job.Seq, result)

		slog.Info("Scan job completed",
			"worker_id", id,
			"repository", job.Repository,
			"vulnerabilities", result.Summary.Total)
	}
}

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

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

	// Step 1: Extract image layers from hold via presigned URLs
	slog.Info("Extracting image layers", "repository", job.Repository)
	imageDir, cleanup, err := extractLayers(job, wp.cfg.TmpDir)
	if err != nil {
		return nil, fmt.Errorf("failed to extract layers: %w", err)
	}
	defer cleanup()

	// Step 2: Generate SBOM with Syft
	slog.Info("Generating SBOM", "repository", job.Repository)
	sbomResult, sbomJSON, sbomDigest, err := generateSBOM(ctx, imageDir)
	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.VulnEnabled {
		slog.Info("Scanning for vulnerabilities", "repository", job.Repository)
		vulnJSON, vulnDigest, summary, err := scanVulnerabilities(ctx, sbomResult, wp.cfg.VulnDBPath)
		if err != nil {
			return nil, fmt.Errorf("failed to scan vulnerabilities: %w", err)
		}
		result.VulnReport = vulnJSON
		result.VulnDigest = vulnDigest
		result.Summary = &summary
	}

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