package gateway import ( "context" "database/sql" "encoding/json" "log/slog" "sync" "time" ) const ( // idleTimeout is the duration of inactivity before a user is set to idle. idleTimeout = 5 * time.Minute // idleCheckInterval is how often the hub checks for idle users. idleCheckInterval = 30 * time.Second ) // PresenceData is the data payload for PRESENCE_UPDATE events. type PresenceData struct { UserID string `json:"user_id"` Username string `json:"username"` Status string `json:"status"` } // Hub maintains the set of active clients and broadcasts messages to them. type Hub struct { mu sync.RWMutex clients map[*Client]struct{} presence map[string]string // userID -> status lastActivity map[string]time.Time // userID -> last activity time userServers map[string]map[string]struct{} // userID -> set of serverIDs register chan *Client unregister chan *Client broadcast chan []byte db *sql.DB logger *slog.Logger } // NewHub creates a new Hub. func NewHub(db *sql.DB, logger *slog.Logger) *Hub { return &Hub{ clients: make(map[*Client]struct{}), presence: make(map[string]string), lastActivity: make(map[string]time.Time), userServers: make(map[string]map[string]struct{}), register: make(chan *Client), unregister: make(chan *Client), broadcast: make(chan []byte, 256), db: db, logger: logger, } } // Run starts the hub's main event loop. Call this in a goroutine. func (h *Hub) Run() { idleTicker := time.NewTicker(idleCheckInterval) defer idleTicker.Stop() for { select { case client := <-h.register: h.mu.Lock() h.clients[client] = struct{}{} h.presence[client.UserID] = "online" h.lastActivity[client.UserID] = time.Now() h.mu.Unlock() h.RefreshUserServers(client.UserID) h.setUserStatus(client.UserID, "online") h.broadcastPresence(client.UserID, "online") h.logger.Info("client connected", "user_id", client.UserID) case client := <-h.unregister: h.mu.Lock() if _, ok := h.clients[client]; ok { delete(h.clients, client) close(client.send) } // Check if user has any remaining connected clients. hasOther := h.hasClientForUser(client.UserID) if !hasOther { delete(h.presence, client.UserID) delete(h.lastActivity, client.UserID) delete(h.userServers, client.UserID) } h.mu.Unlock() if !hasOther { h.setUserStatus(client.UserID, "offline") h.broadcastPresence(client.UserID, "offline") } h.logger.Info("client disconnected", "user_id", client.UserID) case message := <-h.broadcast: h.mu.RLock() for client := range h.clients { select { case client.send <- message: default: go func(c *Client) { h.unregister <- c }(client) } } h.mu.RUnlock() case <-idleTicker.C: h.checkIdleUsers() } } } // hasClientForUser returns true if any connected client belongs to the given // user. Caller must hold h.mu (at least for reading). func (h *Hub) hasClientForUser(userID string) bool { for c := range h.clients { if c.UserID == userID { return true } } return false } // RecordActivity updates the last-activity timestamp for a user. If the user // was idle, it resets their status back to online and broadcasts the change. // Safe to call from any goroutine. func (h *Hub) RecordActivity(userID string) { h.mu.Lock() h.lastActivity[userID] = time.Now() wasIdle := h.presence[userID] == "idle" if wasIdle { h.presence[userID] = "online" } h.mu.Unlock() if wasIdle { h.setUserStatus(userID, "online") h.broadcastPresence(userID, "online") } } // checkIdleUsers scans all tracked users and marks as idle any who have been // inactive for longer than idleTimeout. func (h *Hub) checkIdleUsers() { h.mu.Lock() now := time.Now() var idleUsers []string for userID, last := range h.lastActivity { if h.presence[userID] == "online" && now.Sub(last) >= idleTimeout { h.presence[userID] = "idle" idleUsers = append(idleUsers, userID) } } h.mu.Unlock() for _, userID := range idleUsers { h.setUserStatus(userID, "idle") h.broadcastPresence(userID, "idle") } } // setUserStatus persists the user's status to the database. func (h *Hub) setUserStatus(userID, status string) { setUserStatus(h.db, h.logger, userID, status) } // broadcastPresence sends a PRESENCE_UPDATE event to all connected clients. func (h *Hub) broadcastPresence(userID, status string) { h.BroadcastEvent(Event{ Type: EventPresenceUpdate, Data: PresenceData{ UserID: userID, Status: status, }, }) } // SetUserStatus is the exported form of setUserStatus, used by non-gateway // handlers that need to persist a status change. func (h *Hub) SetUserStatus(userID, status string) { h.setUserStatus(userID, status) } // BroadcastPresence sends a PRESENCE_UPDATE event including the username. func (h *Hub) BroadcastPresence(userID, username, status string) { h.mu.Lock() h.presence[userID] = status h.lastActivity[userID] = time.Now() h.mu.Unlock() h.BroadcastEvent(Event{ Type: EventPresenceUpdate, Data: PresenceData{ UserID: userID, Username: username, Status: status, }, }) } // Register queues a client for registration with the hub. func (h *Hub) Register(client *Client) { h.register <- client } // Unregister queues a client for removal from the hub. func (h *Hub) Unregister(client *Client) { h.unregister <- client } // BroadcastEvent serializes an Event and broadcasts it to all connected clients. func (h *Hub) BroadcastEvent(event Event) { data, err := json.Marshal(event) if err != nil { h.logger.Error("failed to marshal event", "type", event.Type, "error", err) return } h.broadcast <- data } // ClientCount returns the number of connected clients. func (h *Hub) ClientCount() int { h.mu.RLock() defer h.mu.RUnlock() return len(h.clients) } // SendToUser sends an event only to the specified user's connected client(s). func (h *Hub) SendToUser(userID string, event Event) { data, err := json.Marshal(event) if err != nil { h.logger.Error("failed to marshal event", "type", event.Type, "error", err) return } h.mu.RLock() defer h.mu.RUnlock() for client := range h.clients { if client.UserID == userID { select { case client.send <- data: default: go func(c *Client) { h.unregister <- c }(client) } } } } // RefreshUserServers loads the server membership for a user from the database // and caches it. Call when a client connects or when server membership changes. func (h *Hub) RefreshUserServers(userID string) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() rows, err := h.db.QueryContext(ctx, `SELECT server_id FROM members WHERE user_id = $1`, userID) if err != nil { h.logger.Error("failed to load user servers", "user_id", userID, "error", err) return } defer rows.Close() servers := make(map[string]struct{}) for rows.Next() { var sid string if err := rows.Scan(&sid); err == nil { servers[sid] = struct{}{} } } h.mu.Lock() h.userServers[userID] = servers h.mu.Unlock() } // BroadcastToServer sends an event only to clients who are members of the // given server. Use this for server-scoped events (messages, channels, // reactions) to prevent cross-server data leaks. func (h *Hub) BroadcastToServer(serverID string, event Event) { data, err := json.Marshal(event) if err != nil { h.logger.Error("failed to marshal event", "type", event.Type, "error", err) return } h.mu.RLock() defer h.mu.RUnlock() for client := range h.clients { servers, ok := h.userServers[client.UserID] if !ok { continue } if _, member := servers[serverID]; member { select { case client.send <- data: default: go func(c *Client) { h.unregister <- c }(client) } } } } // BroadcastToConversation sends an event to all clients who are members of a // given direct-message conversation. func (h *Hub) BroadcastToConversation(convID string, event Event) { data, err := json.Marshal(event) if err != nil { h.logger.Error("failed to marshal event", "type", event.Type, "error", err) return } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() rows, err := h.db.QueryContext(ctx, `SELECT user_id FROM conversation_members WHERE conversation_id = $1`, convID) if err != nil { h.logger.Error("failed to load conversation members", "conversation_id", convID, "error", err) return } defer rows.Close() members := make(map[string]struct{}) for rows.Next() { var uid string if err := rows.Scan(&uid); err == nil { members[uid] = struct{}{} } } h.mu.RLock() defer h.mu.RUnlock() for client := range h.clients { if _, ok := members[client.UserID]; ok { select { case client.send <- data: default: go func(c *Client) { h.unregister <- c }(client) } } } } // ServerIDForChannel looks up the server_id that owns the given channel. // Used by handlers to route events to the correct server scope. func (h *Hub) ServerIDForChannel(ctx context.Context, channelID string) (string, error) { var serverID string err := h.db.QueryRowContext(ctx, `SELECT server_id FROM channels WHERE id = $1`, channelID).Scan(&serverID) return serverID, err } func setUserStatus(db *sql.DB, logger *slog.Logger, userID, status string) { if db == nil { return } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _, err := db.ExecContext(ctx, `UPDATE users SET status = $1 WHERE id = $2`, status, userID) if err != nil { logger.Error("failed to update user status", "user_id", userID, "status", status, "error", err) } }