at-container-registry/cmd/firehose-tap/main.go

// firehose-tap subscribes to an ATProto com.atproto.sync.subscribeRepos endpoint
// and pretty-prints every event it sees, with extra emphasis on Sync 1.1 fields
// (per-op prev CIDs, commit-level prevData) so you can verify a hold or relay
// is emitting compliant events.
//
// Typical use:
//
//	# Watch new events as you trigger pushes/deletes on production
//	go run ./cmd/firehose-tap wss://hold01.atcr.io
//
//	# Replay from the beginning of the persisted buffer
//	go run ./cmd/firehose-tap --cursor 0 wss://hold01.atcr.io
//
//	# CI mode: subscribe for 30s, exit non-zero if any update/delete op is
//	# missing prev (Sync 1.1 inductive-firehose requirement).
//	go run ./cmd/firehose-tap --validate --duration 30s wss://hold01.atcr.io
//
// The endpoint may be a single hold or a relay. The wss:// scheme is required.
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

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

// ANSI color codes (disabled via --no-color or NO_COLOR env)
var (
	cRed    = "\033[31m"
	cGreen  = "\033[32m"
	cYellow = "\033[33m"
	cCyan   = "\033[36m"
	cBold   = "\033[1m"
	cDim    = "\033[2m"
	cReset  = "\033[0m"
)

func disableColors() {
	cRed, cGreen, cYellow, cCyan, cBold, cDim, cReset = "", "", "", "", "", "", ""
}

type config struct {
	endpoint  string
	cursor    int64
	cursorSet bool
	didFilter string
	validate  bool
	duration  time.Duration
	maxEvents int64
	verbose   bool
	noColor   bool
}

// violationCounts tracks Sync 1.1 compliance issues seen so far.
type violationCounts struct {
	missingPrev     atomic.Int64 // update/delete op without prev
	createWithPrev  atomic.Int64 // create op with non-nil prev (spec forbids)
	deleteWithCid   atomic.Int64 // delete op with non-nil cid (spec forbids)
	missingPrevData atomic.Int64 // commit with since but no prevData
}

func (v *violationCounts) total() int64 {
	return v.missingPrev.Load() + v.createWithPrev.Load() + v.deleteWithCid.Load() + v.missingPrevData.Load()
}

func main() {
	cfg := parseFlags()
	if cfg.noColor || os.Getenv("NO_COLOR") != "" {
		disableColors()
	}

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()
	if cfg.duration > 0 {
		var timeoutCancel context.CancelFunc
		ctx, timeoutCancel = context.WithTimeout(ctx, cfg.duration)
		defer timeoutCancel()
	}

	if err := run(ctx, cfg); err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) {
		fmt.Fprintf(os.Stderr, "%sfirehose-tap: %v%s\n", cRed, err, cReset)
		os.Exit(2)
	}
}

func parseFlags() *config {
	cfg := &config{}
	flag.Int64Var(&cfg.cursor, "cursor", -1, "subscribeRepos cursor (-1 = new events only, 0 = replay from start)")
	flag.StringVar(&cfg.didFilter, "did", "", "only show events for this DID (case-sensitive match on commit.repo)")
	flag.BoolVar(&cfg.validate, "validate", false, "exit non-zero if any Sync 1.1 violation is observed")
	flag.DurationVar(&cfg.duration, "duration", 0, "exit cleanly after this duration (e.g. 30s)")
	flag.Int64Var(&cfg.maxEvents, "max-events", 0, "exit cleanly after seeing this many events (0 = unlimited)")
	flag.BoolVar(&cfg.verbose, "verbose", false, "print full ops for create/no-prev-needed ops too")
	flag.BoolVar(&cfg.noColor, "no-color", false, "disable ANSI colors")

	// Detect whether --cursor was explicitly set so we can omit it from the URL
	// when the user didn't pass one (default -1 means "live, no backfill").
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] <subscribeRepos-endpoint>\n\n", os.Args[0])
		fmt.Fprintln(os.Stderr, "Endpoint examples:")
		fmt.Fprintln(os.Stderr, "  wss://hold01.atcr.io")
		fmt.Fprintln(os.Stderr, "  wss://relay1.us-east.bsky.network")
		fmt.Fprintln(os.Stderr, "  ws://localhost:8080")
		fmt.Fprintln(os.Stderr)
		fmt.Fprintln(os.Stderr, "Flags:")
		flag.PrintDefaults()
	}
	flag.Parse()

	flag.Visit(func(f *flag.Flag) {
		if f.Name == "cursor" {
			cfg.cursorSet = true
		}
	})

	args := flag.Args()
	if len(args) != 1 {
		flag.Usage()
		os.Exit(2)
	}
	cfg.endpoint = args[0]
	return cfg
}

func run(ctx context.Context, cfg *config) error {
	wsURL, err := buildWSURL(cfg.endpoint, cfg.cursor, cfg.cursorSet)
	if err != nil {
		return err
	}

	fmt.Fprintf(os.Stderr, "%sfirehose-tap%s connecting to %s%s%s\n", cBold, cReset, cCyan, wsURL, cReset)
	if cfg.validate {
		fmt.Fprintf(os.Stderr, "%svalidate mode: process will exit non-zero on Sync 1.1 violations%s\n", cDim, cReset)
	}

	header := http.Header{}
	header.Set("User-Agent", "atcr-firehose-tap/1.0")

	dialer := websocket.DefaultDialer
	dialer.HandshakeTimeout = 30 * time.Second
	conn, resp, err := dialer.DialContext(ctx, wsURL, header)
	if err != nil {
		if resp != nil {
			return fmt.Errorf("dial %s: %w (HTTP %d)", wsURL, err, resp.StatusCode)
		}
		return fmt.Errorf("dial %s: %w", wsURL, err)
	}
	defer conn.Close()

	v := &violationCounts{}
	var eventsSeen atomic.Int64

	callbacks := &events.RepoStreamCallbacks{
		RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
			n := eventsSeen.Add(1)
			if cfg.didFilter == "" || cfg.didFilter == evt.Repo {
				printCommit(evt, v, cfg)
			}
			if cfg.maxEvents > 0 && n >= cfg.maxEvents {
				return errMaxEvents
			}
			return nil
		},
		RepoIdentity: func(evt *comatproto.SyncSubscribeRepos_Identity) error {
			eventsSeen.Add(1)
			if cfg.didFilter == "" || cfg.didFilter == evt.Did {
				printIdentity(evt)
			}
			return nil
		},
		RepoAccount: func(evt *comatproto.SyncSubscribeRepos_Account) error {
			eventsSeen.Add(1)
			if cfg.didFilter == "" || cfg.didFilter == evt.Did {
				printAccount(evt)
			}
			return nil
		},
		RepoSync: func(evt *comatproto.SyncSubscribeRepos_Sync) error {
			eventsSeen.Add(1)
			if cfg.didFilter == "" || cfg.didFilter == evt.Did {
				printSync(evt)
			}
			return nil
		},
		RepoInfo: func(evt *comatproto.SyncSubscribeRepos_Info) error {
			fmt.Printf("%s#info%s name=%s message=%s\n", cYellow, cReset, evt.Name, derefStr(evt.Message))
			return nil
		},
		Error: func(evt *events.ErrorFrame) error {
			fmt.Printf("%s#error%s %s: %s\n", cRed, cReset, evt.Error, evt.Message)
			return nil
		},
	}

	sched := sequential.NewScheduler("firehose-tap", callbacks.EventHandler)
	streamErr := events.HandleRepoStream(ctx, conn, sched, nil)
	if errors.Is(streamErr, errMaxEvents) {
		streamErr = nil
	}

	printSummary(eventsSeen.Load(), v, cfg)

	if cfg.validate && v.total() > 0 {
		os.Exit(1)
	}
	return streamErr
}

// errMaxEvents is returned from the commit handler to short-circuit
// HandleRepoStream once --max-events has been reached.
var errMaxEvents = errors.New("max events reached")

func buildWSURL(endpoint string, cursor int64, cursorSet bool) (string, error) {
	if !strings.Contains(endpoint, "://") {
		endpoint = "wss://" + endpoint
	}
	u, err := url.Parse(endpoint)
	if err != nil {
		return "", fmt.Errorf("parse endpoint: %w", err)
	}
	switch u.Scheme {
	case "wss", "ws":
		// already a websocket URL
	case "https":
		u.Scheme = "wss"
	case "http":
		u.Scheme = "ws"
	default:
		return "", fmt.Errorf("unsupported scheme %q (use ws:// or wss://)", u.Scheme)
	}

	if !strings.Contains(u.Path, "subscribeRepos") {
		u.Path = strings.TrimRight(u.Path, "/") + "/xrpc/com.atproto.sync.subscribeRepos"
	}

	if cursorSet {
		q := u.Query()
		q.Set("cursor", fmt.Sprintf("%d", cursor))
		u.RawQuery = q.Encode()
	}
	return u.String(), nil
}

func printCommit(evt *comatproto.SyncSubscribeRepos_Commit, v *violationCounts, cfg *config) {
	prevData := "(none)"
	if evt.PrevData != nil {
		prevData = shortCID(evt.PrevData.String())
	}
	since := "(none)"
	if evt.Since != nil {
		since = *evt.Since
	}

	// Spec: commits beyond the very first one should carry prevData.
	if evt.Since != nil && evt.PrevData == nil {
		v.missingPrevData.Add(1)
	}

	fmt.Printf("%s#commit%s seq=%d repo=%s rev=%s since=%s prevData=%s ops=%d\n",
		cGreen, cReset, evt.Seq, evt.Repo, evt.Rev, since, prevData, len(evt.Ops))

	for _, op := range evt.Ops {
		printOp(op, v, cfg)
	}
}

func printOp(op *comatproto.SyncSubscribeRepos_RepoOp, v *violationCounts, cfg *config) {
	action := op.Action

	cidStr := "(nil)"
	if op.Cid != nil {
		cidStr = shortCID(op.Cid.String())
	}
	prevStr := "(nil)"
	if op.Prev != nil {
		prevStr = shortCID(op.Prev.String())
	}

	violated := false
	var noteParts []string

	switch action {
	case "update", "delete":
		if op.Prev == nil {
			v.missingPrev.Add(1)
			violated = true
			noteParts = append(noteParts, "MISSING prev (Sync 1.1)")
		}
	case "create":
		if op.Prev != nil {
			v.createWithPrev.Add(1)
			violated = true
			noteParts = append(noteParts, "create MUST NOT have prev")
		}
	}
	if action == "delete" && op.Cid != nil {
		v.deleteWithCid.Add(1)
		violated = true
		noteParts = append(noteParts, "delete MUST have nil cid")
	}

	// In non-verbose mode, only print ops that updated/deleted (the interesting
	// Sync 1.1 cases) or that violated a rule.
	if !cfg.verbose && action == "create" && !violated {
		return
	}

	actionColor := cCyan
	if violated {
		actionColor = cRed
	}
	note := ""
	if len(noteParts) > 0 {
		note = fmt.Sprintf(" %s[%s]%s", cRed, strings.Join(noteParts, "; "), cReset)
	}
	fmt.Printf("    %s%-6s%s %s cid=%s prev=%s%s\n",
		actionColor, action, cReset, op.Path, cidStr, prevStr, note)
}

func printIdentity(evt *comatproto.SyncSubscribeRepos_Identity) {
	handle := "(nil)"
	if evt.Handle != nil {
		handle = *evt.Handle
	}
	fmt.Printf("%s#identity%s seq=%d did=%s handle=%s\n",
		cYellow, cReset, evt.Seq, evt.Did, handle)
}

func printAccount(evt *comatproto.SyncSubscribeRepos_Account) {
	status := "(nil)"
	if evt.Status != nil {
		status = *evt.Status
	}
	fmt.Printf("%s#account%s seq=%d did=%s active=%v status=%s\n",
		cYellow, cReset, evt.Seq, evt.Did, evt.Active, status)
}

func printSync(evt *comatproto.SyncSubscribeRepos_Sync) {
	fmt.Printf("%s#sync%s seq=%d did=%s rev=%s\n",
		cYellow, cReset, evt.Seq, evt.Did, evt.Rev)
}

func printSummary(seen int64, v *violationCounts, cfg *config) {
	fmt.Fprintln(os.Stderr)
	fmt.Fprintf(os.Stderr, "%s--- summary ---%s\n", cBold, cReset)
	fmt.Fprintf(os.Stderr, "events seen: %d\n", seen)
	if v.total() == 0 {
		fmt.Fprintf(os.Stderr, "%sno Sync 1.1 violations%s\n", cGreen, cReset)
		return
	}
	fmt.Fprintf(os.Stderr, "%sviolations:%s\n", cRed, cReset)
	if n := v.missingPrev.Load(); n > 0 {
		fmt.Fprintf(os.Stderr, "  %d update/delete op(s) missing prev\n", n)
	}
	if n := v.createWithPrev.Load(); n > 0 {
		fmt.Fprintf(os.Stderr, "  %d create op(s) with non-nil prev\n", n)
	}
	if n := v.deleteWithCid.Load(); n > 0 {
		fmt.Fprintf(os.Stderr, "  %d delete op(s) with non-nil cid\n", n)
	}
	if n := v.missingPrevData.Load(); n > 0 {
		fmt.Fprintf(os.Stderr, "  %d commit(s) with since but missing prevData\n", n)
	}
	if !cfg.validate {
		fmt.Fprintf(os.Stderr, "%s(--validate not set, exiting 0 anyway)%s\n", cDim, cReset)
	}
}

func shortCID(s string) string {
	if len(s) <= 16 {
		return s
	}
	return s[:8] + "…" + s[len(s)-6:]
}

func derefStr(s *string) string {
	if s == nil {
		return ""
	}
	return *s
}