package gateway import ( "context" "database/sql" "encoding/json" "log/slog" "net/http" "strings" "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 }, } // SetAllowedOrigins replaces the default upgrader with one that validates // the Origin header. Accepts if: // 1. Exact match in origins list, OR // 2. Hostname matches a trusted hostname, OR // 3. Hostname matches the request's own Host header (same-origin), OR // 4. Localhost variant func SetAllowedOrigins(origins []string) { originSet := make(map[string]struct{}, len(origins)) hostSet := make(map[string]struct{}, len(origins)) for _, o := range origins { originSet[o] = struct{}{} h := o if idx := strings.Index(h, "://"); idx >= 0 { h = h[idx+3:] } if idx := strings.Index(h, "/"); idx >= 0 { h = h[:idx] } if idx := strings.LastIndex(h, ":"); idx >= 0 { h = h[:idx] } if h != "" { hostSet[h] = struct{}{} } } upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r *http.Request) bool { origin := r.Header.Get("Origin") if origin == "" { return true } if _, ok := originSet[origin]; ok { return true } oh := extractHostname(origin) // Match trusted hostnames if _, ok := hostSet[oh]; ok { return true } // Match request's own Host header (same-origin) reqHost := extractHostname(r.Host) if reqHost != "" && oh == reqHost { return true } // Localhost variants if oh == "localhost" || oh == "127.0.0.1" || oh == "::1" { return true } return false }, } } func extractHostname(s string) string { h := s if idx := strings.Index(h, "://"); idx >= 0 { h = h[idx+3:] } if idx := strings.Index(h, "/"); idx >= 0 { h = h[:idx] } if idx := strings.LastIndex(h, ":"); idx >= 0 { h = h[:idx] } h = strings.TrimPrefix(h, "[") h = strings.TrimSuffix(h, "]") return h } // Client is a middleman between the websocket connection and the hub. type Client struct { Hub *Hub Conn *websocket.Conn UserID string Username 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 } c.Hub.RecordActivity(c.UserID) switch event.Type { case EventTypingStart: // Enrich with sender info so receivers know who is typing var data map[string]interface{} if raw, ok := event.Data.(json.RawMessage); ok { json.Unmarshal(raw, &data) } if data == nil { data = make(map[string]interface{}) } data["user_id"] = c.UserID data["username"] = c.Username c.Hub.BroadcastEvent(Event{Type: event.Type, Data: data}) case EventPresenceUpdate: c.Hub.BroadcastEvent(event) case EventVoiceJoin, EventVoiceLeave, EventVoiceMute, EventVoiceDeafen: // Broadcast voice state changes to all clients with sender info var vData map[string]interface{} if raw, ok := event.Data.(json.RawMessage); ok { json.Unmarshal(raw, &vData) } if vData == nil { vData = make(map[string]interface{}) } vData["user_id"] = c.UserID vData["username"] = c.Username c.Hub.BroadcastEvent(Event{Type: event.Type, Data: vData}) case EventVoiceWhisper: // Forward voice whispers only to the target user, not broadcast var whisperData struct { TargetUserID string `json:"target_user_id"` FromUserID string `json:"from_user_id"` FromUsername string `json:"from_username"` Message string `json:"message,omitempty"` } if raw, ok := event.Data.(json.RawMessage); ok { if err := json.Unmarshal(raw, &whisperData); err == nil { whisperData.FromUserID = c.UserID c.Hub.SendToUser(whisperData.TargetUserID, Event{ Type: EventVoiceWhisper, Data: whisperData, }) } } 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) 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 } } } } // authMessage is the first frame the client must send after connecting (for TUI/bot clients). type authMessage struct { Token string `json:"token"` } // ServeWS handles websocket requests from the peer. // It authenticates via cookie (browser) or message-frame (TUI/bots). func ServeWS(db *sql.DB, hub *Hub, logger *slog.Logger, w http.ResponseWriter, r *http.Request, cookieName string) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { logger.Error("websocket upgrade failed", "error", err) return } // Try cookie-based auth first (browser clients). var token string if cookie, cookieErr := r.Cookie(cookieName); cookieErr == nil && cookie.Value != "" { token = cookie.Value } // Fall back to message-frame auth (TUI, bots). if token == "" { conn.SetReadDeadline(time.Now().Add(10 * time.Second)) _, raw, readErr := conn.ReadMessage() if readErr != nil { logger.Warn("ws auth: failed to read auth message", "error", readErr) conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"auth timeout"}`)) conn.Close() return } var auth authMessage if jsonErr := json.Unmarshal(raw, &auth); jsonErr != nil || auth.Token == "" { logger.Warn("ws auth: invalid auth message") conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"missing token"}`)) conn.Close() return } token = auth.Token } var userID, username string err = db.QueryRowContext(context.Background(), `SELECT u.id, u.username FROM sessions s JOIN users u ON u.id = s.user_id WHERE s.token = $1 AND s.expires_at > NOW()`, token, ).Scan(&userID, &username) if err != nil { logger.Warn("ws auth: invalid session", "error", err) conn.WriteMessage(websocket.TextMessage, []byte(`{"error":"invalid session"}`)) conn.Close() return } conn.SetReadDeadline(time.Time{}) conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"ready"}`)) client := &Client{ Hub: hub, Conn: conn, UserID: userID, Username: username, send: make(chan []byte, 256), } hub.Register(client) go client.writePump() go client.readPump() }