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

package client

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

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

// upgradeOnly accepts a WebSocket and hands the server side to onConn, which
// decides how lifelike the peer is. The returned URL is what a HoldClient dials.
func upgradeOnly(t *testing.T, onConn func(*websocket.Conn)) string {
	t.Helper()

	upgrader := websocket.Upgrader{}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		onConn(conn)
	}))
	t.Cleanup(srv.Close)
	return "ws" + strings.TrimPrefix(srv.URL, "http")
}

func newTestHoldClient(t *testing.T, url string) *HoldClient {
	t.Helper()

	c := NewHoldClient(url, "secret", queue.NewJobQueue(8))
	c.pingInterval = 20 * time.Millisecond
	c.pongWait = 200 * time.Millisecond
	c.writeWait = time.Second
	t.Cleanup(func() {
		defer func() { _ = recover() }() // Close panics if already closed
		c.Close()
	})
	return c
}

// TestHoldClient_DetectsHalfOpenConnection is the scanner half of F7.
//
// connectOnce sat in ReadMessage with no deadline and sent no pings, so a hold
// that went away without closing the TCP connection left this scanner
// permanently "connected": it never reconnects, never re-drains the hold's
// pending rows, and every result it computes is written into a socket that goes
// nowhere. The peer here upgrades and then does nothing at all — it never reads,
// so it never pongs.
func TestHoldClient_DetectsHalfOpenConnection(t *testing.T) {
	held := make(chan struct{})
	url := upgradeOnly(t, func(conn *websocket.Conn) {
		<-held // hold the connection open, answering nothing
		_ = conn.Close()
	})
	t.Cleanup(func() { close(held) })

	c := newTestHoldClient(t, url)

	errCh := make(chan error, 1)
	go func() { errCh <- c.connectOnce(-1) }()

	select {
	case err := <-errCh:
		if err == nil {
			t.Fatal("connectOnce returned nil for a peer that answered nothing")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("a hold that answers no pings never disconnected the scanner: " +
			"the client waits in ReadMessage with no deadline and sends no pings")
	}
}

// The other side of the rule. An idle hold is the normal state — the dfd604b
// commit message records a hold and scanner that shared one WebSocket for a
// week — so silence alone must not drop the connection. gorilla answers a ping
// from inside ReadMessage, which is what the real hold's reader does.
func TestHoldClient_StaysConnectedWhileHoldAnswers(t *testing.T) {
	url := upgradeOnly(t, func(conn *websocket.Conn) {
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				_ = conn.Close()
				return
			}
		}
	})

	c := newTestHoldClient(t, url)

	errCh := make(chan error, 1)
	go func() { errCh <- c.connectOnce(-1) }()

	// Many pongWait periods with no traffic but pings.
	select {
	case err := <-errCh:
		t.Fatalf("an idle but responsive hold was dropped: %v", err)
	case <-time.After(time.Second):
	}

	c.Close()
	select {
	case <-errCh:
	case <-time.After(5 * time.Second):
		t.Fatal("connectOnce did not return after Close")
	}
}

// A job frame is small — a config descriptor and a layer list — so the read
// limit that bounds a hostile or broken hold must still admit anything the real
// one sends with room to spare.
func TestHoldClient_AcceptsARealisticJobFrame(t *testing.T) {
	delivered := make(chan struct{})
	url := upgradeOnly(t, func(conn *websocket.Conn) {
		// A 128-layer image, which is well past the OCI-practical limit.
		var layers []string
		for i := 0; i < 128; i++ {
			layers = append(layers,
				`{"digest":"sha256:`+strings.Repeat("a", 64)+`","size":123456,"mediaType":"application/vnd.oci.image.layer.v1.tar+gzip"}`)
		}
		frame := `{"type":"job","seq":1,"manifestDigest":"sha256:` + strings.Repeat("b", 64) + `",` +
			`"repository":"repo","tag":"latest","userDid":"did:plc:user","userHandle":"user.example.com",` +
			`"holdDid":"did:web:hold.example.com","holdEndpoint":"https://hold.example.com","tier":"deckhand",` +
			`"config":{"digest":"sha256:` + strings.Repeat("c", 64) + `","size":1024,"mediaType":"application/vnd.oci.image.config.v1+json"},` +
			`"layers":[` + strings.Join(layers, ",") + `]}`
		if err := conn.WriteMessage(websocket.TextMessage, []byte(frame)); err != nil {
			return
		}
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				_ = conn.Close()
				return
			}
			select {
			case delivered <- struct{}{}:
			default:
			}
		}
	})

	c := newTestHoldClient(t, url)
	c.pongWait = 5 * time.Second
	go func() { _ = c.connectOnce(-1) }()

	select {
	case <-delivered:
	case <-time.After(5 * time.Second):
		t.Fatal("a realistic job frame was never acknowledged: the read limit rejected it")
	}

	if c.queue.Len() != 1 {
		t.Errorf("queue depth = %d, want 1: the job frame did not survive the read", c.queue.Len())
	}
}