package events import ( "encoding/json" "log" "net/http" "strings" "sync" "time" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "switchboard-core/metrics" ) const ( writeWait = 10 * time.Second pongWait = 60 * time.Second pingPeriod = (pongWait * 9) / 10 maxMsgSize = 4096 ) // Hub manages WebSocket connections and bridges them to the Bus. type Hub struct { bus *Bus mu sync.RWMutex conns map[string]*Conn // connID → Conn upgrader websocket.Upgrader } // Conn represents a single WebSocket connection. type Conn struct { id string userID string rooms map[string]bool // rooms this connection is subscribed to ws *websocket.Conn send chan []byte hub *Hub unsubs []func() // bus unsubscribe functions } // NewHub creates a Hub bound to an event bus. // allowedOrigins controls the WebSocket upgrader's CheckOrigin behavior. // Empty or "*" allows all origins (dev/test). Comma-separated list // restricts to those origins (production). func NewHub(bus *Bus, allowedOrigins string) *Hub { h := &Hub{ bus: bus, conns: make(map[string]*Conn), } h.upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: h.makeCheckOrigin(allowedOrigins), } return h } // makeCheckOrigin returns a CheckOrigin function that validates the // request Origin header against the configured allowed origins. func (h *Hub) makeCheckOrigin(allowedOrigins string) func(r *http.Request) bool { origins := strings.TrimSpace(allowedOrigins) // No restriction: dev/test or explicit "*" if origins == "" || origins == "*" { return func(r *http.Request) bool { return true } } // Parse comma-separated origin list var allowed []string for _, o := range strings.Split(origins, ",") { o = strings.TrimSpace(o) if o != "" { allowed = append(allowed, o) } } return func(r *http.Request) bool { origin := r.Header.Get("Origin") if origin == "" { // No Origin header — same-origin request (non-browser client). return true } for _, a := range allowed { if a == origin { return true } } log.Printf("[ws] rejected WebSocket from origin %q (allowed: %s)", origin, origins) return false } } // HandleWebSocket is a Gin handler for WebSocket upgrade. // Expects userID set in context (by auth middleware or query param). func (h *Hub) HandleWebSocket(c *gin.Context) { // Auth: check token from query param userID := c.GetString("user_id") if userID == "" { c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication required"}) return } ws, err := h.upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { log.Printf("[ws] upgrade failed: %v", err) return } connID := userID + "-" + time.Now().Format("150405.000") conn := &Conn{ id: connID, userID: userID, rooms: make(map[string]bool), ws: ws, send: make(chan []byte, 64), hub: h, } h.mu.Lock() h.conns[connID] = conn h.mu.Unlock() log.Printf("[ws] connected: %s (user=%s, total=%d)", connID, userID, h.ConnCount()) metrics.WebSocketConnections.Inc() // Subscribe this connection to bus events destined for clients conn.subscribeToBus() // Publish presence h.bus.Publish(Event{ Label: "user.presence", Payload: mustJSON(map[string]any{"user_id": userID, "status": "online"}), Ts: time.Now().UnixMilli(), SenderID: userID, }) // Start read/write pumps go conn.writePump() go conn.readPump() } // ConnCount returns the number of active connections. func (h *Hub) ConnCount() int { h.mu.RLock() defer h.mu.RUnlock() return len(h.conns) } // ConnsByUser returns all connections for a given user. func (h *Hub) ConnsByUser(userID string) []*Conn { h.mu.RLock() defer h.mu.RUnlock() var result []*Conn for _, c := range h.conns { if c.userID == userID { result = append(result, c) } } return result } // IsConnected returns true if the user has at least one active WebSocket. func (h *Hub) IsConnected(userID string) bool { h.mu.RLock() defer h.mu.RUnlock() for _, c := range h.conns { if c.userID == userID { return true } } return false } // PublishToUser sends an event to a specific user via the bus. // All replicas receive the event; only the one with the user's connection delivers it. func (h *Hub) PublishToUser(userID string, event Event) { event.TargetUserID = userID h.bus.Publish(event) } // sendToUserLocal sends an event directly to local WebSocket connections for a user. // Unexported — callers should use PublishToUser for cross-pod delivery. func (h *Hub) sendToUserLocal(userID string, event Event) { data, err := json.Marshal(event) if err != nil { return } h.mu.RLock() defer h.mu.RUnlock() for _, c := range h.conns { if c.userID == userID { select { case c.send <- data: default: } } } } // GetBus returns the underlying event bus. func (h *Hub) GetBus() *Bus { return h.bus } // removeConn cleans up a connection. func (h *Hub) removeConn(conn *Conn) { h.mu.Lock() delete(h.conns, conn.id) h.mu.Unlock() // Unsubscribe from bus for _, unsub := range conn.unsubs { unsub() } log.Printf("[ws] disconnected: %s (total=%d)", conn.id, h.ConnCount()) metrics.WebSocketConnections.Dec() // Publish presence offline (only if no other conns for this user) if len(h.ConnsByUser(conn.userID)) == 0 { h.bus.Publish(Event{ Label: "user.presence", Payload: mustJSON(map[string]any{"user_id": conn.userID, "status": "offline"}), Ts: time.Now().UnixMilli(), SenderID: conn.userID, }) } } // subscribeToBus registers this conn as a bus subscriber for client-facing events. func (c *Conn) subscribeToBus() { unsub := c.hub.bus.Subscribe("*", func(e Event) { // Only forward events destined for clients if !ShouldSendToClient(e.Label) { return } // Targeted events skip room filtering (user-scoped, not room-scoped). if e.TargetUserID != "" && e.TargetUserID != c.userID { return } // must not be forwarded to WebSocket clients — they sent it. if strings.HasPrefix(e.Label, "tool.result.") { return } // Room filtering: if event has a room, only send if conn is in that room if e.Room != "" && !c.rooms[e.Room] { return } data, err := json.Marshal(e) if err != nil { return } select { case c.send <- data: default: // send buffer full, drop event } }) c.unsubs = append(c.unsubs, unsub) } // JoinRoom adds this connection to a named room. func (c *Conn) JoinRoom(room string) { c.rooms[room] = true } // LeaveRoom removes this connection from a named room. func (c *Conn) LeaveRoom(room string) { delete(c.rooms, room) } // ── Read/Write Pumps ───────────────────────── func (c *Conn) readPump() { defer func() { c.hub.removeConn(c) c.ws.Close() }() c.ws.SetReadLimit(maxMsgSize) c.ws.SetReadDeadline(time.Now().Add(pongWait)) c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)) return nil }) for { _, message, err := c.ws.ReadMessage() if err != nil { break } var event Event if err := json.Unmarshal(message, &event); err != nil { continue } // Handle ping if event.Label == "ping" { pong, _ := json.Marshal(Event{Label: "pong", Ts: time.Now().UnixMilli()}) select { case c.send <- pong: default: } continue } if event.Label == "room.subscribe" { var req struct { Room string `json:"room"` } if json.Unmarshal(event.Payload, &req) == nil && req.Room != "" && len(c.rooms) < 100 { c.JoinRoom(req.Room) } continue } if event.Label == "room.unsubscribe" { var req struct { Room string `json:"room"` } if json.Unmarshal(event.Payload, &req) == nil && req.Room != "" { c.LeaveRoom(req.Room) } continue } // Validate: only accept events allowed from clients if !ShouldAcceptFromClient(event.Label) { continue } // Tag with sender info event.SenderID = c.userID event.ConnID = c.id if event.Ts == 0 { event.Ts = time.Now().UnixMilli() } // Publish into bus c.hub.bus.Publish(event) } } func (c *Conn) writePump() { ticker := time.NewTicker(pingPeriod) defer func() { ticker.Stop() c.ws.Close() }() for { select { case msg, ok := <-c.send: c.ws.SetWriteDeadline(time.Now().Add(writeWait)) if !ok { c.ws.WriteMessage(websocket.CloseMessage, []byte{}) return } if err := c.ws.WriteMessage(websocket.TextMessage, msg); err != nil { return } case <-ticker.C: c.ws.SetWriteDeadline(time.Now().Add(writeWait)) if err := c.ws.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } } // ── Helpers ────────────────────────────────── func mustJSON(v any) json.RawMessage { data, err := json.Marshal(v) if err != nil { return json.RawMessage(`{}`) } return data }