at-container-registry/scanner/internal/e2e/deadline_test.go

package e2e

import (
	"context"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync/atomic"
	"testing"
	"time"

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

// This file covers the per-job deadline: the bound on how long one scan may
// occupy a worker, and what the scanner says when it gives up.
//
// The ordering that makes any of it worth having is with the hold. The hold
// stamps started_at when a worker sends 'started' and fails the row
// scanningTimeout after that (10 minutes, pkg/hold/pds/scan_broadcaster.go).
// If the scanner's own deadline fires first it sends a terminal message and
// the hold records the real reason; if the hold's fires first it reclaims a
// row the scanner is still working on and the work is done twice. Every
// default here exists to keep the scanner strictly first.

// --- a stall that can be observed being aborted ----------------------------

// abortableGate is newGate with one difference that is the whole point: its
// handler watches the request context as well as the release channel, so a
// test can tell an aborted request from one that merely finished.
//
// The plain gate cannot: it blocks on <-release unconditionally, so a
// cancelled request looks identical to a served one from the server side.
type abortableGate struct {
	srv      *httptest.Server
	entered  chan struct{}
	release  chan struct{}
	aborted  atomic.Int64
	released atomic.Int64
}

func newAbortableGate(t *testing.T) *abortableGate {
	t.Helper()
	g := &abortableGate{entered: make(chan struct{}, 8), release: make(chan struct{})}
	g.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case g.entered <- struct{}{}:
		default:
		}
		select {
		case <-r.Context().Done():
			g.aborted.Add(1)
			return
		case <-g.release:
			g.released.Add(1)
			http.Error(w, "gate released", http.StatusNotFound)
		}
	}))
	t.Cleanup(func() {
		g.openOnce()
		g.srv.Close()
	})
	return g
}

func (g *abortableGate) waitEntered(t *testing.T, timeout time.Duration) {
	t.Helper()
	select {
	case <-g.entered:
	case <-time.After(timeout):
		t.Fatal("worker never reached the stalled blob fetch")
	}
}

func (g *abortableGate) openOnce() {
	defer func() { _ = recover() }()
	close(g.release)
}

// waitAborted waits for the server to observe the client hanging up.
func (g *abortableGate) waitAborted(t *testing.T, timeout time.Duration) bool {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if g.aborted.Load() > 0 {
			return true
		}
		time.Sleep(10 * time.Millisecond)
	}
	return false
}

// --- shutdown ---------------------------------------------------------------

// TestShutdownAbortsAnInFlightDownload is the cancellability precondition for
// everything else in this file. A per-job deadline is a context deadline, and
// a context is worth nothing on a code path that never looks at one.
//
// The blob fetches go through client.GetBlobPresignedURL and
// client.DownloadBlob. Both used a package-level http.Client and built their
// requests with http.NewRequest, so no context reached the transport at all
// and the only bound was the client's own five-minute per-request Timeout.
// Cancelling the pool context (what cmd/scanner does on SIGTERM, before
// WorkerPool.Wait) therefore did nothing: the worker stayed inside the fetch,
// Wait blocked for as long as the fetch did, and a 30 second termination grace
// period turned into a SIGKILL with the scan directory left behind.
func TestShutdownAbortsAnInFlightDownload(t *testing.T) {
	hold := mockhold.New(mockhold.NewMemory(), mockhold.WithSecret(testSecret))
	t.Cleanup(hold.Close)

	cfg := config.DefaultConfig()
	cfg.Hold.URL = hold.URL()
	cfg.Hold.Secret = testSecret
	cfg.Scanner.Workers = 1
	cfg.Vuln.Enabled = false
	cfg.Vuln.TmpDir = t.TempDir()

	origTmpDir, hadTmpDir := os.LookupEnv("TMPDIR")
	t.Cleanup(func() {
		if hadTmpDir {
			os.Setenv("TMPDIR", origTmpDir)
			return
		}
		os.Unsetenv("TMPDIR")
	})

	restoreCooldown := scan.JobCooldown
	scan.JobCooldown = testJobCooldown

	q := queue.NewJobQueue(cfg.Scanner.QueueSize)
	c := client.NewHoldClient(cfg.Hold.URL, cfg.Hold.Secret, q)
	pool := scan.NewWorkerPool(cfg, q, c)

	ctx, cancel := context.WithCancel(context.Background())
	pool.Start(ctx)
	go c.Connect()

	var closed bool
	t.Cleanup(func() {
		cancel()
		if !closed {
			c.Close()
		}
		q.Close()
		pool.Wait()
		scan.JobCooldown = restoreCooldown
	})

	if err := hold.WaitForScanner(10 * time.Second); err != nil {
		t.Fatalf("scanner never connected: %v", err)
	}

	g := newAbortableGate(t)
	job := synthJob("probe/shutdown")
	job.HoldEndpoint = g.srv.URL
	if _, err := hold.SendJob(job); err != nil {
		t.Fatalf("send job: %v", err)
	}
	g.waitEntered(t, 15*time.Second)

	// The shutdown sequence cmd/scanner runs on SIGTERM.
	cancel()
	c.Close()
	closed = true
	q.Close()

	waited := make(chan struct{})
	go func() { pool.Wait(); close(waited) }()

	select {
	case <-waited:
	case <-time.After(10 * time.Second):
		t.Fatal("worker did not exit on context cancellation: the download " +
			"path still ignores ctx, so shutdown waits out the HTTP client's " +
			"own timeout instead")
	}

	if !g.waitAborted(t, 5*time.Second) {
		t.Error("the server never saw the request hang up; the fetch was not " +
			"actually cancelled, the worker just stopped waiting for it")
	}

	entries, err := os.ReadDir(cfg.Vuln.TmpDir)
	if err != nil {
		t.Fatalf("read tmp dir: %v", err)
	}
	for _, e := range entries {
		t.Errorf("scan directory left behind after a cancelled download: %s", e.Name())
	}
}

// --- the per-job deadline ---------------------------------------------------

// withJobTimeout shortens the per-job deadline. Production defaults to
// something that must accommodate a genuinely large image; these scenarios
// need it to fire in milliseconds.
func withJobTimeout(d time.Duration) Option {
	return func(c *config.Config) { c.Scanner.JobTimeout = d }
}

// TestJobDeadlineFailsAStalledPresignRequest is the headline behaviour: a job
// that wedges in the download stage is given up on, on its own, without the
// hold having to notice.
//
// The stall is the presigned-URL request, which is the first network call a
// job makes and the one buildOCILayout used to make with no context at all.
// The assertion that matters is the timing: the terminal message has to arrive
// while the server is still holding the request open, or the scanner did not
// give up, it just outlasted the fault.
func TestJobDeadlineFailsAStalledPresignRequest(t *testing.T) {
	const jobTimeout = 400 * time.Millisecond

	hold := newHold(mockhold.NewMemory())
	h := startScanner(t, hold, withJobTimeout(jobTimeout))

	g := newAbortableGate(t)
	job := synthJob("probe/deadline")
	job.HoldEndpoint = g.srv.URL

	started := time.Now()
	seq, err := h.Hold.SendJob(job)
	if err != nil {
		t.Fatalf("send job: %v", err)
	}
	g.waitEntered(t, 15*time.Second)

	msg := h.AwaitTerminal(t, seq, 10*time.Second)
	elapsed := time.Since(started)

	// A timeout is not a property of the image: the same bytes may well scan
	// on a faster host, or on this one when it is not contending with a hold
	// doing garbage collection. 'skipped' is terminal at the hold and would
	// retire the image permanently on the evidence of one slow afternoon.
	if msg.Type != "error" {
		t.Fatalf("timeout reported as %q (%s%s), want \"error\" so the hold retries it",
			msg.Type, msg.Error, msg.Reason)
	}
	if !strings.Contains(msg.Error, "timed out") {
		t.Errorf("error does not say the job timed out, so the hold records a "+
			"transport failure instead: %q", msg.Error)
	}
	if elapsed > 5*time.Second {
		t.Errorf("gave up after %s, far past the %s deadline: the deadline is "+
			"not what ended the job", elapsed, jobTimeout)
	}
	if g.released.Load() > 0 {
		t.Error("the stalled request was served rather than abandoned; the " +
			"scanner waited the fault out")
	}
	if !g.waitAborted(t, 5*time.Second) {
		t.Error("the server never saw the request hang up: the deadline ended " +
			"the job but left the fetch running")
	}

	// The scan directory is created before the first download and removed by
	// buildOCILayout's own cleanup on failure. A deadline that fires mid-
	// download must not be the one path that skips it.
	assertNoLeakedScanDirs(t, h)
}

// TestJobDeadlineFailsAStalledBlobBody is the same deadline one call later.
//
// GetBlobPresignedURL and DownloadBlob are separate requests through the same
// client, and only the second one streams bytes. A hold that answers the
// presign promptly and then dribbles the body is the more realistic stall, and
// it exercises the io.Copy rather than the round trip.
func TestJobDeadlineFailsAStalledBlobBody(t *testing.T) {
	const jobTimeout = 400 * time.Millisecond

	stalled := make(chan struct{})
	var aborted atomic.Int64
	hold := newHold(mockhold.NewMemory(), mockhold.WithBlobResponseHook(
		func(w http.ResponseWriter, r *http.Request, digest string) bool {
			w.Header().Set("Content-Length", "4096")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("partial"))
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			select {
			case <-r.Context().Done():
				aborted.Add(1)
			case <-stalled:
			}
			return true
		}))
	t.Cleanup(func() { close(stalled) })

	h := startScanner(t, hold, withJobTimeout(jobTimeout))

	cfgBytes := unparseableConfig()
	started := time.Now()
	seq, err := h.Hold.SendJob(jobFor(desc(digestOf(cfgBytes), int64(len(cfgBytes)), configType)))
	if err != nil {
		t.Fatalf("send job: %v", err)
	}

	msg := h.AwaitTerminal(t, seq, 10*time.Second)
	elapsed := time.Since(started)

	if msg.Type != "error" {
		t.Fatalf("want error, got %s: %s%s", msg.Type, msg.Error, msg.Reason)
	}
	if !strings.Contains(msg.Error, "timed out") {
		t.Errorf("error does not say the job timed out: %q", msg.Error)
	}
	if elapsed > 5*time.Second {
		t.Errorf("gave up after %s against a %s deadline", elapsed, jobTimeout)
	}
	deadline := time.Now().Add(5 * time.Second)
	for aborted.Load() == 0 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if aborted.Load() == 0 {
		t.Error("the body stream was never cancelled; the download outlived the job")
	}
	assertNoLeakedScanDirs(t, h)
}

// TestJobDeadlineFreesTheWorkerForTheNextJob is the reason the deadline exists.
//
// At the shipped default of one worker, a job with no bound stops that scanner
// entirely: everything behind it sits acked in the priority queue while the
// hold's dispatch budget stays spent. The deadline is only worth having if the
// worker comes back.
func TestJobDeadlineFreesTheWorkerForTheNextJob(t *testing.T) {
	hold := newHold(mockhold.NewMemory())
	h := startScanner(t, hold, withJobTimeout(400*time.Millisecond))

	g := newAbortableGate(t)
	wedged := synthJob("probe/wedged")
	wedged.HoldEndpoint = g.srv.URL
	wedgedSeq, err := h.Hold.SendJob(wedged)
	if err != nil {
		t.Fatalf("send wedged job: %v", err)
	}
	g.waitEntered(t, 15*time.Second)

	// A helm chart is refused on its config media type alone, so it reaches a
	// terminal message without touching the network: whether it is answered
	// depends only on whether a worker ever dequeues it.
	next := helmJob(t)
	nextSeq, err := h.Hold.SendJob(next)
	if err != nil {
		t.Fatalf("send follow-up job: %v", err)
	}

	if msg := h.AwaitTerminal(t, wedgedSeq, 10*time.Second); msg.Type != "error" {
		t.Fatalf("wedged job ended as %s, want error", msg.Type)
	}
	if msg := h.AwaitTerminal(t, nextSeq, 10*time.Second); msg.Type != "skipped" {
		t.Fatalf("follow-up job ended as %s (%s%s), want skipped: the worker "+
			"never got back to it", msg.Type, msg.Error, msg.Reason)
	}
}

// TestJobDeadlineDefaultLeavesTheHoldRoomToHearAboutIt pins the ordering the
// whole design rests on.
//
// The hold fails a row scanningTimeout after the worker's 'started' message
// (10 minutes, pkg/hold/pds/scan_broadcaster.go) and hands it back to the
// pending pool. The scanner sends 'started' at the same instant it starts its
// own clock, so the two budgets measure the same interval and the scanner's
// has to be the shorter one — with enough margin left for the terminal message
// to cross the WebSocket and be recorded.
//
// Misconfigured the other way (job_timeout at or above the hold's budget) the
// hold wins the race: it marks the row failed and re-dispatches while this
// scanner is still working, which is the duplicate-scan behaviour the deadline
// was added to prevent. Nothing here can enforce that across two processes;
// this only pins the shipped default.
func TestJobDeadlineDefaultLeavesTheHoldRoomToHearAboutIt(t *testing.T) {
	// pkg/hold/pds/scan_broadcaster.go: scanningTimeout.
	const holdScanningTimeout = 10 * time.Minute

	got := config.DefaultConfig().Scanner.JobTimeout
	if got <= 0 {
		t.Fatalf("scanner.job_timeout defaults to %s: shipped with no deadline at all", got)
	}
	if got >= holdScanningTimeout {
		t.Fatalf("scanner.job_timeout defaults to %s, at or past the hold's %s "+
			"scanning budget: the hold reclaims the row first and the scan is "+
			"run twice", got, holdScanningTimeout)
	}
	if margin := holdScanningTimeout - got; margin < time.Minute {
		t.Errorf("only %s between the scanner's deadline and the hold's; too "+
			"little room for the terminal message to land", margin)
	}
	// A node:22 scan measured 16s end to end on a fast workstation, 81% of it
	// stereoscope extraction, on hardware far better than production. The
	// default has to leave room for that on a much smaller host.
	if got < 2*time.Minute {
		t.Errorf("scanner.job_timeout defaults to %s, too short for a large "+
			"image on a small host", got)
	}
}