package ws import ( "encoding/base64" "encoding/json" "fmt" "log" "net/http" "strings" "sync" "time" "github.com/gorilla/websocket" "github.com/sandrodach/mindsift/internal/domain" ) var upgrader = websocket.Upgrader{ ReadBufferSize: 4096, WriteBufferSize: 4096, CheckOrigin: func(r *http.Request) bool { return true // Self-hosted reverse proxy compatible }, } type Client struct { hub *Hub conn *websocket.Conn roomID string userID string send chan []byte subscribedRooms map[string]bool } type SSEClient struct { id string userID string roomID string send chan []byte subscribedRooms map[string]bool } type docLockItem struct { docID string userID string userName string lockedAt time.Time expiresAt time.Time } type CRDTStore interface { GetDocumentCRDTState(docID string) ([]byte, [][]byte, error) AppendDocumentCRDTUpdate(docID string, update []byte, userID string) error SaveDocumentCRDTState(docID string, stateBlob []byte, userID string) error } type Hub struct { mu sync.RWMutex rooms map[string]map[*Client]bool userRooms map[string]map[*Client]bool allClients map[*Client]bool sseClients map[*SSEClient]bool sseUserRooms map[string]map[*SSEClient]bool sseRooms map[string]map[*SSEClient]bool locks map[string]*docLockItem crdtStore CRDTStore } func NewHub() *Hub { h := &Hub{ rooms: make(map[string]map[*Client]bool), userRooms: make(map[string]map[*Client]bool), allClients: make(map[*Client]bool), sseClients: make(map[*SSEClient]bool), sseUserRooms: make(map[string]map[*SSEClient]bool), sseRooms: make(map[string]map[*SSEClient]bool), locks: make(map[string]*docLockItem), } go h.cleanupLocksLoop() return h } func (h *Hub) SetCRDTStore(store CRDTStore) { h.mu.Lock() defer h.mu.Unlock() h.crdtStore = store } func (h *Hub) Register(c *Client) { h.mu.Lock() h.allClients[c] = true if c.roomID != "" { if h.rooms[c.roomID] == nil { h.rooms[c.roomID] = make(map[*Client]bool) } h.rooms[c.roomID][c] = true } if c.userID != "" && c.userID != "anonymous" { if h.userRooms[c.userID] == nil { h.userRooms[c.userID] = make(map[*Client]bool) } h.userRooms[c.userID][c] = true } onlineCount := len(h.allClients) + len(h.sseClients) h.mu.Unlock() if c.roomID != "" { h.Broadcast(c.roomID, "", "USER_PRESENCE", map[string]interface{}{ "onlineCount": onlineCount, }) } } func (h *Hub) Unregister(c *Client) { h.mu.Lock() delete(h.allClients, c) if c.roomID != "" { if room, ok := h.rooms[c.roomID]; ok { delete(room, c) if len(room) == 0 { delete(h.rooms, c.roomID) } } } var docRoomsLeft []string if c.subscribedRooms != nil { for rID := range c.subscribedRooms { if room, ok := h.rooms[rID]; ok { delete(room, c) if len(room) == 0 { delete(h.rooms, rID) } } if strings.HasPrefix(rID, "doc:") { docRoomsLeft = append(docRoomsLeft, strings.TrimPrefix(rID, "doc:")) } } } userRemaining := 0 if c.userID != "" && c.userID != "anonymous" { if uRoom, ok := h.userRooms[c.userID]; ok { delete(uRoom, c) userRemaining += len(uRoom) if len(uRoom) == 0 { delete(h.userRooms, c.userID) } } if sseRoom, ok := h.sseUserRooms[c.userID]; ok { userRemaining += len(sseRoom) } } close(c.send) // If this user has no other active client connections, release their document locks var releasedDocIDs []string if userRemaining == 0 && c.userID != "" { now := time.Now() for docID, l := range h.locks { if l.userID == c.userID || l.expiresAt.Before(now) { releasedDocIDs = append(releasedDocIDs, docID) delete(h.locks, docID) } } } h.mu.Unlock() // Broadcast unlock events for released locks for _, docID := range releasedDocIDs { h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID, "doc_id": docID}) } // Broadcast awareness leave for doc rooms for _, docID := range docRoomsLeft { h.BroadcastRoom("doc:"+docID, "CRDT_AWARENESS_LEAVE", map[string]interface{}{ "docId": docID, "userId": c.userID, }, nil) } } func (h *Hub) RegisterSSE(c *SSEClient) { h.mu.Lock() h.sseClients[c] = true if c.roomID != "" { if h.sseRooms[c.roomID] == nil { h.sseRooms[c.roomID] = make(map[*SSEClient]bool) } h.sseRooms[c.roomID][c] = true } if c.userID != "" && c.userID != "anonymous" { if h.sseUserRooms[c.userID] == nil { h.sseUserRooms[c.userID] = make(map[*SSEClient]bool) } h.sseUserRooms[c.userID][c] = true } if c.subscribedRooms != nil { for rID := range c.subscribedRooms { if h.sseRooms[rID] == nil { h.sseRooms[rID] = make(map[*SSEClient]bool) } h.sseRooms[rID][c] = true } } onlineCount := len(h.allClients) + len(h.sseClients) h.mu.Unlock() if c.roomID != "" { h.Broadcast(c.roomID, "", "USER_PRESENCE", map[string]interface{}{ "onlineCount": onlineCount, }) } } func (h *Hub) UnregisterSSE(c *SSEClient) { h.mu.Lock() delete(h.sseClients, c) if c.roomID != "" { if room, ok := h.sseRooms[c.roomID]; ok { delete(room, c) if len(room) == 0 { delete(h.sseRooms, c.roomID) } } } var docRoomsLeft []string if c.subscribedRooms != nil { for rID := range c.subscribedRooms { if room, ok := h.sseRooms[rID]; ok { delete(room, c) if len(room) == 0 { delete(h.sseRooms, rID) } } if strings.HasPrefix(rID, "doc:") { docRoomsLeft = append(docRoomsLeft, strings.TrimPrefix(rID, "doc:")) } } } userRemaining := 0 if c.userID != "" && c.userID != "anonymous" { if uRoom, ok := h.sseUserRooms[c.userID]; ok { delete(uRoom, c) userRemaining += len(uRoom) if len(uRoom) == 0 { delete(h.sseUserRooms, c.userID) } } if wsRoom, ok := h.userRooms[c.userID]; ok { userRemaining += len(wsRoom) } } close(c.send) var releasedDocIDs []string if userRemaining == 0 && c.userID != "" { now := time.Now() for docID, l := range h.locks { if l.userID == c.userID || l.expiresAt.Before(now) { releasedDocIDs = append(releasedDocIDs, docID) delete(h.locks, docID) } } } h.mu.Unlock() for _, docID := range releasedDocIDs { h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID, "doc_id": docID}) } for _, docID := range docRoomsLeft { h.BroadcastRoom("doc:"+docID, "CRDT_AWARENESS_LEAVE", map[string]interface{}{ "docId": docID, "userId": c.userID, }, nil) } } func (h *Hub) JoinSSERoomByID(sseID, roomID string) bool { h.mu.Lock() defer h.mu.Unlock() for client := range h.sseClients { if client.id == sseID { if h.sseRooms[roomID] == nil { h.sseRooms[roomID] = make(map[*SSEClient]bool) } h.sseRooms[roomID][client] = true if client.subscribedRooms == nil { client.subscribedRooms = make(map[string]bool) } client.subscribedRooms[roomID] = true return true } } return false } func (h *Hub) LeaveSSERoomByID(sseID, roomID string) bool { h.mu.Lock() defer h.mu.Unlock() for client := range h.sseClients { if client.id == sseID { if room, ok := h.sseRooms[roomID]; ok { delete(room, client) if len(room) == 0 { delete(h.sseRooms, roomID) } } if client.subscribedRooms != nil { delete(client.subscribedRooms, roomID) } return true } } return false } // Document Locking Methods func (h *Hub) AcquireDocLock(docID, userID, userName string, ttl time.Duration) (*domain.DocLockInfo, bool, *domain.DocLockInfo) { h.mu.Lock() defer h.mu.Unlock() now := time.Now() existing, exists := h.locks[docID] if exists && existing.expiresAt.After(now) { if existing.userID != userID { // Locked by another user info := &domain.DocLockInfo{ DocID: existing.docID, UserID: existing.userID, UserName: existing.userName, ExpiresAt: existing.expiresAt.Format(time.RFC3339), } return nil, false, info } // Same user extending lock existing.expiresAt = now.Add(ttl) if userName != "" { existing.userName = userName } info := &domain.DocLockInfo{ DocID: existing.docID, UserID: existing.userID, UserName: existing.userName, ExpiresAt: existing.expiresAt.Format(time.RFC3339), } return info, true, nil } // Create fresh lock item := &docLockItem{ docID: docID, userID: userID, userName: userName, lockedAt: now, expiresAt: now.Add(ttl), } h.locks[docID] = item info := &domain.DocLockInfo{ DocID: item.docID, UserID: item.userID, UserName: item.userName, ExpiresAt: item.expiresAt.Format(time.RFC3339), } return info, true, nil } func (h *Hub) HeartbeatDocLock(docID, userID string, ttl time.Duration) bool { h.mu.Lock() defer h.mu.Unlock() now := time.Now() item, ok := h.locks[docID] if !ok || item.userID != userID || item.expiresAt.Before(now) { return false } item.expiresAt = now.Add(ttl) return true } func (h *Hub) ReleaseDocLock(docID, userID string, isAdmin bool) bool { h.mu.Lock() defer h.mu.Unlock() item, ok := h.locks[docID] if !ok { return false } if item.userID == userID || isAdmin || userID == "admin" { delete(h.locks, docID) return true } return false } func (h *Hub) GetDocLock(docID string) *domain.DocLockInfo { h.mu.RLock() defer h.mu.RUnlock() item, ok := h.locks[docID] if !ok || item.expiresAt.Before(time.Now()) { return nil } return &domain.DocLockInfo{ DocID: item.docID, UserID: item.userID, UserName: item.userName, ExpiresAt: item.expiresAt.Format(time.RFC3339), } } func (h *Hub) cleanupLocksLoop() { ticker := time.NewTicker(2 * time.Second) defer ticker.Stop() for range ticker.C { h.mu.Lock() now := time.Now() var expiredDocIDs []string for docID, l := range h.locks { if l.expiresAt.Before(now) { expiredDocIDs = append(expiredDocIDs, docID) delete(h.locks, docID) } } h.mu.Unlock() for _, docID := range expiredDocIDs { h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID, "doc_id": docID}) } } } // Broadcasting Methods func (h *Hub) BroadcastToAll(event string, payload interface{}) { msg := domain.WSMessage{ Event: event, Payload: payload, Timestamp: domain.NowISO(), } data, err := json.Marshal(msg) if err != nil { return } h.mu.RLock() defer h.mu.RUnlock() for client := range h.allClients { select { case client.send <- data: default: } } for sseClient := range h.sseClients { select { case sseClient.send <- data: default: } } } func (h *Hub) BroadcastToUser(userID, event string, payload interface{}) { msg := domain.WSMessage{ Event: event, Payload: payload, Timestamp: domain.NowISO(), } data, err := json.Marshal(msg) if err != nil { return } h.mu.RLock() defer h.mu.RUnlock() if room, ok := h.userRooms[userID]; ok { for client := range room { select { case client.send <- data: default: } } } if sseRoom, ok := h.sseUserRooms[userID]; ok { for sseClient := range sseRoom { select { case sseClient.send <- data: default: } } } } func (h *Hub) BroadcastTaskEvent(userID, event string, payload interface{}) { h.BroadcastToUser(userID, event, payload) } func (h *Hub) BroadcastDocEvent(docID, event string, payload interface{}) { h.BroadcastToAll(event, payload) } func (h *Hub) Broadcast(spaceID, listID, event string, payload interface{}) { msg := domain.WSMessage{ Event: event, ListID: listID, Payload: payload, Timestamp: domain.NowISO(), } data, err := json.Marshal(msg) if err != nil { return } h.mu.RLock() defer h.mu.RUnlock() targetRoom := listID if targetRoom == "" { targetRoom = spaceID } if room, ok := h.rooms[targetRoom]; ok { for client := range room { select { case client.send <- data: default: // Slow client buffer full } } } if sseRoom, ok := h.sseRooms[targetRoom]; ok { for sseClient := range sseRoom { select { case sseClient.send <- data: default: } } } } func (h *Hub) BroadcastListEvent(listID, event string, payload interface{}) { h.Broadcast("", listID, event, payload) } func (h *Hub) JoinRoom(c *Client, roomID string) { h.mu.Lock() defer h.mu.Unlock() if h.rooms[roomID] == nil { h.rooms[roomID] = make(map[*Client]bool) } h.rooms[roomID][c] = true if c.subscribedRooms == nil { c.subscribedRooms = make(map[string]bool) } c.subscribedRooms[roomID] = true } func (h *Hub) LeaveRoom(c *Client, roomID string) { h.mu.Lock() defer h.mu.Unlock() if room, ok := h.rooms[roomID]; ok { delete(room, c) if len(room) == 0 { delete(h.rooms, roomID) } } if c.subscribedRooms != nil { delete(c.subscribedRooms, roomID) } } func (h *Hub) BroadcastRoom(roomID, event string, payload interface{}, excludeClient *Client) { msg := domain.WSMessage{ Event: event, Payload: payload, Timestamp: domain.NowISO(), } data, err := json.Marshal(msg) if err != nil { return } h.mu.RLock() defer h.mu.RUnlock() if room, ok := h.rooms[roomID]; ok { for client := range room { if excludeClient != nil && client == excludeClient { continue } select { case client.send <- data: default: } } } if sseRoom, ok := h.sseRooms[roomID]; ok { for sseClient := range sseRoom { select { case sseClient.send <- data: default: } } } } func (c *Client) SendEvent(event string, payload interface{}) { msg := domain.WSMessage{ Event: event, Payload: payload, Timestamp: domain.NowISO(), } data, err := json.Marshal(msg) if err == nil { select { case c.send <- data: default: } } } func (h *Hub) ServeWS(w http.ResponseWriter, r *http.Request, roomID, userID string) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("[ws] upgrade error: %v", err) return } client := &Client{ hub: h, conn: conn, roomID: roomID, userID: userID, send: make(chan []byte, 128), } h.Register(client) go client.writePump() go client.readPump() } func (h *Hub) ServeSSE(w http.ResponseWriter, r *http.Request, roomID, userID string, initialRooms []string) { flusher, ok := w.(http.Flusher) if !ok { http.Error(w, "Streaming unsupported!", http.StatusInternalServerError) return } w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache, no-transform") w.Header().Set("Connection", "keep-alive") w.Header().Set("X-Accel-Buffering", "no") w.Header().Set("Access-Control-Allow-Origin", "*") client := &SSEClient{ id: domain.NewID("sse"), userID: userID, roomID: roomID, send: make(chan []byte, 128), subscribedRooms: make(map[string]bool), } for _, room := range initialRooms { if strings.TrimSpace(room) != "" { client.subscribedRooms[strings.TrimSpace(room)] = true } } h.RegisterSSE(client) defer h.UnregisterSSE(client) // Send initial connected comment and event fmt.Fprintf(w, ": connected (transport: sse, id: %s)\n\n", client.id) connectedMsg, _ := json.Marshal(domain.WSMessage{ Event: "SSE_CONNECTED", Payload: map[string]interface{}{ "sseId": client.id, "transport": "sse", "userId": userID, }, Timestamp: domain.NowISO(), }) fmt.Fprintf(w, "data: %s\n\n", connectedMsg) flusher.Flush() // If initial rooms had doc rooms, send CRDT_INIT if h.crdtStore != nil { for rID := range client.subscribedRooms { if strings.HasPrefix(rID, "doc:") { docID := strings.TrimPrefix(rID, "doc:") stateBlob, updates, err := h.crdtStore.GetDocumentCRDTState(docID) if err == nil && (len(stateBlob) > 0 || len(updates) > 0) { var b64Updates []string for _, u := range updates { b64Updates = append(b64Updates, base64.StdEncoding.EncodeToString(u)) } b64State := "" if len(stateBlob) > 0 { b64State = base64.StdEncoding.EncodeToString(stateBlob) } initMsg := domain.WSMessage{ Event: "CRDT_INIT", Payload: map[string]interface{}{ "docId": docID, "state": b64State, "updates": b64Updates, }, Timestamp: domain.NowISO(), } if initBytes, err := json.Marshal(initMsg); err == nil { fmt.Fprintf(w, "data: %s\n\n", initBytes) flusher.Flush() } } } } } ticker := time.NewTicker(15 * time.Second) defer ticker.Stop() notify := r.Context().Done() for { select { case <-notify: return case msg, ok := <-client.send: if !ok { return } fmt.Fprintf(w, "data: %s\n\n", msg) flusher.Flush() case <-ticker.C: // Periodic ping comment to keep HTTP connection alive and prevent proxy timeouts fmt.Fprintf(w, ": ping\n\n") flusher.Flush() } } } func (c *Client) readPump() { defer func() { c.hub.Unregister(c) _ = c.conn.Close() }() c.conn.SetReadLimit(4 * 1024 * 1024) // 4MB for large document CRDT updates _ = 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 { break } // Handle client-initiated lock, CRDT, room subscription, and presence messages var incoming struct { Event string `json:"event"` DocID string `json:"docId"` Room string `json:"room"` Payload json.RawMessage `json:"payload"` } if json.Unmarshal(message, &incoming) == nil && c.userID != "" && c.userID != "anonymous" { docID := incoming.DocID if docID == "" && len(incoming.Payload) > 0 { var p struct { DocID string `json:"docId"` } _ = json.Unmarshal(incoming.Payload, &p) docID = p.DocID } // Room subscription without requiring docID if incoming.Event == "JOIN_ROOM" { room := incoming.Room if room == "" { room = docID } if room != "" { c.hub.JoinRoom(c, room) } } else if incoming.Event == "LEAVE_ROOM" { room := incoming.Room if room == "" { room = docID } if room != "" { c.hub.LeaveRoom(c, room) } } if docID != "" { switch incoming.Event { case "LOCK_DOC": if info, ok, _ := c.hub.AcquireDocLock(docID, c.userID, "", 30*time.Second); ok && info != nil { c.hub.BroadcastDocEvent(docID, "DOC_LOCKED", info) } case "HEARTBEAT_DOC": c.hub.HeartbeatDocLock(docID, c.userID, 30*time.Second) case "UNLOCK_DOC": if c.hub.ReleaseDocLock(docID, c.userID, false) { c.hub.BroadcastDocEvent(docID, "DOC_UNLOCKED", map[string]string{"docId": docID, "doc_id": docID}) } case "CRDT_JOIN": c.hub.JoinRoom(c, "doc:"+docID) if c.hub.crdtStore != nil { stateBlob, updates, err := c.hub.crdtStore.GetDocumentCRDTState(docID) if err == nil && (len(stateBlob) > 0 || len(updates) > 0) { var b64Updates []string for _, u := range updates { b64Updates = append(b64Updates, base64.StdEncoding.EncodeToString(u)) } b64State := "" if len(stateBlob) > 0 { b64State = base64.StdEncoding.EncodeToString(stateBlob) } c.SendEvent("CRDT_INIT", map[string]interface{}{ "docId": docID, "state": b64State, "updates": b64Updates, }) } } case "CRDT_UPDATE": var p struct { Update string `json:"update"` } _ = json.Unmarshal(incoming.Payload, &p) if p.Update != "" { c.hub.BroadcastRoom("doc:"+docID, "CRDT_UPDATE", map[string]interface{}{ "docId": docID, "update": p.Update, "fromUserId": c.userID, }, c) if c.hub.crdtStore != nil { if raw, err := base64.StdEncoding.DecodeString(p.Update); err == nil && len(raw) > 0 { _ = c.hub.crdtStore.AppendDocumentCRDTUpdate(docID, raw, c.userID) } } } case "CRDT_COMPACT": var p struct { State string `json:"state"` } _ = json.Unmarshal(incoming.Payload, &p) if p.State != "" && c.hub.crdtStore != nil { if raw, err := base64.StdEncoding.DecodeString(p.State); err == nil && len(raw) > 0 { _ = c.hub.crdtStore.SaveDocumentCRDTState(docID, raw, c.userID) } } case "CRDT_AWARENESS": c.hub.BroadcastRoom("doc:"+docID, "CRDT_AWARENESS", incoming.Payload, c) case "CRDT_LEAVE": c.hub.LeaveRoom(c, "doc:"+docID) c.hub.BroadcastRoom("doc:"+docID, "CRDT_AWARENESS_LEAVE", map[string]interface{}{ "docId": docID, "userId": c.userID, }, c) } } } } } func (c *Client) writePump() { ticker := time.NewTicker(30 * 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 } if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil { return } case <-ticker.C: _ = c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } }