at-container-registry/pkg/hold/labeler/subscriber.go

package labeler

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"strings"
	"time"

	"atcr.io/pkg/atproto"

	comatproto "github.com/bluesky-social/indigo/api/atproto"
	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/bluesky-social/indigo/events"
	"github.com/gorilla/websocket"
)

// TakedownLabelValue is the label value the hold treats as a takedown trigger.
// Mirrors what pkg/labeler/takedown.go emits.
const TakedownLabelValue = "!takedown"

// Purger is the subset of HoldPDS the subscriber needs to act on takedowns.
// Defined as an interface so tests can substitute a stub without standing up
// a full PDS, and to avoid an import cycle with pkg/hold/pds.
type Purger interface {
	PurgeManifestRecords(ctx context.Context, manifestURI string) (PurgeOutcome, error)
	PurgeUserManifests(ctx context.Context, userDID string) (PurgeOutcome, error)
}

// PurgeOutcome mirrors pds.PurgeResult without creating an import cycle.
type PurgeOutcome struct {
	LayersDeleted      int
	ScanDeleted        bool
	ImageConfigDeleted bool
}

// Subscriber connects to a labeler's subscribeLabels endpoint, mirrors
// takedowns into the local cache, and triggers record purges on the hold.
type Subscriber struct {
	labelerDID string
	cache      *Cache
	purger     Purger
	stopCh     chan struct{}
}

// NewSubscriber builds a subscriber for the given labeler DID (did:plc or
// did:web). The websocket endpoint is resolved on each (re)connect through
// the shared identity directory's #atproto_labeler service entry, so the
// labeler can move (or fix a misconfigured endpoint) without clients
// redeploying.
func NewSubscriber(labelerDID string, cache *Cache, purger Purger) *Subscriber {
	return &Subscriber{
		labelerDID: labelerDID,
		cache:      cache,
		purger:     purger,
		stopCh:     make(chan struct{}),
	}
}

// LabelerDID returns the DID derived from the labeler URL. Useful for the
// caller to log the trusted source.
func (s *Subscriber) LabelerDID() string { return s.labelerDID }

// Start runs the subscription loop in a goroutine.
func (s *Subscriber) Start() {
	go s.run()
}

// Stop signals the subscriber to shut down. Safe to call once.
func (s *Subscriber) Stop() {
	close(s.stopCh)
}

func (s *Subscriber) run() {
	backoff := time.Second
	for {
		select {
		case <-s.stopCh:
			return
		default:
		}

		if err := s.connect(); err != nil {
			slog.Warn("Hold labeler subscription error, reconnecting",
				"labeler", s.labelerDID,
				"error", err,
				"backoff", backoff,
			)
			select {
			case <-s.stopCh:
				return
			case <-time.After(backoff):
			}
			if backoff < 30*time.Second {
				backoff *= 2
			}
		} else {
			backoff = time.Second
		}
	}
}

func (s *Subscriber) connect() error {
	cursor, err := s.cache.GetCursor(s.labelerDID)
	if err != nil {
		return fmt.Errorf("get cursor: %w", err)
	}

	resolveCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	httpURL, err := resolveLabelerURL(resolveCtx, s.labelerDID)
	cancel()
	if err != nil {
		return fmt.Errorf("resolve labeler endpoint: %w", err)
	}

	wsURL := toWebSocketURL(httpURL) + "/xrpc/com.atproto.label.subscribeLabels"
	if cursor > 0 {
		wsURL += fmt.Sprintf("?cursor=%d", cursor)
	}

	slog.Info("Hold connecting to labeler", "url", wsURL, "cursor", cursor)
	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		return fmt.Errorf("websocket dial: %w", err)
	}
	defer conn.Close()
	slog.Info("Hold connected to labeler", "url", httpURL)

	for {
		select {
		case <-s.stopCh:
			return nil
		default:
		}

		mt, payload, err := conn.ReadMessage()
		if err != nil {
			return fmt.Errorf("read: %w", err)
		}
		if mt != websocket.BinaryMessage {
			slog.Warn("Hold labeler: ignoring non-binary frame", "type", mt)
			continue
		}

		seq, labels, err := decodeFrame(payload)
		if err != nil {
			if errors.Is(err, errInfoFrame) {
				continue
			}
			return fmt.Errorf("decode frame: %w", err)
		}

		for _, lbl := range labels {
			s.applyLabel(seq, lbl)
		}

		if err := s.cache.SetCursor(s.labelerDID, seq); err != nil {
			slog.Warn("Hold labeler: failed to persist cursor", "seq", seq, "error", err)
		}
	}
}

// applyLabel processes one label. Only takedown labels from a trusted source
// trigger cache mutations and record purges; everything else is ignored.
func (s *Subscriber) applyLabel(seq int64, lbl *comatproto.LabelDefs_Label) {
	if lbl == nil {
		return
	}
	if lbl.Val != TakedownLabelValue {
		return
	}
	if !s.trustsSource(lbl.Src) {
		slog.Debug("Hold labeler: ignoring untrusted source",
			"src", lbl.Src, "uri", lbl.Uri, "seq", seq)
		return
	}

	cts, _ := time.Parse(time.RFC3339, lbl.Cts)
	negated := lbl.Neg != nil && *lbl.Neg

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	if negated {
		if err := s.cache.Negate(lbl.Uri); err != nil {
			slog.Warn("Hold labeler: failed to drop takedown from cache",
				"uri", lbl.Uri, "error", err)
			return
		}
		slog.Info("Hold labeler: takedown reversed", "uri", lbl.Uri, "seq", seq)
		return
	}

	if err := s.cache.Set(lbl.Uri, lbl.Src, cts); err != nil {
		slog.Warn("Hold labeler: failed to record takedown",
			"uri", lbl.Uri, "error", err)
		return
	}

	// User-level vs per-record. The labeler emits per-record labels for every
	// individual manifest in a repo-level takedown plus a summary; for a
	// user-level takedown only at://<did> is emitted. We dispatch on shape so
	// we don't try to PurgeManifestRecords on a non-manifest URI.
	switch shape := classifyURI(lbl.Uri); shape.Kind {
	case uriKindManifest:
		out, err := s.purger.PurgeManifestRecords(ctx, lbl.Uri)
		if err != nil {
			slog.Warn("Hold labeler: purge failed", "uri", lbl.Uri, "error", err)
			return
		}
		slog.Info("Hold labeler: purged manifest on takedown",
			"uri", lbl.Uri, "layers", out.LayersDeleted,
			"scan", out.ScanDeleted, "config", out.ImageConfigDeleted)
	case uriKindUser:
		out, err := s.purger.PurgeUserManifests(ctx, shape.DID)
		if err != nil {
			slog.Warn("Hold labeler: user-level purge failed",
				"did", shape.DID, "error", err)
			return
		}
		slog.Info("Hold labeler: purged all manifests for user on takedown",
			"did", shape.DID, "layers", out.LayersDeleted)
	default:
		// Repo-level summary URI (at://did/io.atcr.repo/<repo>) and other
		// non-record subjects: cached for GC's reachability check, but no
		// records to purge directly. The per-manifest labels in the same
		// stream do the actual record removal.
		slog.Debug("Hold labeler: cached non-record takedown",
			"uri", lbl.Uri, "kind", string(shape.Kind))
	}
}

// uriKind classifies the shape of a label subject URI.
type uriKind string

const (
	uriKindManifest uriKind = "manifest"
	uriKindUser     uriKind = "user"
	uriKindOther    uriKind = "other"
)

type uriShape struct {
	Kind uriKind
	DID  string
}

// classifyURI inspects an at:// URI and reports whether it points at a single
// manifest record, an entire user, or something else (repo summary etc.).
func classifyURI(uri string) uriShape {
	const prefix = "at://"
	if !strings.HasPrefix(uri, prefix) {
		return uriShape{Kind: uriKindOther}
	}
	rest := uri[len(prefix):]
	parts := strings.SplitN(rest, "/", 3)
	if len(parts) == 0 || parts[0] == "" {
		return uriShape{Kind: uriKindOther}
	}
	did := parts[0]
	if len(parts) == 1 {
		return uriShape{Kind: uriKindUser, DID: did}
	}
	if len(parts) == 3 && parts[1] == "io.atcr.manifest" && parts[2] != "" {
		return uriShape{Kind: uriKindManifest, DID: did}
	}
	return uriShape{Kind: uriKindOther, DID: did}
}

// errInfoFrame signals that the frame was an info frame and the caller should
// continue without treating it as a label or an error.
var errInfoFrame = errors.New("hold labeler: info frame")

// decodeFrame parses a single subscribeLabels binary frame.
//
// ATProto event-stream framing is two concatenated CBOR objects: a {op,t}
// header and a body. We dispatch on op/t; for op=1, t="#labels" we return the
// labels body. #info frames are logged and signal errInfoFrame so the caller
// loops; error frames become Go errors so the run loop reconnects.
func decodeFrame(payload []byte) (int64, []*comatproto.LabelDefs_Label, error) {
	r := bytes.NewReader(payload)
	var header events.EventHeader
	if err := header.UnmarshalCBOR(r); err != nil {
		return 0, nil, fmt.Errorf("unmarshal header: %w", err)
	}

	switch {
	case header.Op == events.EvtKindErrorFrame:
		var ef events.ErrorFrame
		if err := ef.UnmarshalCBOR(r); err != nil {
			return 0, nil, fmt.Errorf("unmarshal error frame: %w", err)
		}
		return 0, nil, fmt.Errorf("labeler error frame: %s — %s", ef.Error, ef.Message)
	case header.Op == events.EvtKindMessage && header.MsgType == "#labels":
		var body comatproto.LabelSubscribeLabels_Labels
		if err := body.UnmarshalCBOR(r); err != nil {
			return 0, nil, fmt.Errorf("unmarshal labels body: %w", err)
		}
		return body.Seq, body.Labels, nil
	case header.Op == events.EvtKindMessage && header.MsgType == "#info":
		var info comatproto.LabelSubscribeLabels_Info
		if err := info.UnmarshalCBOR(r); err != nil {
			return 0, nil, fmt.Errorf("unmarshal info body: %w", err)
		}
		message := ""
		if info.Message != nil {
			message = *info.Message
		}
		slog.Info("Hold labeler: info frame", "name", info.Name, "message", message)
		return 0, nil, errInfoFrame
	default:
		return 0, nil, fmt.Errorf("unexpected frame op=%d t=%q", header.Op, header.MsgType)
	}
}

// trustsSource reports whether labels with the given Src DID should be
// honored. Today this is "matches the configured labeler DID" — a single
// trusted source. If we ever want a list, this becomes a set membership
// check without touching callers.
func (s *Subscriber) trustsSource(src string) bool {
	return src == s.labelerDID
}

// resolveLabelerURL resolves a labeler DID to its HTTP(S) endpoint by looking
// up the #atproto_labeler service in the shared identity directory: did:plc
// via plc.directory, did:web via /.well-known/did.json. The directory is the
// source of truth — clients don't need redeploying when the labeler moves or
// fixes a misconfigured endpoint.
func resolveLabelerURL(ctx context.Context, labelerDID string) (string, error) {
	parsed, err := syntax.ParseDID(labelerDID)
	if err != nil {
		return "", fmt.Errorf("labeler: invalid DID %q: %w", labelerDID, err)
	}
	ident, err := atproto.GetDirectory().LookupDID(ctx, parsed)
	if err != nil {
		return "", fmt.Errorf("labeler: failed to resolve %s: %w", labelerDID, err)
	}
	endpoint := ident.GetServiceEndpoint("atproto_labeler")
	if endpoint == "" {
		return "", fmt.Errorf("labeler: %s has no #atproto_labeler service endpoint", labelerDID)
	}
	return endpoint, nil
}

// toWebSocketURL converts an HTTP URL to a WebSocket URL. http→ws, https→wss.
func toWebSocketURL(httpURL string) string {
	u, err := url.Parse(httpURL)
	if err != nil {
		return httpURL
	}
	switch u.Scheme {
	case "https":
		u.Scheme = "wss"
	default:
		u.Scheme = "ws"
	}
	return u.String()
}