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 (
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"

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

// httpClient is used for blob downloads and presigned URL requests
// with a timeout to prevent stalled connections from leaking memory.
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{}
}

// 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{}),
	}
}

// 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))
	}
	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")

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

		var raw scanner.ScanJobRaw
		if err := json.Unmarshal(data, &raw); err != nil {
			slog.Error("Failed to unmarshal message", "error", err)
			continue
		}

		if raw.Type != "job" {
			slog.Warn("Unknown message type from hold", "type", raw.Type)
			continue
		}

		// Parse config and layers from raw JSON
		var config scanner.BlobDescriptor
		if err := json.Unmarshal(raw.Config, &config); err != nil {
			slog.Error("Failed to unmarshal config", "seq", raw.Seq, "error", err)
			continue
		}

		var layers []scanner.BlobDescriptor
		if err := json.Unmarshal(raw.Layers, &layers); err != nil {
			slog.Error("Failed to unmarshal layers", "seq", raw.Seq, "error", err)
			continue
		}

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

// SendAck sends an acknowledgement for a received job
func (c *HoldClient) SendAck(seq int64) {
	c.sendJSON(scanner.AckMessage{Type: "ack", 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})
}

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.WriteJSON(v); err != nil {
		slog.Error("Failed to send WebSocket message", "error", err)
	}
}

// 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.
func GetBlobPresignedURL(holdEndpoint, holdDID, 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))

	req, err := http.NewRequest("GET", 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
}

// DownloadBlob downloads a blob from a presigned URL to a local file
func DownloadBlob(presignedURL, destPath string) error {
	resp, err := httpClient.Get(presignedURL)
	if err != nil {
		return fmt.Errorf("failed to download blob: %w", err)
	}
	defer resp.Body.Close()

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

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

	if _, err := io.Copy(out, resp.Body); err != nil {
		return fmt.Errorf("failed to write blob: %w", err)
	}

	return nil
}