at-container-registry/scanner/internal/client/hold.go

// Package client implements the bidirectional WebSocket client for communicating
// with the hold service, plus HTTP helpers for downloading blobs.
package client

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"

	scanner "atcr.io/scanner"
	"atcr.io/scanner/internal/queue"
	"github.com/gorilla/websocket"
)

const (
	// holdPingInterval is how often the scanner pings an idle hold connection.
	holdPingInterval = 30 * time.Second

	// holdPongWait is how long the connection may go without a pong (or any
	// other frame) before the scanner treats the hold as gone. Three missed
	// pings.
	holdPongWait = 90 * time.Second

	// holdWriteWait bounds a single write to the hold.
	holdWriteWait = 30 * time.Second

	// holdMaxMessageSize caps one frame from the hold. gorilla's default is
	// unlimited, so without it a broken or hostile hold decides how much this
	// process allocates.
	//
	// A job frame is small — a config descriptor and a layer list — so 16 MiB
	// is several orders of magnitude of headroom over anything the real hold
	// sends. Note the asymmetry with the hold's own limit: results travel the
	// other way and carry whole SBOMs inline, so the hold's ceiling has to be
	// far higher than this one.
	holdMaxMessageSize = 16 << 20
)

// httpClient is used for blob downloads and presigned URL requests.
//
// Its Timeout is a per-request backstop and nothing more: a job that fetches a
// config plus nineteen layers gets twenty of these in series, so on its own it
// bounds one request at five minutes and a job at a hundred. The bound that
// matters is the caller's context, which every request below carries; this
// only stops a single wedged request from outliving a job whose deadline the
// operator has disabled.
var httpClient = &http.Client{Timeout: 5 * time.Minute}

// HoldClient manages the WebSocket connection to a hold service
type HoldClient struct {
	holdURL string
	secret  string
	queue   *queue.JobQueue
	conn    *websocket.Conn
	mu      sync.Mutex // protects conn writes
	done    chan struct{}

	// instanceID identifies this scanner process to the hold, and is sent on
	// every connect. It is deliberately per-process and not persisted: a hold
	// hands a reconnecting instance back the jobs its workers never stopped
	// running, and a process that actually restarted has lost that work and
	// must not claim it.
	instanceID string

	// workers is how many scans this process runs at once. The hold uses it as
	// the dispatch budget for this connection; a hold that does not know the
	// parameter ignores it and assumes one.
	workers int

	// WebSocket liveness knobs. Zero means "use the package default"; tests
	// shrink them so a half-open connection can be provoked in milliseconds.
	pingInterval time.Duration
	pongWait     time.Duration
	writeWait    time.Duration
}

func (c *HoldClient) pingEvery() time.Duration {
	if c.pingInterval > 0 {
		return c.pingInterval
	}
	return holdPingInterval
}

func (c *HoldClient) pongDeadline() time.Duration {
	if c.pongWait > 0 {
		return c.pongWait
	}
	return holdPongWait
}

func (c *HoldClient) writeDeadline() time.Duration {
	if c.writeWait > 0 {
		return c.writeWait
	}
	return holdWriteWait
}

// NewHoldClient creates a new hold client
func NewHoldClient(holdURL, secret string, q *queue.JobQueue) *HoldClient {
	return &HoldClient{
		holdURL:    holdURL,
		secret:     secret,
		queue:      q,
		done:       make(chan struct{}),
		instanceID: newInstanceID(),
		workers:    1,
	}
}

// SetWorkers tells the hold how many scans this process runs concurrently, so
// it can keep that many jobs in flight instead of one.
//
// This is a setter rather than a constructor argument on purpose: the
// connection is entirely usable without it, and a caller that never calls it
// gets the single-worker default the hold assumes anyway.
func (c *HoldClient) SetWorkers(n int) {
	if n > 0 {
		c.workers = n
	}
}

func newInstanceID() string {
	b := make([]byte, 8)
	if _, err := rand.Read(b); err != nil {
		// The hold falls back to a per-connection identity when none is sent,
		// which costs resumption on reconnect and nothing else.
		return ""
	}
	return hex.EncodeToString(b)
}

// Connect establishes the WebSocket connection with auto-reconnect
func (c *HoldClient) Connect() {
	var cursor int64 = -1

	for {
		select {
		case <-c.done:
			return
		default:
		}

		err := c.connectOnce(cursor)
		if err != nil {
			slog.Error("WebSocket connection failed, reconnecting",
				"error", err)
		}

		// Exponential backoff with max 30s
		select {
		case <-c.done:
			return
		case <-time.After(5 * time.Second):
		}
	}
}

func (c *HoldClient) connectOnce(cursor int64) error {
	// Build WebSocket URL
	u, err := url.Parse(c.holdURL)
	if err != nil {
		return fmt.Errorf("invalid hold URL: %w", err)
	}

	// Convert http(s) to ws(s)
	switch u.Scheme {
	case "http":
		u.Scheme = "ws"
	case "https":
		u.Scheme = "wss"
	case "ws", "wss":
		// Already correct
	}

	u.Path = "/xrpc/io.atcr.hold.subscribeScanJobs"
	q := u.Query()
	q.Set("secret", c.secret)
	if cursor >= 0 {
		q.Set("cursor", fmt.Sprintf("%d", cursor))
	}
	if c.instanceID != "" {
		q.Set("instance", c.instanceID)
	}
	if c.workers > 0 {
		q.Set("workers", fmt.Sprintf("%d", c.workers))
	}
	u.RawQuery = q.Encode()

	slog.Info("Connecting to hold service", "url", u.Host)

	conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		return fmt.Errorf("dial failed: %w", err)
	}
	defer conn.Close()

	c.mu.Lock()
	c.conn = conn
	c.mu.Unlock()

	slog.Info("Connected to hold service")

	// Liveness. This loop used to sit in ReadMessage with no deadline and send
	// nothing, so a hold that went away without closing the TCP connection left
	// this process permanently "connected": never reconnecting, so never
	// re-draining the hold's pending rows, and writing every result it computed
	// into a socket that went nowhere.
	conn.SetReadLimit(holdMaxMessageSize)
	refreshRead := func() {
		if err := conn.SetReadDeadline(time.Now().Add(c.pongDeadline())); err != nil {
			slog.Debug("Failed to set hold read deadline", "error", err)
		}
	}
	refreshRead()
	conn.SetPongHandler(func(string) error {
		refreshRead()
		return nil
	})

	stopPing := make(chan struct{})
	defer close(stopPing)
	go c.pingLoop(conn, stopPing)

	// Read messages from hold
	for {
		_, data, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				slog.Error("WebSocket read error", "error", err)
			}
			return err
		}
		// Any frame is proof of life, not just a pong.
		refreshRead()

		c.handleFrame(data)
		refreshRead()
	}
}

// pingLoop keeps the connection provably alive from this end. gorilla permits
// WriteControl concurrently with every other method, so this needs none of the
// write mutex and cannot be delayed behind a result being sent.
func (c *HoldClient) pingLoop(conn *websocket.Conn, stop <-chan struct{}) {
	ticker := time.NewTicker(c.pingEvery())
	defer ticker.Stop()

	for {
		select {
		case <-stop:
			return
		case <-c.done:
			return
		case <-ticker.C:
			if err := conn.WriteControl(websocket.PingMessage, nil,
				time.Now().Add(c.writeDeadline())); err != nil {
				// Closing is what unblocks the read loop, which then
				// reconnects. Returning quietly would leave it parked.
				slog.Warn("Failed to ping hold, dropping the connection", "error", err)
				_ = conn.Close()
				return
			}
		}
	}
}

// handleFrame decodes one frame from the hold and either enqueues the job or
// answers for it.
//
// The invariant, which every branch below keeps: a frame carrying a usable seq
// is never dropped in silence. The hold writes status='assigned' for that seq
// before it writes the frame, and its only escape from that row is an
// acknowledgement or a verdict from us. A frame we drop instead sits assigned
// until the five-minute ack timeout, is re-offered within thirty seconds — to
// this same scanner, which cannot decode it this time either — and holds the
// hold's single proactive dispatch slot the entire time. One such frame stops
// proactive scanning for every user of that hold, permanently.
func (c *HoldClient) handleFrame(data []byte) {
	var raw scanner.ScanJobRaw
	if err := json.Unmarshal(data, &raw); err != nil {
		// encoding/json records the first type error and keeps decoding the
		// rest of the object, so the seq that addresses the row is often still
		// usable even when the frame as a whole was rejected.
		c.rejectJob(raw.Seq, "malformed job frame", err)
		return
	}

	if raw.Type != "job" {
		// Not a job, so no row of ours is waiting on an answer. The hold sends
		// nothing but "job" today; a future message type is better ignored by
		// an old scanner than answered with a verdict about a job.
		slog.Warn("Unknown message type from hold", "type", raw.Type, "seq", raw.Seq)
		return
	}

	// Parse config and layers from raw JSON
	var config scanner.BlobDescriptor
	if err := json.Unmarshal(raw.Config, &config); err != nil {
		c.rejectJob(raw.Seq, "malformed job config", err)
		return
	}

	var layers []scanner.BlobDescriptor
	if err := json.Unmarshal(raw.Layers, &layers); err != nil {
		c.rejectJob(raw.Seq, "malformed job layers", err)
		return
	}

	job := &scanner.ScanJob{
		Seq:            raw.Seq,
		ManifestDigest: raw.ManifestDigest,
		Repository:     raw.Repository,
		Tag:            raw.Tag,
		UserDID:        raw.UserDID,
		UserHandle:     raw.UserHandle,
		HoldDID:        raw.HoldDID,
		HoldEndpoint:   raw.HoldEndpoint,
		Tier:           raw.Tier,
		Config:         config,
		Layers:         layers,
	}

	// Send ack immediately
	c.SendAck(job.Seq)

	// Enqueue into priority queue
	if !c.queue.Enqueue(job) {
		slog.Warn("Queue full, sending error",
			"seq", job.Seq,
			"repository", job.Repository)
		c.SendError(job.Seq, "scanner queue full")
	}
}

// rejectJob answers a job frame this scanner cannot decode.
//
// "skipped" and not "error": the hold treats a failure as transient and
// re-queues it on the rescan interval, but a frame that does not decode will
// not decode on the next attempt either. A skip is terminal — the hold marks
// the row completed, writes a scan record carrying the reason, and releases
// the manifest from its in-flight set — so the job leaves the rotation and the
// reason reaches the user instead of vanishing into this process's log.
//
// A seq of zero is the one frame that cannot be answered: the hold's rows
// start at 1, so there is no job to address. Nothing is stranded by staying
// quiet there either, because the hold assigns the row by seq before it sends,
// and a seq that never survived the wire is not the seq it assigned.
func (c *HoldClient) rejectJob(seq int64, reason string, err error) {
	if seq <= 0 {
		slog.Error("Dropping undecodable frame with no usable seq",
			"reason", reason, "error", err)
		return
	}

	slog.Error("Rejecting undecodable job frame",
		"seq", seq, "reason", reason, "error", err)
	c.SendSkipped(seq, fmt.Sprintf("%s: %v", reason, err))
}

// SendAck sends an acknowledgement for a received job
func (c *HoldClient) SendAck(seq int64) {
	c.sendJSON(scanner.AckMessage{Type: "ack", Seq: seq})
}

// SendStarted tells the hold a worker has begun this scan.
//
// The ack this job already got was sent from the reader goroutine on receipt,
// before the job was queued. This is the signal the hold measures its scanning
// deadline from; a hold too old to know the message logs and ignores it, and
// falls back to its own budget measured from dispatch.
func (c *HoldClient) SendStarted(seq int64) {
	c.sendJSON(scanner.StartedMessage{Type: "started", Seq: seq})
}

// SendResult sends scan results back to the hold
func (c *HoldClient) SendResult(seq int64, result *scanner.ScanResult) {
	msg := scanner.ResultMessage{
		Type:    "result",
		Seq:     seq,
		Summary: result.Summary,
	}
	if result.SBOM != nil {
		msg.SBOM = string(result.SBOM)
	}
	if result.VulnReport != nil {
		msg.VulnReport = string(result.VulnReport)
	}
	c.sendJSON(msg)
}

// SendError sends an error message for a failed scan
func (c *HoldClient) SendError(seq int64, errMsg string) {
	c.sendJSON(scanner.ErrorMessage{Type: "error", Seq: seq, Error: errMsg})
}

// SendSkipped sends a skipped message for an artifact the scanner intentionally
// won't process (e.g., helm charts). Distinct from SendError so the hold can
// distinguish a permanent skip from a retryable failure.
func (c *HoldClient) SendSkipped(seq int64, reason string) {
	c.sendJSON(scanner.SkippedMessage{Type: "skipped", Seq: seq, Reason: reason})
}

// sendJSON writes one message to the hold under the write mutex.
//
// The deadline is what stops a wedged write from taking the process with it:
// a result is megabytes of SBOM, and a hold that has stopped reading fills the
// socket buffer and blocks the write. Without a deadline that write blocks
// until the kernel gives up on the connection, which is on the order of hours,
// with the mutex held the whole time — so every worker that finishes a scan
// queues up behind it and the scanner goes quiet without ever disconnecting.
//
// A failed write closes the connection rather than just logging. The read loop
// is then released, connectOnce returns, and Connect redials; leaving it open
// means the next write blocks the same way against the same dead socket.
func (c *HoldClient) sendJSON(v any) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn == nil {
		slog.Warn("Cannot send, no connection")
		return
	}

	if err := c.conn.SetWriteDeadline(time.Now().Add(c.writeDeadline())); err != nil {
		slog.Error("Failed to set WebSocket write deadline", "error", err)
	}
	if err := c.conn.WriteJSON(v); err != nil {
		slog.Error("Failed to send WebSocket message, dropping the connection", "error", err)
		_ = c.conn.Close()
	}
}

// Close shuts down the client
func (c *HoldClient) Close() {
	close(c.done)
	c.mu.Lock()
	if c.conn != nil {
		c.conn.Close()
	}
	c.mu.Unlock()
}

// GetBlobPresignedURL gets a presigned download URL from the hold service.
// If secret is non-empty, it is sent as a Bearer token for private hold access.
//
// The digest is a scanner.Digest rather than a string so that only a validated
// "sha256:<hex>" can ever be asked for. Callers parse once, at the boundary,
// and the same value then names the blob on the wire and the file on disk.
func GetBlobPresignedURL(ctx context.Context, holdEndpoint, holdDID string, digest scanner.Digest, secret string) (string, error) {
	reqURL := fmt.Sprintf("%s/xrpc/com.atproto.sync.getBlob?did=%s&cid=%s&method=GET",
		holdEndpoint,
		url.QueryEscape(holdDID),
		url.QueryEscape(digest.String()))

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
	if err != nil {
		return "", fmt.Errorf("failed to create request: %w", err)
	}
	if secret != "" {
		req.Header.Set("Authorization", "Bearer "+secret)
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("failed to get presigned URL: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("hold returned status %d: %s", resp.StatusCode, string(body))
	}

	var result struct {
		URL string `json:"url"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", fmt.Errorf("failed to decode response: %w", err)
	}

	return result.URL, nil
}

// Errors a blob download can fail with that no retry can fix. Both are
// decided by bytes the scanner has already seen, so a caller mapping them to a
// permanent skip is not guessing.
var (
	// ErrBlobCorrupt means the bytes that arrived are not the bytes the
	// descriptor described. It deliberately does not distinguish "the digest
	// is wrong" from "the size is wrong": both fields come from the same
	// user-writable manifest record, so a disagreement condemns the whole
	// descriptor rather than proving anything about one field.
	ErrBlobCorrupt = errors.New("digest mismatch")

	// ErrBlobTooLarge means the transfer would take the job past the byte
	// ceiling the configuration set.
	ErrBlobTooLarge = errors.New("image too large")
)

// BlobExpectation is what a download has to turn out to be.
type BlobExpectation struct {
	// Digest is the validated digest the bytes must hash to. This is the
	// authority: it is checked always, and it is what makes serving the wrong
	// object detectable.
	Digest scanner.Digest

	// DeclaredSize is the descriptor's size field, corroborating evidence
	// rather than authority. Zero or negative means the record claimed no
	// size, in which case the digest alone decides.
	DeclaredSize int64

	// MaxBytes caps this transfer. Negative means unbounded. Zero is a real
	// ceiling of zero bytes: it is what remains when an earlier blob in the
	// same job has spent the whole budget.
	MaxBytes int64
}

// DownloadBlob downloads a blob from a presigned URL to a local file and
// verifies it against want, returning the number of bytes written.
//
// Verification happens while streaming, not by re-reading the file: the hash
// and the byte count are both accumulated by the same io.Copy that writes the
// blob, so a scan can never catalog bytes that were not checked, and a blob
// over the ceiling stops costing bandwidth one byte past it.
//
// The context bounds the whole transfer, not just the round trip: the request
// body is read under it, so a hold that answers promptly and then dribbles
// bytes forever is cancelled by the job deadline rather than by the client's
// five-minute per-request ceiling.
func DownloadBlob(ctx context.Context, presignedURL, destPath string, want BlobExpectation) (int64, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, presignedURL, nil)
	if err != nil {
		return 0, fmt.Errorf("failed to create download request: %w", err)
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		return 0, fmt.Errorf("failed to download blob: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return 0, fmt.Errorf("download returned status %d", resp.StatusCode)
	}

	out, err := os.Create(destPath)
	if err != nil {
		return 0, fmt.Errorf("failed to create file: %w", err)
	}
	defer out.Close()

	// One byte past the ceiling is all the proof needed that the ceiling was
	// broken, and reading no further is the point: the claimed sizes that let
	// the job through the pre-check are the attacker's to choose.
	body := io.Reader(resp.Body)
	if want.MaxBytes >= 0 {
		body = io.LimitReader(resp.Body, want.MaxBytes+1)
	}

	hasher := sha256.New()
	n, err := io.Copy(io.MultiWriter(out, hasher), body)
	if err != nil {
		return n, fmt.Errorf("failed to write blob: %w", err)
	}

	if want.MaxBytes >= 0 && n > want.MaxBytes {
		return n, fmt.Errorf("%w: blob %s exceeds the %d bytes left in the job's budget",
			ErrBlobTooLarge, want.Digest, want.MaxBytes)
	}

	gotHex := hex.EncodeToString(hasher.Sum(nil))
	sizeDisagrees := want.DeclaredSize > 0 && n != want.DeclaredSize
	if gotHex != want.Digest.Hex || sizeDisagrees {
		claimedSize := "no declared size"
		if want.DeclaredSize > 0 {
			claimedSize = fmt.Sprintf("%d bytes", want.DeclaredSize)
		}
		return n, fmt.Errorf("%w: descriptor claims %s and %s, received %d bytes hashing to %s:%s",
			ErrBlobCorrupt, want.Digest, claimedSize, n, want.Digest.Algorithm, gotHex)
	}

	return n, nil
}