add support for wantedCollections

main.go

···
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"log/slog"
 	"net/http"
···
 	// 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()
 
···
 	}
 }
 
-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
···
 	}
 }
 
+// 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) {
···
 		}
 		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