at-container-registry/pkg/labeler/subscribe.go

package labeler

import (
	"encoding/json"
	"log/slog"
	"net/http"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

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

// LabelOutput is the ATProto label format for subscribeLabels/queryLabels output.
type LabelOutput 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"`
}

func labelToOutput(l Label) LabelOutput {
	out := LabelOutput{
		Src: l.Src,
		URI: l.URI,
		CID: l.CID,
		Val: l.Val,
		Neg: l.Neg,
		Cts: l.Cts.UTC().Format(time.RFC3339),
	}
	if l.Exp != nil {
		out.Exp = l.Exp.UTC().Format(time.RFC3339)
	}
	return out
}

// handleSubscribeLabels implements com.atproto.label.subscribeLabels (WebSocket).
func (s *Server) handleSubscribeLabels(w http.ResponseWriter, r *http.Request) {
	cursorStr := r.URL.Query().Get("cursor")
	var cursor int64
	if cursorStr != "" {
		var err error
		cursor, err = strconv.ParseInt(cursorStr, 10, 64)
		if err != nil {
			http.Error(w, "invalid cursor", http.StatusBadRequest)
			return
		}
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		slog.Error("WebSocket upgrade failed", "error", err)
		return
	}
	defer conn.Close()

	slog.Info("subscribeLabels client connected", "cursor", cursor)

	// Send historical labels since cursor
	labels, err := GetLabelsSince(s.db, cursor, 1000)
	if err != nil {
		slog.Error("Failed to get labels", "error", err)
		return
	}

	for _, l := range labels {
		msg := LabelsMessage{
			Seq:    l.ID,
			Labels: []LabelOutput{labelToOutput(l)},
		}
		if err := conn.WriteJSON(msg); err != nil {
			return
		}
		cursor = l.ID
	}

	// Poll for new labels
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	// Read pump (detect client disconnect)
	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}()

	for {
		select {
		case <-done:
			return
		case <-ticker.C:
			labels, err := GetLabelsSince(s.db, cursor, 100)
			if err != nil {
				slog.Error("Failed to poll labels", "error", err)
				continue
			}
			for _, l := range labels {
				msg := LabelsMessage{
					Seq:    l.ID,
					Labels: []LabelOutput{labelToOutput(l)},
				}
				if err := conn.WriteJSON(msg); err != nil {
					return
				}
				cursor = l.ID
			}
		}
	}
}

// handleQueryLabels implements com.atproto.label.queryLabels (HTTP GET).
func (s *Server) handleQueryLabels(w http.ResponseWriter, r *http.Request) {
	uriPatterns := r.URL.Query()["uriPatterns"]
	cursorStr := r.URL.Query().Get("cursor")
	limitStr := r.URL.Query().Get("limit")

	var cursor int64
	if cursorStr != "" {
		cursor, _ = strconv.ParseInt(cursorStr, 10, 64)
	}
	limit := 50
	if limitStr != "" {
		if l, err := strconv.Atoi(limitStr); err == nil && l > 0 && l <= 250 {
			limit = l
		}
	}

	labels, err := GetLabelsSince(s.db, cursor, limit)
	if err != nil {
		http.Error(w, "failed to query labels", http.StatusInternalServerError)
		return
	}

	// Filter by URI patterns if provided
	var filtered []LabelOutput
	for _, l := range labels {
		if len(uriPatterns) == 0 || matchesAnyPattern(l.URI, uriPatterns) {
			filtered = append(filtered, labelToOutput(l))
		}
	}

	var nextCursor string
	if len(labels) > 0 {
		nextCursor = strconv.FormatInt(labels[len(labels)-1].ID, 10)
	}

	resp := struct {
		Cursor string        `json:"cursor,omitempty"`
		Labels []LabelOutput `json:"labels"`
	}{
		Cursor: nextCursor,
		Labels: filtered,
	}
	if resp.Labels == nil {
		resp.Labels = []LabelOutput{}
	}

	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(resp)
}

func matchesAnyPattern(uri string, patterns []string) bool {
	for _, p := range patterns {
		// Simple prefix matching (ATProto spec allows glob-like patterns)
		if p == uri || (len(p) > 0 && p[len(p)-1] == '*' && len(uri) >= len(p)-1 && uri[:len(p)-1] == p[:len(p)-1]) {
			return true
		}
	}
	return false
}