package realtime import ( "context" "encoding/json" "log/slog" "net/http" "sync" "github.com/coder/websocket" "messenger/internal/platform/metrics" ) // MembershipProvider resolves conversation membership and read state. type MembershipProvider interface { ConversationIDsForUser(ctx context.Context, userID int64) ([]int64, error) IsMember(ctx context.Context, conversationID, userID int64) (bool, error) UpdateLastRead(ctx context.Context, conversationID, userID int64) error } // Hub tracks WebSocket connections and fans out events to the // users subscribed to a conversation. type Hub struct { mu sync.Mutex conns map[int64]map[*client]struct{} subs map[int64]map[int64]struct{} members MembershipProvider } // NewHub creates a hub that resolves membership through members. func NewHub(members MembershipProvider) *Hub { return &Hub{ conns: make(map[int64]map[*client]struct{}), subs: make(map[int64]map[int64]struct{}), members: members, } } // HandleWS upgrades the HTTP request to a WebSocket connection and // runs the client loops until the connection closes. func (h *Hub) HandleWS(w http.ResponseWriter, r *http.Request, userID int64) { conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{OriginPatterns: []string{"*"}}) if err != nil { slog.Error("ws upgrade failed", "error", err) return } conn.SetReadLimit(32 << 10) ctx := r.Context() conversationIDs, err := h.members.ConversationIDsForUser(ctx, userID) if err != nil { slog.Error("ws: load memberships", "user_id", userID, "error", err) conn.Close(websocket.StatusInternalError, "membership lookup failed") return } c := newClient(h, conn, userID) h.register(c, conversationIDs) defer h.unregister(c, conversationIDs) h.broadcastPresence(userID, conversationIDs, StatusOnline) defer h.broadcastPresence(userID, conversationIDs, StatusOffline) go c.writeLoop() if err := c.readLoop(); err != nil { slog.Debug("ws connection closed", "user_id", userID, "error", err) } } // PublishMessageCreated delivers a stored message to everyone subscribed // to the conversation, including the sender. func (h *Hub) PublishMessageCreated(conversationID int64, evt MessageCreated) { h.sendToConversation(conversationID, envelope{Type: EventMessageCreated, Data: mustMarshal(evt)}, 0) } // PublishTyping delivers a typing state change to everyone except the sender. func (h *Hub) PublishTyping(conversationID, userID int64, isTyping bool) { evt := Typing{ConversationID: conversationID, UserID: userID, IsTyping: isTyping} h.sendToConversation(conversationID, envelope{Type: EventTyping, Data: mustMarshal(evt)}, userID) } // PublishRead delivers a read receipt to everyone except the reader. func (h *Hub) PublishRead(conversationID, userID, messageID int64) { evt := Read{ConversationID: conversationID, UserID: userID, MessageID: messageID} h.sendToConversation(conversationID, envelope{Type: EventRead, Data: mustMarshal(evt)}, userID) } // Subscribe adds a user to a conversation's subscriber set so they receive its // events without reconnecting (used when a member is added to a group). func (h *Hub) Subscribe(userID, conversationID int64) { h.mu.Lock() defer h.mu.Unlock() if h.subs[conversationID] == nil { h.subs[conversationID] = make(map[int64]struct{}) } h.subs[conversationID][userID] = struct{}{} } // Unsubscribe removes a user from a conversation's subscriber set. func (h *Hub) Unsubscribe(userID, conversationID int64) { h.mu.Lock() defer h.mu.Unlock() users := h.subs[conversationID] delete(users, userID) if len(users) == 0 { delete(h.subs, conversationID) } } // SendToUser delivers an event to every connection of a single user. func (h *Hub) SendToUser(userID int64, evtType string, data any) { payload, err := json.Marshal(envelope{Type: evtType, Data: mustMarshal(data)}) if err != nil { slog.Error("ws: marshal direct event", "error", err) return } h.mu.Lock() defer h.mu.Unlock() for c := range h.conns[userID] { c.enqueueRaw(payload) } } // PublishMemberChanged tells a conversation's members that someone joined or left. func (h *Hub) PublishMemberChanged(conversationID, userID int64, added bool) { evtType := EventMemberAdded if !added { evtType = EventMemberRemoved } evt := MemberChanged{ConversationID: conversationID, UserID: userID} h.sendToConversation(conversationID, envelope{Type: evtType, Data: mustMarshal(evt)}, 0) } // PublishReactionChanged tells a conversation's members that a reaction changed. func (h *Hub) PublishReactionChanged(conversationID, messageID, userID int64, emoji string, added bool) { evtType := EventReactionAdded if !added { evtType = EventReactionRemoved } evt := ReactionChanged{ ConversationID: conversationID, MessageID: messageID, UserID: userID, Emoji: emoji, } h.sendToConversation(conversationID, envelope{Type: evtType, Data: mustMarshal(evt)}, 0) } // IsOnline reports whether a user has at least one active WebSocket connection. func (h *Hub) IsOnline(userID int64) bool { h.mu.Lock() defer h.mu.Unlock() return len(h.conns[userID]) > 0 } func (h *Hub) register(c *client, conversationIDs []int64) { h.mu.Lock() defer h.mu.Unlock() if h.conns[c.userID] == nil { h.conns[c.userID] = make(map[*client]struct{}) } h.conns[c.userID][c] = struct{}{} for _, cid := range conversationIDs { if h.subs[cid] == nil { h.subs[cid] = make(map[int64]struct{}) } h.subs[cid][c.userID] = struct{}{} } metrics.RecordWS(1) } func (h *Hub) unregister(c *client, conversationIDs []int64) { h.mu.Lock() defer h.mu.Unlock() if conns := h.conns[c.userID]; conns != nil { delete(conns, c) if len(conns) == 0 { delete(h.conns, c.userID) } } for _, cid := range conversationIDs { if users := h.subs[cid]; users != nil { delete(users, c.userID) if len(users) == 0 { delete(h.subs, cid) } } } metrics.RecordWS(-1) close(c.done) } func (h *Hub) broadcastPresence(userID int64, conversationIDs []int64, status string) { evt := Presence{UserID: userID, Status: status} env := envelope{Type: EventPresence, Data: mustMarshal(evt)} h.mu.Lock() defer h.mu.Unlock() for _, cid := range conversationIDs { for uid := range h.subs[cid] { if uid == userID { continue } for c := range h.conns[uid] { c.enqueue(env) } } } } func (h *Hub) sendToConversation(conversationID int64, env envelope, exceptUserID int64) { payload, err := json.Marshal(env) if err != nil { slog.Error("ws: marshal event", "error", err) return } h.mu.Lock() defer h.mu.Unlock() for uid := range h.subs[conversationID] { if uid == exceptUserID { continue } for c := range h.conns[uid] { c.enqueueRaw(payload) } } } func mustMarshal(v any) json.RawMessage { data, err := json.Marshal(v) if err != nil { panic(err) } return data }