at-container-registry/pkg/hold/scanner/queue.go

package scanner

import (
	"context"
	"fmt"
	"log/slog"
	"sync"

	"atcr.io/pkg/atproto"
)

// Queue manages a pool of workers for scanning container images
type Queue struct {
	jobs    chan *ScanJob
	results chan *ScanResult
	workers int
	wg      sync.WaitGroup
	ctx     context.Context
	cancel  context.CancelFunc
}

// NewQueue creates a new scanner queue with the specified number of workers
func NewQueue(workers int, bufferSize int) *Queue {
	ctx, cancel := context.WithCancel(context.Background())

	return &Queue{
		jobs:    make(chan *ScanJob, bufferSize),
		results: make(chan *ScanResult, bufferSize),
		workers: workers,
		ctx:     ctx,
		cancel:  cancel,
	}
}

// Start starts the worker pool
// The workerFunc is called for each job to perform the actual scanning
func (q *Queue) Start(workerFunc func(context.Context, *ScanJob) (*ScanResult, error)) {
	slog.Info("Starting scanner worker pool", "workers", q.workers)

	for i := 0; i < q.workers; i++ {
		q.wg.Add(1)
		go q.worker(i, workerFunc)
	}

	// Start result handler goroutine
	q.wg.Add(1)
	go q.resultHandler()
}

// worker processes jobs from the queue
func (q *Queue) worker(id int, workerFunc func(context.Context, *ScanJob) (*ScanResult, error)) {
	defer q.wg.Done()

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

	for {
		select {
		case <-q.ctx.Done():
			slog.Info("Scanner worker shutting down", "worker_id", id)
			return

		case job, ok := <-q.jobs:
			if !ok {
				slog.Info("Scanner worker: jobs channel closed", "worker_id", id)
				return
			}

			slog.Info("Scanner worker processing job",
				"worker_id", id,
				"repository", job.Repository,
				"tag", job.Tag,
				"digest", job.ManifestDigest)

			result, err := workerFunc(q.ctx, job)
			if err != nil {
				slog.Error("Scanner worker failed to process job",
					"worker_id", id,
					"repository", job.Repository,
					"tag", job.Tag,
					"error", err)
				continue
			}

			// Send result to results channel
			select {
			case q.results <- result:
				slog.Info("Scanner worker completed job",
					"worker_id", id,
					"repository", job.Repository,
					"tag", job.Tag,
					"vulnerabilities", result.Summary.Total)
			case <-q.ctx.Done():
				return
			}
		}
	}
}

// resultHandler processes scan results (for logging and metrics)
func (q *Queue) resultHandler() {
	defer q.wg.Done()

	for {
		select {
		case <-q.ctx.Done():
			return

		case result, ok := <-q.results:
			if !ok {
				return
			}

			// Log the result
			slog.Info("Scan completed",
				"repository", result.Job.Repository,
				"tag", result.Job.Tag,
				"digest", result.Job.ManifestDigest,
				"critical", result.Summary.Critical,
				"high", result.Summary.High,
				"medium", result.Summary.Medium,
				"low", result.Summary.Low,
				"total", result.Summary.Total,
				"scanner", result.ScannerVersion)
		}
	}
}

// Enqueue adds a job to the queue
func (q *Queue) Enqueue(jobAny any) error {
	// Type assert to ScanJob (can be map or struct from HandleNotifyManifest)
	var job *ScanJob

	switch v := jobAny.(type) {
	case *ScanJob:
		job = v
	case map[string]interface{}:
		// Convert map to ScanJob (from HandleNotifyManifest)
		job = &ScanJob{
			ManifestDigest: v["manifestDigest"].(string),
			Repository:     v["repository"].(string),
			Tag:            v["tag"].(string),
			UserDID:        v["userDID"].(string),
			UserHandle:     v["userHandle"].(string),
		}

		// Parse config blob reference
		if configMap, ok := v["config"].(map[string]interface{}); ok {
			job.Config = atproto.BlobReference{
				Digest:    configMap["digest"].(string),
				Size:      convertToInt64(configMap["size"]),
				MediaType: configMap["mediaType"].(string),
			}
		}

		// Parse layers
		if layersSlice, ok := v["layers"].([]interface{}); ok {
			slog.Info("Parsing layers from scan job",
				"layersFound", len(layersSlice))
			job.Layers = make([]atproto.BlobReference, len(layersSlice))
			for i, layerAny := range layersSlice {
				if layerMap, ok := layerAny.(map[string]interface{}); ok {
					job.Layers[i] = atproto.BlobReference{
						Digest:    layerMap["digest"].(string),
						Size:      convertToInt64(layerMap["size"]),
						MediaType: layerMap["mediaType"].(string),
					}
				}
			}
		} else {
			slog.Warn("No layers found in scan job map",
				"layersType", fmt.Sprintf("%T", v["layers"]),
				"layersValue", v["layers"])
		}
	default:
		return fmt.Errorf("invalid job type: %T", jobAny)
	}

	select {
	case q.jobs <- job:
		slog.Info("Enqueued scan job",
			"repository", job.Repository,
			"tag", job.Tag,
			"digest", job.ManifestDigest)
		return nil
	case <-q.ctx.Done():
		return q.ctx.Err()
	}
}

// Shutdown gracefully shuts down the queue, waiting for all workers to finish
func (q *Queue) Shutdown() {
	slog.Info("Shutting down scanner queue")

	// Close the jobs channel to signal no more jobs
	close(q.jobs)

	// Wait for all workers to finish
	q.wg.Wait()

	// Close results channel
	close(q.results)

	// Cancel context
	q.cancel()

	slog.Info("Scanner queue shut down complete")
}

// Len returns the number of jobs currently in the queue
func (q *Queue) Len() int {
	return len(q.jobs)
}

// convertToInt64 converts an interface{} number to int64, handling both float64 and int64
func convertToInt64(v interface{}) int64 {
	switch n := v.(type) {
	case float64:
		return int64(n)
	case int64:
		return n
	case int:
		return int64(n)
	default:
		return 0
	}
}