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

// Package labeler provides a subscription client for consuming labels
// from an ATProto labeler service.
package labeler

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

	"atcr.io/pkg/appview/db"
	"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 only label value the appview honors.
const TakedownLabelValue = "!takedown"

// Subscriber connects to a labeler's subscribeLabels endpoint and mirrors
// the current set of active takedowns into the appview database.
type Subscriber struct {
	labelerDID string
	database   *sql.DB
	stopCh     chan struct{}
}

// NewSubscriber creates a new labeler subscriber. labelerDID is a did:plc or
// did:web identifier. The websocket endpoint is resolved on each (re)connect
// via 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, database *sql.DB) *Subscriber {
	return &Subscriber{
		labelerDID: labelerDID,
		database:   database,
		stopCh:     make(chan struct{}),
	}
}

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

// Stop signals the subscriber to shut down.
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("Labeler subscription error, reconnecting",
				"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 := db.GetCursor(s.database, s.labelerDID)
	if err != nil {
		return fmt.Errorf("failed to 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("Connecting to labeler", "url", wsURL, "cursor", cursor)

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

	slog.Info("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 error: %w", err)
		}
		// Per the ATProto event-stream spec each frame is a binary message; reject text.
		if mt != websocket.BinaryMessage {
			slog.Warn("Ignoring non-binary frame from labeler", "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 _, le := range labels {
			s.applyLabel(le)
		}

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

// applyLabel processes a single label. The appview only honors !takedown labels
// from the configured labeler, and only at the granularity it can enforce —
// user-level (at://<did>) and repo summary (at://<did>/io.atcr.repo/<repo>).
// Per-record labels (per manifest, tag, repo-page) are dropped; the registry
// middleware gates per (did, repo) so finer granularity has no effect.
func (s *Subscriber) applyLabel(le *comatproto.LabelDefs_Label) {
	if le == nil {
		return
	}
	if le.Val != TakedownLabelValue {
		return
	}
	if le.Src != s.labelerDID {
		slog.Debug("Ignoring label from untrusted source", "src", le.Src, "uri", le.Uri)
		return
	}

	shape := classifyURI(le.Uri)
	if shape.kind == uriOther {
		slog.Debug("Skipping non-enforced label", "uri", le.Uri)
		return
	}

	negated := le.Neg != nil && *le.Neg
	if negated {
		if err := db.RemoveTakedown(s.database, le.Src, shape.did, shape.repo); err != nil {
			slog.Warn("Failed to remove takedown", "uri", le.Uri, "error", err)
			return
		}
		slog.Info("Mirrored takedown reversal",
			"src", le.Src, "did", shape.did, "repo", shape.repo)
		return
	}

	cts, _ := time.Parse(time.RFC3339, le.Cts)
	if err := db.SetTakedown(s.database, le.Src, shape.did, shape.repo, cts); err != nil {
		slog.Warn("Failed to record takedown", "uri", le.Uri, "error", err)
		return
	}
	slog.Info("Mirrored takedown",
		"src", le.Src, "did", shape.did, "repo", shape.repo)
}

// uriShape captures the parts of a label subject URI that the appview cares about.
type uriShape struct {
	kind uriKind
	did  string
	repo string
}

type uriKind int

const (
	uriOther uriKind = iota
	uriUserLevel
	uriRepoSummary
)

// classifyURI reports whether the URI is a user-level subject (at://<did>),
// a repo summary (at://<did>/io.atcr.repo/<repo>), or something else
// (per-record manifest/tag/repo-page labels we don't enforce).
func classifyURI(uri string) uriShape {
	const prefix = "at://"
	if !strings.HasPrefix(uri, prefix) {
		return uriShape{}
	}
	rest := uri[len(prefix):]
	parts := strings.SplitN(rest, "/", 3)
	if len(parts) == 0 || parts[0] == "" {
		return uriShape{}
	}
	did := parts[0]
	if len(parts) == 1 {
		return uriShape{kind: uriUserLevel, did: did}
	}
	if len(parts) == 3 && parts[1] == "io.atcr.repo" && parts[2] != "" {
		return uriShape{kind: uriRepoSummary, did: did, repo: parts[2]}
	}
	return uriShape{kind: uriOther, did: did}
}

// errInfoFrame is returned by decodeFrame when the frame is informational and the
// caller should just continue to the next message.
var errInfoFrame = errors.New("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 the
// header op/t pair and return the labels body for op=1, t="#labels". For #info frames
// we log and signal errInfoFrame so the caller skips. Error frames (op=-1) become Go
// errors so the run loop reconnects with backoff.
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("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)
	}
}

// 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.
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()
}

// SubscriberFromConfig creates a Subscriber from a labeler DID config value.
// Returns nil if labelerDID is empty.
func SubscriberFromConfig(labelerDID string, database *sql.DB) *Subscriber {
	if labelerDID == "" {
		return nil
	}
	return NewSubscriber(labelerDID, database)
}