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