package gateway import ( "context" "database/sql" "encoding/json" "log/slog" "net/http" "time" "github.com/gorilla/websocket" ) const ( writeWait = 10 * time.Second pongWait = 60 * time.Second pingPeriod = (pongWait * 9) / 10 maxMessageSize = 4096 ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r *http.Request) bool { return true }, } // Client is a middleman between the websocket connection and the hub. type Client struct { Hub *Hub Conn *websocket.Conn UserID string send chan []byte } // readPump pumps messages from the websocket connection to the hub. func (c *Client) readPump() { defer func() { c.Hub.Unregister(c) c.Conn.Close() }() c.Conn.SetReadLimit(maxMessageSize) c.Conn.SetReadDeadline(time.Now().Add(pongWait)) c.Conn.SetPongHandler(func(string) error { c.Conn.SetReadDeadline(time.Now().Add(pongWait)) return nil }) for { _, raw, err := c.Conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) { c.Hub.logger.Warn("websocket read error", "error", err, "user_id", c.UserID) } break } var event Event if err := json.Unmarshal(raw, &event); err != nil { c.Hub.logger.Warn("invalid event JSON", "error", err, "user_id", c.UserID) continue } // Record activity for idle detection. c.Hub.RecordActivity(c.UserID) // Handle client-sent events (e.g. TYPING_START, PRESENCE_UPDATE) // For now, broadcast them to all clients switch event.Type { case EventTypingStart, EventPresenceUpdate: c.Hub.BroadcastEvent(event) default: c.Hub.logger.Info("received event from client", "type", event.Type, "user_id", c.UserID) } } } // writePump pumps messages from the hub to the websocket connection. func (c *Client) writePump() { ticker := time.NewTicker(pingPeriod) defer func() { ticker.Stop() c.Conn.Close() }() for { select { case message, ok := <-c.send: c.Conn.SetWriteDeadline(time.Now().Add(writeWait)) if !ok { c.Conn.WriteMessage(websocket.CloseMessage, []byte{}) return } w, err := c.Conn.NextWriter(websocket.TextMessage) if err != nil { return } w.Write(message) // Drain queued messages into the current write n := len(c.send) for i := 0; i < n; i++ { w.Write([]byte{'\n'}) w.Write(<-c.send) } if err := w.Close(); err != nil { return } case <-ticker.C: c.Conn.SetWriteDeadline(time.Now().Add(writeWait)) if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } } // ServeWS handles websocket requests from the peer. It authenticates the // client via a session token passed as the "token" query parameter. func ServeWS(db *sql.DB, hub *Hub, logger *slog.Logger, w http.ResponseWriter, r *http.Request) { token := r.URL.Query().Get("token") if token == "" { http.Error(w, `{"error":"missing token"}`, http.StatusUnauthorized) return } var userID string err := db.QueryRowContext(context.Background(), `SELECT user_id FROM sessions WHERE token = $1 AND expires_at > NOW()`, token, ).Scan(&userID) if err != nil { http.Error(w, `{"error":"invalid or expired session"}`, http.StatusUnauthorized) return } conn, err := upgrader.Upgrade(w, r, nil) if err != nil { logger.Error("websocket upgrade failed", "error", err) return } client := &Client{ Hub: hub, Conn: conn, UserID: userID, send: make(chan []byte, 256), } hub.Register(client) go client.writePump() go client.readPump() }