package ws import ( "encoding/json" "fmt" "log/slog" "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 sendMu sync.RWMutex closed bool } func (c *Client) closeSend() { c.sendMu.Lock() defer c.sendMu.Unlock() if !c.closed { c.closed = true close(c.send) } } func (c *Client) trySend(data []byte) { c.sendMu.RLock() defer c.sendMu.RUnlock() if c.closed { return } select { case c.send <- data: default: } } type SSEClient struct { id string userID string roomID string send chan []byte subscribedRooms map[string]bool sendMu sync.RWMutex closed bool } func (c *SSEClient) closeSend() { c.sendMu.Lock() defer c.sendMu.Unlock() if !c.closed { c.closed = true close(c.send) } } func (c *SSEClient) trySend(data []byte) { c.sendMu.RLock() defer c.sendMu.RUnlock() if c.closed { return } select { case c.send <- data: default: } } type docLockItem struct { docID string userID string userName string lockedAt time.Time expiresAt time.Time } type blockLockItem struct { docID string blockID string userID string userName string userColor string lockedAt time.Time expiresAt time.Time } type liveSessionItem struct { docID string initiatorID string initiatorName string startedAt time.Time expiresAt time.Time } type DocStore interface { LoadDocumentState(docID string) (string, []domain.DocumentBlock, error) FlushDocumentState(docID string, title string, blocks []domain.DocumentBlock) 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 blockLocks map[string]*blockLockItem // key: docID + ":" + blockID liveSessions map[string]*liveSessionItem // key: docID docRooms map[string]*DocRoom docStore DocStore stopCleanup chan struct{} closed bool } 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), blockLocks: make(map[string]*blockLockItem), liveSessions: make(map[string]*liveSessionItem), docRooms: make(map[string]*DocRoom), stopCleanup: make(chan struct{}), } go h.cleanupLocksLoop() return h } func (h *Hub) SetDocStore(store DocStore) { h.mu.Lock() defer h.mu.Unlock() h.docStore = 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) } } c.closeSend() // 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) } } c.closeSend() 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), } } // Block Locking Methods (Option B) func (h *Hub) AcquireBlockLock(docID, blockID, userID, userName, userColor string, ttl time.Duration) (*domain.BlockLockInfo, bool, *domain.BlockLockInfo) { h.mu.Lock() defer h.mu.Unlock() now := time.Now() key := docID + ":" + blockID existing, exists := h.blockLocks[key] if exists && existing.expiresAt.After(now) { if existing.userID != userID { info := &domain.BlockLockInfo{ DocID: existing.docID, BlockID: existing.blockID, UserID: existing.userID, UserName: existing.userName, UserColor: existing.userColor, ExpiresAt: existing.expiresAt.Format(time.RFC3339), } return nil, false, info } // Refresh lease existing.expiresAt = now.Add(ttl) if userName != "" { existing.userName = userName } if userColor != "" { existing.userColor = userColor } info := &domain.BlockLockInfo{ DocID: existing.docID, BlockID: existing.blockID, UserID: existing.userID, UserName: existing.userName, UserColor: existing.userColor, ExpiresAt: existing.expiresAt.Format(time.RFC3339), } return info, true, nil } item := &blockLockItem{ docID: docID, blockID: blockID, userID: userID, userName: userName, userColor: userColor, lockedAt: now, expiresAt: now.Add(ttl), } h.blockLocks[key] = item info := &domain.BlockLockInfo{ DocID: item.docID, BlockID: item.blockID, UserID: item.userID, UserName: item.userName, UserColor: item.userColor, ExpiresAt: item.expiresAt.Format(time.RFC3339), } return info, true, nil } func (h *Hub) ReleaseBlockLock(docID, blockID, userID string) bool { h.mu.Lock() defer h.mu.Unlock() key := docID + ":" + blockID existing, exists := h.blockLocks[key] if !exists { return false } if userID != "" && existing.userID != userID && userID != "admin" { return false } delete(h.blockLocks, key) return true } func (h *Hub) GetDocBlockLocks(docID string) []domain.BlockLockInfo { h.mu.RLock() defer h.mu.RUnlock() now := time.Now() var result []domain.BlockLockInfo prefix := docID + ":" for k, l := range h.blockLocks { if strings.HasPrefix(k, prefix) && l.expiresAt.After(now) { result = append(result, domain.BlockLockInfo{ DocID: l.docID, BlockID: l.blockID, UserID: l.userID, UserName: l.userName, UserColor: l.userColor, ExpiresAt: l.expiresAt.Format(time.RFC3339), }) } } return result } // Live Session Methods (Opt-In Live Collaboration) func (h *Hub) StartLiveSession(docID, userID, userName string, ttl time.Duration) (*domain.LiveSessionInfo, bool, *domain.LiveSessionInfo) { h.mu.Lock() defer h.mu.Unlock() now := time.Now() existing, exists := h.liveSessions[docID] if exists && existing.expiresAt.After(now) { if existing.initiatorID != userID { info := &domain.LiveSessionInfo{ DocID: existing.docID, InitiatorID: existing.initiatorID, InitiatorName: existing.initiatorName, StartedAt: existing.startedAt.Format(time.RFC3339), ExpiresAt: existing.expiresAt.Format(time.RFC3339), } return nil, false, info } // Refresh lease existing.expiresAt = now.Add(ttl) info := &domain.LiveSessionInfo{ DocID: existing.docID, InitiatorID: existing.initiatorID, InitiatorName: existing.initiatorName, StartedAt: existing.startedAt.Format(time.RFC3339), ExpiresAt: existing.expiresAt.Format(time.RFC3339), } return info, true, nil } item := &liveSessionItem{ docID: docID, initiatorID: userID, initiatorName: userName, startedAt: now, expiresAt: now.Add(ttl), } h.liveSessions[docID] = item info := &domain.LiveSessionInfo{ DocID: item.docID, InitiatorID: item.initiatorID, InitiatorName: item.initiatorName, StartedAt: item.startedAt.Format(time.RFC3339), ExpiresAt: item.expiresAt.Format(time.RFC3339), } return info, true, nil } func (h *Hub) EndLiveSession(docID, userID string) bool { h.mu.Lock() defer h.mu.Unlock() existing, exists := h.liveSessions[docID] if !exists { return false } if userID != "" && existing.initiatorID != userID && userID != "admin" { return false } delete(h.liveSessions, docID) return true } func (h *Hub) GetLiveSession(docID string) *domain.LiveSessionInfo { h.mu.RLock() defer h.mu.RUnlock() now := time.Now() existing, exists := h.liveSessions[docID] if !exists || existing.expiresAt.Before(now) { return nil } return &domain.LiveSessionInfo{ DocID: existing.docID, InitiatorID: existing.initiatorID, InitiatorName: existing.initiatorName, StartedAt: existing.startedAt.Format(time.RFC3339), ExpiresAt: existing.expiresAt.Format(time.RFC3339), } } func (h *Hub) cleanupLocksLoop() { ticker := time.NewTicker(2 * time.Second) defer ticker.Stop() for { select { case <-h.stopCleanup: return case <-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) } } type expBlock struct { docID string blockID string } var expiredBlocks []expBlock for k, bl := range h.blockLocks { if bl.expiresAt.Before(now) { expiredBlocks = append(expiredBlocks, expBlock{docID: bl.docID, blockID: bl.blockID}) delete(h.blockLocks, k) } } var expiredLiveSessions []string for docID, ls := range h.liveSessions { if ls.expiresAt.Before(now) { expiredLiveSessions = append(expiredLiveSessions, docID) delete(h.liveSessions, docID) } } h.mu.Unlock() for _, docID := range expiredDocIDs { h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID, "doc_id": docID}) } for _, b := range expiredBlocks { h.BroadcastRoom("doc:"+b.docID, "DOC_BLOCK_UNLOCKED", map[string]string{ "docId": b.docID, "blockId": b.blockID, }, nil) } for _, docID := range expiredLiveSessions { h.BroadcastRoom("doc:"+docID, "DOC_LIVE_ENDED", map[string]string{ "docId": docID, }, nil) } } } } // Close gracefully shuts down the Hub, stops background tickers, flushes dirty document rooms to DB, // and immediately disconnects all WebSocket and SSE connections so HTTP shutdown completes in milliseconds. func (h *Hub) Close() { h.mu.Lock() if h.closed { h.mu.Unlock() return } h.closed = true if h.stopCleanup != nil { select { case <-h.stopCleanup: default: close(h.stopCleanup) } } docRooms := make([]*DocRoom, 0, len(h.docRooms)) for _, r := range h.docRooms { docRooms = append(docRooms, r) } allWsClients := make([]*Client, 0, len(h.allClients)) for c := range h.allClients { allWsClients = append(allWsClients, c) } allSSEClients := make([]*SSEClient, 0, len(h.sseClients)) for c := range h.sseClients { allSSEClients = append(allSSEClients, c) } h.mu.Unlock() // 1. Flush any active dirty document rooms and close doc WS connections for _, r := range docRooms { r.Flush() r.mu.Lock() if r.debounceTimer != nil { r.debounceTimer.Stop() r.debounceTimer = nil } for c := range r.wsClients { _ = c.conn.Close() } r.wsClients = make(map[*DocWSClient]bool) r.mu.Unlock() } // 2. Close all WebSocket connections for _, c := range allWsClients { _ = c.conn.Close() c.closeSend() } // 3. Close all SSE client channels so ServeSSE select loop exits immediately for _, c := range allSSEClients { c.closeSend() } } // 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() if h.closed { h.mu.RUnlock() return } defer h.mu.RUnlock() for client := range h.allClients { client.trySend(data) } for sseClient := range h.sseClients { sseClient.trySend(data) } } 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() if h.closed { h.mu.RUnlock() return } defer h.mu.RUnlock() if room, ok := h.userRooms[userID]; ok { for client := range room { client.trySend(data) } } if sseRoom, ok := h.sseUserRooms[userID]; ok { for sseClient := range sseRoom { sseClient.trySend(data) } } } func (h *Hub) BroadcastTaskEvent(userID, event string, payload interface{}) { h.BroadcastToUser(userID, event, payload) } func (h *Hub) BroadcastDocEvent(docID, event string, payload interface{}) { if docID != "" { h.BroadcastRoom("doc:"+docID, event, payload, nil) } else { 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() if h.closed { h.mu.RUnlock() return } defer h.mu.RUnlock() targetRoom := listID if targetRoom == "" { targetRoom = spaceID } if room, ok := h.rooms[targetRoom]; ok { for client := range room { client.trySend(data) } } if sseRoom, ok := h.sseRooms[targetRoom]; ok { for sseClient := range sseRoom { sseClient.trySend(data) } } } 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() if h.closed { h.mu.RUnlock() return } defer h.mu.RUnlock() if room, ok := h.rooms[roomID]; ok { for client := range room { if excludeClient != nil && client == excludeClient { continue } client.trySend(data) } } if sseRoom, ok := h.sseRooms[roomID]; ok { for sseClient := range sseRoom { sseClient.trySend(data) } } } 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 { c.trySend(data) } } func (h *Hub) ServeWS(w http.ResponseWriter, r *http.Request, roomID, userID string) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { slog.Error("WebSocket upgrade error", "room_id", roomID, "user_id", userID, "error", 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() 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}) } } } } } } 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 } } } }