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 (
	"database/sql"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/url"
	"strings"
	"time"

	"atcr.io/pkg/appview/db"

	"github.com/gorilla/websocket"
)

// LabelsMessage is the wire format for subscribeLabels events.
type LabelsMessage struct {
	Seq    int64        `json:"seq"`
	Labels []LabelEvent `json:"labels"`
}

// LabelEvent is a single label from the labeler.
type LabelEvent struct {
	Src string `json:"src"`
	URI string `json:"uri"`
	CID string `json:"cid,omitempty"`
	Val string `json:"val"`
	Neg bool   `json:"neg"`
	Cts string `json:"cts"`
	Exp string `json:"exp,omitempty"`
}

// Subscriber connects to a labeler's subscribeLabels endpoint
// and mirrors labels into the appview database.
type Subscriber struct {
	labelerURL string
	database   *sql.DB
	stopCh     chan struct{}
}

// NewSubscriber creates a new labeler subscriber.
func NewSubscriber(labelerURL string, database *sql.DB) *Subscriber {
	return &Subscriber{
		labelerURL: labelerURL,
		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 {
	// Get cursor from DB
	// Use the labeler URL as src identifier
	labelerDID := extractDIDFromURL(s.labelerURL)
	cursor, err := db.GetLabelCursor(s.database, labelerDID)
	if err != nil {
		return fmt.Errorf("failed to get cursor: %w", err)
	}

	// Build WebSocket URL
	wsURL := toWebSocketURL(s.labelerURL) + "/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", s.labelerURL)

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

		var msg LabelsMessage
		if err := conn.ReadJSON(&msg); err != nil {
			return fmt.Errorf("read error: %w", err)
		}

		for _, le := range msg.Labels {
			cts, _ := time.Parse(time.RFC3339, le.Cts)
			did, repo := extractSubjectFromURI(le.URI)

			label := &db.Label{
				Src:         le.Src,
				URI:         le.URI,
				Val:         le.Val,
				Neg:         le.Neg,
				Cts:         cts,
				SubjectDID:  did,
				SubjectRepo: repo,
				Seq:         msg.Seq,
			}

			if err := db.UpsertLabel(s.database, label); err != nil {
				slog.Warn("Failed to upsert label", "uri", le.URI, "error", err)
				continue
			}

			slog.Info("Mirrored label",
				"uri", le.URI,
				"val", le.Val,
				"neg", le.Neg,
				"subject_did", did,
				"subject_repo", repo,
			)
		}
	}
}

// extractSubjectFromURI extracts the DID and repository from an AT URI.
// Examples:
//
//	at://did:plc:xyz                          → (did:plc:xyz, "")
//	at://did:plc:xyz/io.atcr.manifest/abc     → (did:plc:xyz, "")  - repo extracted from record
//	at://did:plc:xyz/io.atcr.repo/myimage     → (did:plc:xyz, "myimage")
func extractSubjectFromURI(uri string) (did, repo string) {
	trimmed := strings.TrimPrefix(uri, "at://")
	parts := strings.SplitN(trimmed, "/", 3)
	if len(parts) == 0 {
		return "", ""
	}
	did = parts[0]

	// For repo-level summary labels: at://did/io.atcr.repo/reponame
	if len(parts) >= 3 && parts[1] == "io.atcr.repo" {
		repo = parts[2]
	}
	return did, repo
}

// extractDIDFromURL derives a did:web from a labeler URL.
func extractDIDFromURL(labelerURL string) string {
	u, err := url.Parse(labelerURL)
	if err != nil {
		return labelerURL
	}
	host := u.Hostname()
	if port := u.Port(); port != "" {
		host += "%3A" + port
	}
	return "did:web:" + host
}

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

// ParseLabelerURL parses a labeler DID or URL into an HTTP URL.
func ParseLabelerURL(labelerDIDOrURL string) string {
	if strings.HasPrefix(labelerDIDOrURL, "http://") || strings.HasPrefix(labelerDIDOrURL, "https://") {
		return labelerDIDOrURL
	}
	if strings.HasPrefix(labelerDIDOrURL, "did:web:") {
		host := strings.TrimPrefix(labelerDIDOrURL, "did:web:")
		host = strings.ReplaceAll(host, "%3A", ":")
		return "https://" + host
	}
	return labelerDIDOrURL
}

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

// DecodeLabelsFromJSON decodes a JSON-encoded labels message.
func DecodeLabelsFromJSON(data []byte) (*LabelsMessage, error) {
	var msg LabelsMessage
	if err := json.Unmarshal(data, &msg); err != nil {
		return nil, err
	}
	return &msg, nil
}