package main import ( "encoding/json" "fmt" "log" "net/url" "strings" "sync" "time" tea "github.com/charmbracelet/bubbletea" "github.com/gorilla/websocket" ) // WSEvent represents a raw websocket event from the server. type WSEvent struct { Type string `json:"type"` Data json.RawMessage `json:"data,omitempty"` } // WebSocket connection manager. type WSConn struct { conn *websocket.Conn mu sync.Mutex done chan struct{} closed bool } // ConnectWS dials the websocket endpoint and authenticates via an initial // message frame (not a URL query parameter, to avoid leaking the token in logs). func ConnectWS(serverURL, token string) (*WSConn, error) { u, err := url.Parse(serverURL) if err != nil { return nil, err } // Switch to ws/wss scheme if u.Scheme == "https" { u.Scheme = "wss" } else { u.Scheme = "ws" } u.Path = "/ws" conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil) if err != nil { return nil, err } // Send auth token as the first message frame. authMsg, _ := json.Marshal(map[string]string{"token": token}) if err := conn.WriteMessage(websocket.TextMessage, authMsg); err != nil { conn.Close() return nil, fmt.Errorf("send auth: %w", err) } // Wait for the server's ready response. conn.SetReadDeadline(time.Now().Add(10 * time.Second)) _, raw, err := conn.ReadMessage() if err != nil { conn.Close() return nil, fmt.Errorf("auth timeout: %w", err) } var resp struct { Type string `json:"type"` Error string `json:"error"` } if err := json.Unmarshal(raw, &resp); err != nil || resp.Type != "ready" { conn.Close() msg := "unexpected response" if resp.Error != "" { msg = resp.Error } return nil, fmt.Errorf("auth failed: %s", msg) } conn.SetReadDeadline(time.Time{}) ws := &WSConn{ conn: conn, done: make(chan struct{}), } return ws, nil } // Listen starts reading messages and sends them as tea.Msgs to the program. // It returns a tea.Cmd that listens for the next message. func (ws *WSConn) Listen() tea.Cmd { return func() tea.Msg { _, raw, err := ws.conn.ReadMessage() if err != nil { return wsErrMsg{err: err} } var event WSEvent if err := json.Unmarshal(raw, &event); err != nil { return wsErrMsg{err: err} } return wsEventMsg(event) } } // Send sends a JSON event to the server. func (ws *WSConn) Send(eventType string, data interface{}) error { ws.mu.Lock() defer ws.mu.Unlock() msg := map[string]interface{}{ "type": eventType, "data": data, } payload, err := json.Marshal(msg) if err != nil { return err } return ws.conn.WriteMessage(websocket.TextMessage, payload) } // Close closes the websocket connection. func (ws *WSConn) Close() { ws.mu.Lock() defer ws.mu.Unlock() if !ws.closed { ws.closed = true ws.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) ws.conn.Close() close(ws.done) } } // StartPing sends periodic pings to keep the connection alive. func (ws *WSConn) StartPing() { go func() { ticker := time.NewTicker(30 * time.Second) defer ticker.Stop() for { select { case <-ticker.C: ws.mu.Lock() if !ws.closed { _ = ws.conn.WriteMessage(websocket.PingMessage, nil) } ws.mu.Unlock() case <-ws.done: return } } }() } // ---- Message types for bubbletea ---- type wsEventMsg WSEvent type wsErrMsg struct { err error } // Parse helpers for ws event data type wsMessageData struct { ID string `json:"id"` ChannelID string `json:"channel_id"` AuthorID string `json:"author_id"` AuthorName string `json:"author_username"` DisplayName *string `json:"author_display_name"` Content string `json:"content"` EditedAt *string `json:"edited_at"` CreatedAt string `json:"created_at"` } type wsPresenceData struct { UserID string `json:"user_id"` Username string `json:"username"` Status string `json:"status"` } type wsTypingData struct { ChannelID string `json:"channel_id"` UserID string `json:"user_id"` Username string `json:"username"` } type wsVoiceData struct { RoomName string `json:"room_name"` UserID string `json:"user_id"` Username string `json:"username"` } // parseWSEventMessage extracts a Message from event data. func parseWSEventMessage(data json.RawMessage) (*Message, error) { var m Message if err := json.Unmarshal(data, &m); err != nil { return nil, err } return &m, nil } // parseWSDeleteData extracts the message ID from a delete event. func parseWSDeleteData(data json.RawMessage) (id string, channelID string) { var d struct { ID string `json:"id"` ChannelID string `json:"channel_id"` } if err := json.Unmarshal(data, &d); err != nil { log.Printf("ws: failed to parse delete data: %v", err) return } return d.ID, d.ChannelID } // parseWSTyping extracts typing event data. func parseWSTyping(data json.RawMessage) (username string, channelID string) { var d struct { ChannelID string `json:"channel_id"` Username string `json:"username"` } if err := json.Unmarshal(data, &d); err != nil { return } return d.Username, d.ChannelID } // trimURLScheme strips http(s):// for display func trimURLScheme(s string) string { s = strings.TrimPrefix(s, "https://") s = strings.TrimPrefix(s, "http://") return s }