at-container-registry/pkg/hold/pds/scan_broadcaster_test.go

package pds

import (
	"context"
	"database/sql"
	"os"
	"path/filepath"
	"testing"
	"time"

	holddb "atcr.io/pkg/hold/db"
	"atcr.io/pkg/s3"
)

// newTestScanBroadcaster builds a ScanBroadcaster with just a database, no
// background goroutines. Subscriber lifecycle is all that is under test here,
// so the constructor's discovery/dispatch/stale loops (and their s3 and PDS
// dependencies) are deliberately skipped.
//
// The database is opened the way NewScanBroadcaster opens it, through
// holddb.OpenLocalDB, so every pooled connection carries busy_timeout and the
// file is in WAL mode. A bare sql.Open gives neither: the tests that poll a
// job's status while the storage goroutine writes it then race the writer's
// commit on a second connection, and on a slow CI disk that read fails with
// "database is locked" instead of waiting.
func newTestScanBroadcaster(t *testing.T) *ScanBroadcaster {
	t.Helper()

	db, err := holddb.OpenLocalDB(filepath.Join(t.TempDir(), "scan.db"))
	if err != nil {
		t.Fatalf("open scan db: %v", err)
	}
	t.Cleanup(func() { _ = db.Close() })

	sb := &ScanBroadcaster{
		db:           db,
		holdDID:      "did:web:hold.example.com",
		holdEndpoint: "https://hold.example.com",
	}
	if err := sb.initSchema(); err != nil {
		t.Fatalf("initSchema: %v", err)
	}
	return sb
}

// newRecordingScanBroadcaster is newTestScanBroadcaster plus the two
// dependencies handleResult needs to finish its work: an embedded PDS to write
// the scan record into, and an S3 stand-in to take the SBOM blob. The bare
// helper leaves both nil, which is fine for row bookkeeping and fatal for
// anything that asserts on what was stored.
func newRecordingScanBroadcaster(t *testing.T) *ScanBroadcaster {
	t.Helper()

	sb := newTestScanBroadcaster(t)
	sb.inflight = make(map[string]struct{})
	sb.ackTimeout = 5 * time.Minute

	pds := setupScanTestPDS(t)
	sb.pds = pds
	sb.holdDID = pds.did
	sb.s3 = &s3.S3Service{Client: s3.NewMockS3Client(""), Bucket: "test-bucket"}

	return sb
}

// setupScanTestPDS is setupTestPDS on a file-backed database rather than
// ":memory:".
//
// go-libsql's in-memory database is per *connection*, not per database, so the
// moment two goroutines write through one *sql.DB the second one gets a
// connection with no tables in it — "no such table: blocks" out of the
// carstore. That did not matter while every scan record was written on the
// caller's goroutine. It does now: terminal messages are stored on the
// subscriber's storage goroutine, so the record write and the test's readback
// are on different goroutines by construction.
func setupScanTestPDS(t *testing.T) *HoldPDS {
	t.Helper()

	ctx := context.Background()
	tmpDir := t.TempDir()

	keyPath := filepath.Join(tmpDir, "signing-key")
	if err := os.WriteFile(keyPath, sharedTestKey, 0600); err != nil {
		t.Fatalf("write signing key: %v", err)
	}

	pds, err := NewHoldPDS(ctx, "did:web:hold.example.com", "https://hold.example.com",
		"https://atcr.io", filepath.Join(tmpDir, "pds.sqlite3"), keyPath, false)
	if err != nil {
		t.Fatalf("create test PDS: %v", err)
	}
	if err := pds.repomgr.InitNewActor(ctx, pds.uid, "", pds.did, "", "", ""); err != nil {
		t.Fatalf("initialize test repo: %v", err)
	}
	t.Cleanup(func() { pds.Close() })

	return pds
}

// newTestScanSubscriber mirrors what Subscribe builds, registered with the
// broadcaster so Unsubscribe finds it.
func newTestScanSubscriber(t *testing.T, sb *ScanBroadcaster, bufSize int) *ScanSubscriber {
	t.Helper()

	sub := &ScanSubscriber{
		conn: nil, // no websocket needed; handleWriter is not exercised here
		send: make(chan *ScanJobEvent, bufSize),
		id:   "test-subscriber",
		done: make(chan struct{}),
		// One worker unless a test says otherwise, which is what a scanner
		// that declares nothing is treated as.
		capacity: 1,
	}

	sb.mu.Lock()
	sb.subscribers = append(sb.subscribers, sub)
	sb.mu.Unlock()

	return sub
}

func seedPendingJobs(t *testing.T, sb *ScanBroadcaster, n int) {
	t.Helper()

	for i := range n {
		_, err := sb.db.Exec(`
			INSERT INTO scan_jobs
				(manifest_digest, repository, tag, user_did, user_handle,
				 hold_did, hold_endpoint, tier, config_json, layers_json, status)
			VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending')
		`, "sha256:deadbeef", "repo", "latest", "did:plc:user", "user.example.com",
			sb.holdDID, sb.holdEndpoint, "deckhand", "{}", "[]")
		if err != nil {
			t.Fatalf("seed job %d: %v", i, err)
		}
	}
}

// TestScanUnsubscribe_IsIdempotent covers a panic reachable on any scanner that
// dropped mid-write. Unsubscribe used to close sub.send unconditionally, but it
// is called from two places — handleWriter on write error, and handleReader in
// its defer — so a write failure followed by the read side unwinding produced
// "panic: close of closed channel". The slice-removal loop had no guard, so the
// second call fell straight through to the close.
func TestScanUnsubscribe_IsIdempotent(t *testing.T) {
	sb := newTestScanBroadcaster(t)
	sub := newTestScanSubscriber(t, sb, 4)

	sb.Unsubscribe(sub)
	sb.Unsubscribe(sub) // must be a no-op, not a second close

	select {
	case <-sub.done:
	default:
		t.Error("done was not closed by Unsubscribe")
	}

	sb.mu.RLock()
	n := len(sb.subscribers)
	sb.mu.RUnlock()
	if n != 0 {
		t.Errorf("expected subscriber removed, got %d remaining", n)
	}
}

// TestScanUnsubscribe_MarksItsOwnJobsOnce verifies the idempotency guard
// protects the disconnect bookkeeping too, and that the bookkeeping is scoped
// to this subscriber's rows.
//
// Unsubscribe used to flip every assigned and processing row straight back to
// 'pending', which handed a running scan to whichever process connected next
// (see TestScanUnsubscribe_DoesNotImmediatelyHandOffWorkStillRunning). It now
// marks the disconnect and leaves the work where it is. Either way the guard
// matters for the same reason: a dropped scanner unwinds both handleWriter and
// handleReader, and each calls Unsubscribe.
func TestScanUnsubscribe_MarksItsOwnJobsOnce(t *testing.T) {
	sb := newTestScanBroadcaster(t)
	sub := newTestScanSubscriber(t, sb, 4)
	seedPendingJobs(t, sb, 1)

	if _, err := sb.db.Exec(`UPDATE scan_jobs SET status='processing', assigned_to=?`, sub.id); err != nil {
		t.Fatalf("assign: %v", err)
	}

	sb.Unsubscribe(sub)

	var (
		status         string
		assignedTo     sql.NullString
		disconnectedAt sql.NullTime
	)
	row := func() {
		t.Helper()
		if err := sb.db.QueryRow(
			`SELECT status, assigned_to, disconnected_at FROM scan_jobs LIMIT 1`,
		).Scan(&status, &assignedTo, &disconnectedAt); err != nil {
			t.Fatalf("query: %v", err)
		}
	}

	row()
	if status != "processing" || assignedTo.String != sub.id {
		t.Errorf("job = %q/%q, want processing/%q: a disconnect is not evidence "+
			"the scanner stopped scanning", status, assignedTo.String, sub.id)
	}
	if !disconnectedAt.Valid {
		t.Error("the row was not marked as belonging to a disconnected scanner")
	}
	firstMark := disconnectedAt.Time

	// Hand the job to a "replacement" scanner, then unsubscribe the dead one
	// again. The guard must stop it from touching a row that has moved on.
	if _, err := sb.db.Exec(
		`UPDATE scan_jobs SET status='assigned', assigned_to=?, disconnected_at=NULL`, "replacement",
	); err != nil {
		t.Fatalf("reassign: %v", err)
	}
	sb.Unsubscribe(sub)

	row()
	if status != "assigned" || assignedTo.String != "replacement" || disconnectedAt.Valid {
		t.Errorf("second Unsubscribe touched the replacement's job: %q/%q "+
			"disconnected=%v", status, assignedTo.String, disconnectedAt.Valid)
	}
	_ = firstMark
}

// TestScanDrainPendingJobs_ConcurrentUnsubscribe is the regression test for
// "panic: send on closed channel" in the scanner path — the same defect as the
// firehose backfill. Subscribe spawns drainPendingJobs in its own goroutine and
// it sends to sub.send without holding sb.mu, so a scanner disconnecting during
// the drain had Unsubscribe close the channel under an in-flight send.
//
// The subscriber uses a 1-slot buffer with no reader so the drain is reliably
// blocked in the send when Unsubscribe fires.
func TestScanDrainPendingJobs_ConcurrentUnsubscribe(t *testing.T) {
	for range 25 {
		func() {
			sb := newTestScanBroadcaster(t)
			sub := newTestScanSubscriber(t, sb, 1)
			sub.capacity = 50 // the drain must walk rows, not stop at capacity
			seedPendingJobs(t, sb, 50)

			done := make(chan struct{})
			go func() {
				defer close(done)
				// Pre-fix this panicked instead of returning.
				sb.drainPendingJobs(sub, 0)
			}()

			sb.Unsubscribe(sub)

			select {
			case <-done:
			case <-time.After(10 * time.Second):
				t.Fatal("drainPendingJobs did not return after Unsubscribe")
			}
		}()
	}
}

// TestScanDispatchJob_AfterUnsubscribe verifies a removed scanner stops being
// dispatched to and that dispatch does not touch its channel.
func TestScanDispatchJob_AfterUnsubscribe(t *testing.T) {
	sb := newTestScanBroadcaster(t)
	sub := newTestScanSubscriber(t, sb, 4)
	seedPendingJobs(t, sb, 1)

	sb.Unsubscribe(sub)

	sb.dispatchJob(&ScanJobEvent{Type: "job", Seq: 1, Repository: "repo"})

	if len(sub.send) != 0 {
		t.Errorf("unsubscribed scanner received %d jobs", len(sub.send))
	}
}