diff --git a/main.go b/main.go index 5bfd18b..e6cfc9f 100644 --- a/main.go +++ b/main.go @@ -2,6 +2,7 @@ package main import ( "context" + "encoding/json" "fmt" "log/slog" "net/http" @@ -24,27 +25,51 @@ var DEFAULT_POOL = []string{ // want yours here? contact me } +// Event represents a Jetstream event +type Event struct { + Did string `json:"did"` + TimeUS int64 `json:"time_us"` + Kind string `json:"kind,omitempty"` + Commit *Commit `json:"commit,omitempty"` +} + +// Commit represents a commit event +type Commit struct { + Rev string `json:"rev,omitempty"` + Operation string `json:"operation,omitempty"` + Collection string `json:"collection,omitempty"` + RKey string `json:"rkey,omitempty"` + Record json.RawMessage `json:"record,omitempty"` + CID string `json:"cid,omitempty"` +} + +// Message wraps a Jetstream event with both parsed and raw forms +type Message struct { + Event *Event + Raw []byte +} + // Broadcaster manages subscribers to Jetstream events type Broadcaster struct { - listeners []chan []byte + listeners []chan *Message mu sync.Mutex connected atomic.Bool lastMessageTime atomic.Int64 // Unix timestamp in seconds } // Subscribe returns a new channel that will receive Jetstream events -func (b *Broadcaster) Subscribe() chan []byte { +func (b *Broadcaster) Subscribe() chan *Message { b.mu.Lock() defer b.mu.Unlock() // firehose can be more-than-1k events per second, // prefer to create a large buffer for the subscribers - ch := make(chan []byte, 10000) + ch := make(chan *Message, 10000) b.listeners = append(b.listeners, ch) return ch } -func (b *Broadcaster) Unsubscribe(ch chan []byte) { +func (b *Broadcaster) Unsubscribe(ch chan *Message) { b.mu.Lock() defer b.mu.Unlock() @@ -57,15 +82,27 @@ func (b *Broadcaster) Unsubscribe(ch chan []byte) { } } -func (b *Broadcaster) Broadcast(message []byte) { +func (b *Broadcaster) Broadcast(rawMessage []byte) { b.lastMessageTime.Store(time.Now().Unix()) + // Parse the event once + var event Event + if err := json.Unmarshal(rawMessage, &event); err != nil { + slog.Debug("Failed to parse event", slog.Any("error", err)) + // Broadcast anyway with nil event + } + + msg := &Message{ + Event: &event, + Raw: rawMessage, + } + b.mu.Lock() defer b.mu.Unlock() for _, ch := range b.listeners { select { - case ch <- message: + case ch <- msg: // event sent successfully. we don't want to block default: // channel full, skip to avoid blocking @@ -120,6 +157,39 @@ func handleHealth(broadcaster *Broadcaster) http.HandlerFunc { } } +// matchesCollection checks if an event matches any of the wanted collections +func matchesCollection(event *Event, wantedCollections []string) bool { + // Always pass through account and identity events + if event.Kind == "account" || event.Kind == "identity" { + return true + } + + // If no wanted collections specified, pass everything + if len(wantedCollections) == 0 { + return true + } + + // For commit events, check the collection + if event.Commit == nil { + return false + } + + collection := event.Commit.Collection + for _, wanted := range wantedCollections { + // Support wildcard matching like "app.bsky.graph.*" + if strings.HasSuffix(wanted, ".*") { + prefix := strings.TrimSuffix(wanted, ".*") + if strings.HasPrefix(collection, prefix+".") || collection == prefix { + return true + } + } else if collection == wanted { + return true + } + } + + return false +} + // handleSubscribe upgrades HTTP connection to websocket and streams events func handleSubscribe(broadcaster *Broadcaster) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { @@ -130,15 +200,34 @@ func handleSubscribe(broadcaster *Broadcaster) http.HandlerFunc { } defer conn.Close() + // Parse wantedCollections from query params + wantedCollections := r.URL.Query()["wantedCollections"] + if len(wantedCollections) > 100 { + slog.Warn("Client requested too many collections, limiting to 100", slog.Int("requested", len(wantedCollections))) + wantedCollections = wantedCollections[:100] + } + // Subscribe to broadcaster ch := broadcaster.Subscribe() defer broadcaster.Unsubscribe(ch) - slog.Info("Client connected", slog.String("remote", r.RemoteAddr)) + if len(wantedCollections) > 0 { + slog.Info("Client connected", slog.String("remote", r.RemoteAddr), slog.Any("wantedCollections", wantedCollections)) + } else { + slog.Info("Client connected", slog.String("remote", r.RemoteAddr)) + } // Stream events to client - for message := range ch { - err := conn.WriteMessage(websocket.TextMessage, message) + for msg := range ch { + // If filtering is enabled, check the event + if len(wantedCollections) > 0 && msg.Event != nil { + // Check if event matches wanted collections + if !matchesCollection(msg.Event, wantedCollections) { + continue + } + } + + err := conn.WriteMessage(websocket.TextMessage, msg.Raw) if err != nil { slog.Debug("Client disconnected", slog.String("remote", r.RemoteAddr), slog.Any("error", err)) break