package websocket import ( "encoding/json" "net/http" "sync" "time" "github.com/gorilla/websocket" "github.com/rs/zerolog" ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r *http.Request) bool { return true // Allow all origins in development }, } // Client represents a WebSocket client connection type Client struct { hub *Hub conn *websocket.Conn send chan []byte tenantID string userID string } // Hub maintains the set of active clients and broadcasts messages type Hub struct { clients map[*Client]bool broadcast chan []byte register chan *Client unregister chan *Client logger zerolog.Logger mu sync.RWMutex } // NewHub creates a new WebSocket hub func NewHub(logger zerolog.Logger) *Hub { return &Hub{ clients: make(map[*Client]bool), broadcast: make(chan []byte), register: make(chan *Client), unregister: make(chan *Client), logger: logger.With().Str("component", "websocket").Logger(), } } // Run starts the hub's event loop func (h *Hub) Run() { for { select { case client := <-h.register: h.mu.Lock() h.clients[client] = true h.mu.Unlock() h.logger.Info().Str("tenant", client.tenantID).Msg("client connected") case client := <-h.unregister: h.mu.Lock() if _, ok := h.clients[client]; ok { delete(h.clients, client) close(client.send) } h.mu.Unlock() h.logger.Info().Str("tenant", client.tenantID).Msg("client disconnected") case message := <-h.broadcast: h.mu.RLock() for client := range h.clients { select { case client.send <- message: default: // Client's send channel is full, close it close(client.send) delete(h.clients, client) } } h.mu.RUnlock() } } } // HandleWebSocket upgrades HTTP connection to WebSocket func (h *Hub) HandleWebSocket(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { h.logger.Error().Err(err).Msg("websocket upgrade failed") return } // Extract tenant and user from query params (in production, verify JWT) tenantID := r.URL.Query().Get("tenant_id") userID := r.URL.Query().Get("user_id") if tenantID == "" { tenantID = "default" } client := &Client{ hub: h, conn: conn, send: make(chan []byte, 256), tenantID: tenantID, userID: userID, } client.hub.register <- client // Start goroutines for reading and writing go client.writePump() go client.readPump() } // Broadcast sends a message to all connected clients func (h *Hub) Broadcast(message []byte) { h.broadcast <- message } // BroadcastToTenant sends a message to clients of a specific tenant func (h *Hub) BroadcastToTenant(tenantID string, message []byte) { h.mu.RLock() defer h.mu.RUnlock() for client := range h.clients { if client.tenantID == tenantID { select { case client.send <- message: default: // Channel full, skip } } } } // readPump handles incoming messages from the client func (c *Client) readPump() { defer func() { c.hub.unregister <- c c.conn.Close() }() c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) return nil }) for { _, message, err := c.conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { c.hub.logger.Error().Err(err).Msg("websocket read error") } break } // Handle incoming messages (e.g., subscribe to events) c.hub.logger.Debug().Str("message", string(message)).Msg("received websocket message") } } // writePump handles outgoing messages to the client func (c *Client) writePump() { ticker := time.NewTicker(54 * time.Second) defer func() { ticker.Stop() c.conn.Close() }() for { select { case message, ok := <-c.send: c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if !ok { c.conn.WriteMessage(websocket.CloseMessage, []byte{}) return } c.conn.WriteMessage(websocket.TextMessage, message) case <-ticker.C: c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } } // Event types for WebSocket messages type EventType string const ( EventDealUpdated EventType = "deal_updated" EventInvoiceCreated EventType = "invoice_created" EventTicketUpdated EventType = "ticket_updated" EventEmployeeUpdated EventType = "employee_updated" EventContractReminder EventType = "contract_reminder" EventReportReady EventType = "report_ready" EventWorkflowRun EventType = "workflow_run" ) // Event represents a real-time event type Event struct { Type EventType `json:"type"` TenantID string `json:"tenant_id"` EntityID string `json:"entity_id"` Data interface{} `json:"data"` Timestamp time.Time `json:"timestamp"` } // SendEvent sends an event to relevant clients func (h *Hub) SendEvent(event Event) { // In production, filter by tenant and user permissions message, _ := json.Marshal(event) h.BroadcastToTenant(event.TenantID, message) }