webex_message_handler/
mercury_socket.rs

1//! Mercury WebSocket connection with auth, heartbeat, and reconnection.
2
3use crate::errors::WebexError;
4use crate::types::MercuryActivity;
5use futures_util::{SinkExt, StreamExt};
6use serde_json::Value;
7use std::sync::Arc;
8use std::time::Duration;
9use tokio::sync::{mpsc, Mutex, Notify};
10use tokio::time;
11use tokio_tungstenite::{connect_async, tungstenite::Message};
12use tracing::{debug, error, info};
13use url::Url;
14use uuid::Uuid;
15
16/// Type alias for the injected WebSocket factory.
17#[allow(dead_code)]
18type WsFactoryFn = Arc<
19    dyn Fn(String) -> std::pin::Pin<
20        Box<dyn std::future::Future<
21            Output = Result<Box<dyn crate::types::InjectedWebSocket>, Box<dyn std::error::Error + Send + Sync>>,
22        > + Send>,
23    > + Send + Sync,
24>;
25
26/// Type alias for the WebSocket write half to reduce type complexity.
27type WsSink = futures_util::stream::SplitSink<
28    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
29    Message,
30>;
31
32/// Events emitted by the Mercury socket.
33#[derive(Debug, Clone)]
34pub enum MercuryEvent {
35    Connected,
36    Disconnected(String),
37    Reconnecting(u32),
38    Activity(Box<MercuryActivity>),
39    KmsResponse(Value),
40    Error(String),
41}
42
43/// Mercury WebSocket connection manager.
44pub struct MercurySocket {
45    #[allow(dead_code)]
46    ws_factory: Option<WsFactoryFn>,
47    ping_interval: Duration,
48    pong_timeout: Duration,
49    reconnect_backoff_max: Duration,
50    max_reconnect_attempts: u32,
51
52    token: Arc<Mutex<String>>,
53    base_url: Arc<Mutex<String>>,
54    connected: Arc<Mutex<bool>>,
55    should_reconnect: Arc<Mutex<bool>>,
56    reconnect_attempts: Arc<Mutex<u32>>,
57    shutdown: Arc<Notify>,
58
59    event_tx: mpsc::UnboundedSender<MercuryEvent>,
60    event_rx: Arc<Mutex<Option<mpsc::UnboundedReceiver<MercuryEvent>>>>,
61}
62
63impl MercurySocket {
64    pub fn new(
65        _ws_factory: Option<WsFactoryFn>,
66        ping_interval: Duration,
67        pong_timeout: Duration,
68        reconnect_backoff_max: Duration,
69        max_reconnect_attempts: u32,
70    ) -> Self {
71        let (event_tx, event_rx) = mpsc::unbounded_channel();
72
73        Self {
74            ws_factory: _ws_factory,
75            ping_interval,
76            pong_timeout,
77            reconnect_backoff_max,
78            max_reconnect_attempts,
79            token: Arc::new(Mutex::new(String::new())),
80            base_url: Arc::new(Mutex::new(String::new())),
81            connected: Arc::new(Mutex::new(false)),
82            should_reconnect: Arc::new(Mutex::new(true)),
83            reconnect_attempts: Arc::new(Mutex::new(0)),
84            shutdown: Arc::new(Notify::new()),
85            event_tx,
86            event_rx: Arc::new(Mutex::new(Some(event_rx))),
87        }
88    }
89
90    /// Take the event receiver. Can only be called once.
91    pub async fn take_event_rx(&self) -> Option<mpsc::UnboundedReceiver<MercuryEvent>> {
92        self.event_rx.lock().await.take()
93    }
94
95    /// Connect to Mercury WebSocket.
96    pub async fn connect(&self, ws_url: &str, token: &str) -> Result<(), WebexError> {
97        *self.token.lock().await = token.to_string();
98        *self.base_url.lock().await = ws_url.to_string();
99        *self.should_reconnect.lock().await = true;
100        *self.reconnect_attempts.lock().await = 0;
101        self.connect_internal().await
102    }
103
104    async fn connect_internal(&self) -> Result<(), WebexError> {
105        let base_url = self.base_url.lock().await.clone();
106        let prepared_url = Self::prepare_url(&base_url)?;
107        debug!("Connecting to Mercury at {prepared_url}");
108
109        *self.connected.lock().await = false;
110
111        let (ws_stream, _) = connect_async(&prepared_url)
112            .await
113            .map_err(|e| WebexError::mercury_connection(format!("Failed to connect: {e}"), None))?;
114
115        let (mut write, mut read) = ws_stream.split();
116
117        // Send authorization
118        let token = self.token.lock().await.clone();
119        let auth_msg = serde_json::json!({
120            "id": Uuid::new_v4().to_string(),
121            "type": "authorization",
122            "data": { "token": format!("Bearer {}", token) }
123        });
124        write
125            .send(Message::Text(auth_msg.to_string()))
126            .await
127            .map_err(|e| WebexError::mercury_connection(format!("Failed to send auth: {e}"), None))?;
128
129        // Wait for connection ready
130        let _ready_timeout = time::timeout(Duration::from_secs(30), async {
131            while let Some(msg) = read.next().await {
132                match msg {
133                    Ok(Message::Text(text)) => {
134                        let text_str: &str = &text;
135                        if let Ok(parsed) = serde_json::from_str::<Value>(text_str) {
136                            if Self::is_connection_ready(&parsed) {
137                                return Ok(parsed);
138                            }
139                        }
140                    }
141                    Err(e) => {
142                        return Err(WebexError::mercury_connection(
143                            format!("WebSocket error during setup: {e}"),
144                            None,
145                        ));
146                    }
147                    _ => {}
148                }
149            }
150            Err(WebexError::mercury_connection("WebSocket closed during setup", None))
151        })
152        .await
153        .map_err(|_| WebexError::mercury_connection("Mercury connection timeout", None))??;
154
155        debug!("Mercury connection ready");
156        *self.connected.lock().await = true;
157
158        // Spawn read loop and ping loop
159        let event_tx = self.event_tx.clone();
160        let connected = self.connected.clone();
161        let should_reconnect = self.should_reconnect.clone();
162        let reconnect_attempts = self.reconnect_attempts.clone();
163        let max_reconnect = self.max_reconnect_attempts;
164        let backoff_max = self.reconnect_backoff_max;
165        let ping_interval = self.ping_interval;
166        let _pong_timeout = self.pong_timeout;
167        let shutdown = self.shutdown.clone();
168        let base_url_clone = self.base_url.clone();
169        let token_clone = self.token.clone();
170
171        let write = Arc::new(Mutex::new(write));
172        let write_clone = write.clone();
173
174        // Ping loop
175        let ping_write = write.clone();
176        let ping_connected = connected.clone();
177        let ping_shutdown = shutdown.clone();
178        let _ping_event_tx = event_tx.clone();
179        tokio::spawn(async move {
180            let mut interval = time::interval(ping_interval);
181            interval.tick().await; // skip first tick
182
183            loop {
184                tokio::select! {
185                    _ = interval.tick() => {
186                        if !*ping_connected.lock().await {
187                            break;
188                        }
189                        let pong_id = Uuid::new_v4().to_string();
190                        let ping_msg = serde_json::json!({
191                            "id": pong_id,
192                            "type": "ping"
193                        });
194                        let mut w = ping_write.lock().await;
195                        if w.send(Message::Text(ping_msg.to_string())).await.is_err() {
196                            break;
197                        }
198                        debug!("Sent ping: {pong_id}");
199                        drop(w);
200
201                        // Pong timeout handled by read loop
202                    }
203                    _ = ping_shutdown.notified() => break,
204                }
205            }
206        });
207
208        // Read loop
209        tokio::spawn(async move {
210            while let Some(msg) = read.next().await {
211                match msg {
212                    Ok(Message::Text(text)) => {
213                        let text_str: &str = &text;
214                        debug!("WS message received ({} bytes)", text_str.len());
215                        if let Ok(parsed) = serde_json::from_str::<Value>(text_str) {
216                            Self::handle_message_static(&parsed, &event_tx, &write_clone).await;
217                        } else {
218                            debug!("Failed to parse WS message as JSON");
219                        }
220                    }
221                    Ok(Message::Close(frame)) => {
222                        let code = frame.as_ref().map(|f| f.code.into()).unwrap_or(1000u16);
223                        let reason = frame.as_ref().map(|f| f.reason.to_string()).unwrap_or_default();
224                        Self::handle_close_static(
225                            code,
226                            &reason,
227                            &connected,
228                            &should_reconnect,
229                            &reconnect_attempts,
230                            max_reconnect,
231                            backoff_max,
232                            &base_url_clone,
233                            &token_clone,
234                            &event_tx,
235                        )
236                        .await;
237                        break;
238                    }
239                    Err(e) => {
240                        error!("WebSocket error: {e}");
241                        let _ = event_tx.send(MercuryEvent::Error(e.to_string()));
242                        *connected.lock().await = false;
243                        break;
244                    }
245                    _ => {}
246                }
247            }
248
249            // Connection ended — handle reconnection if needed
250            if *should_reconnect.lock().await && !*connected.lock().await {
251                // Reconnect logic handled in handle_close_static
252            }
253        });
254
255        Ok(())
256    }
257
258    fn prepare_url(base_url: &str) -> Result<String, WebexError> {
259        let mut url = Url::parse(base_url)
260            .map_err(|e| WebexError::mercury_connection(format!("Invalid URL: {e}"), None))?;
261        url.query_pairs_mut()
262            .append_pair("outboundWireFormat", "text")
263            .append_pair("bufferStates", "true")
264            .append_pair("aliasHttpStatus", "true")
265            .append_pair(
266                "clientTimestamp",
267                &std::time::SystemTime::now()
268                    .duration_since(std::time::UNIX_EPOCH)
269                    .unwrap_or_default()
270                    .as_millis()
271                    .to_string(),
272            );
273        Ok(url.to_string())
274    }
275
276    fn is_connection_ready(message: &Value) -> bool {
277        let event_type = message
278            .get("data")
279            .and_then(|d| d.get("eventType"))
280            .and_then(|e| e.as_str())
281            .unwrap_or("");
282        event_type.contains("mercury.buffer_state") || event_type.contains("mercury.registration_status")
283    }
284
285    async fn handle_message_static(
286        message: &Value,
287        event_tx: &mpsc::UnboundedSender<MercuryEvent>,
288        write: &Arc<Mutex<WsSink>>,
289    ) {
290        let msg_type = message.get("type").and_then(|t| t.as_str()).unwrap_or("");
291
292        match msg_type {
293            "pong" => {
294                let id = message.get("id").and_then(|i| i.as_str()).unwrap_or("");
295                debug!("Received pong: {id}");
296            }
297            "shutdown" => {
298                info!("Received shutdown message from Mercury");
299                // Reconnection will be triggered by connection close
300            }
301            _ => {
302                if let Some(data) = message.get("data") {
303                    if let Some(event_type) = data.get("eventType").and_then(|e| e.as_str()) {
304                        debug!("Mercury eventType: {event_type}");
305
306                        // Send ACK
307                        if let Some(msg_id) = message.get("id").and_then(|i| i.as_str()) {
308                            let ack = serde_json::json!({"messageId": msg_id, "type": "ack"});
309                            let mut w = write.lock().await;
310                            let _ = w.send(Message::Text(ack.to_string())).await;
311                        }
312
313                        if event_type.starts_with("encryption.") {
314                            debug!("Emitting kms:response for eventType: {event_type}");
315                            let _ = event_tx.send(MercuryEvent::KmsResponse(data.clone()));
316                        } else if event_type == "conversation.activity" {
317                            if let Some(activity_raw) = data.get("activity") {
318                                match serde_json::from_value::<MercuryActivity>(activity_raw.clone()) {
319                                    Ok(activity) => {
320                                        debug!("Emitting activity: {}", activity.id);
321                                        let _ = event_tx.send(MercuryEvent::Activity(Box::new(activity)));
322                                    }
323                                    Err(e) => {
324                                        error!("Failed to parse activity: {e}");
325                                        debug!("Raw activity keys: {:?}", activity_raw.as_object().map(|o| o.keys().collect::<Vec<_>>()));
326                                    }
327                                }
328                            }
329                        }
330                    } else {
331                        debug!("Unhandled Mercury message, type={msg_type:?}, keys={:?}", message.as_object().map(|o| o.keys().collect::<Vec<_>>()));
332                    }
333                } else {
334                    debug!("Unhandled Mercury message, type={msg_type:?}, no data field");
335                }
336            }
337        }
338    }
339
340    #[allow(clippy::too_many_arguments)]
341    async fn handle_close_static(
342        code: u16,
343        reason: &str,
344        connected: &Arc<Mutex<bool>>,
345        should_reconnect: &Arc<Mutex<bool>>,
346        reconnect_attempts: &Arc<Mutex<u32>>,
347        max_reconnect: u32,
348        backoff_max: Duration,
349        _base_url: &Arc<Mutex<String>>,
350        _token: &Arc<Mutex<String>>,
351        event_tx: &mpsc::UnboundedSender<MercuryEvent>,
352    ) {
353        info!("WebSocket closed with code {code}: {reason}");
354        *connected.lock().await = false;
355
356        if code == 4401 {
357            error!("Mercury authorization failed");
358            *should_reconnect.lock().await = false;
359            let _ = event_tx.send(MercuryEvent::Error("Mercury authorization failed".into()));
360            let _ = event_tx.send(MercuryEvent::Disconnected("auth-failed".into()));
361            return;
362        }
363
364        if code == 4400 || code == 4403 {
365            error!("Mercury permanent failure (code {code})");
366            *should_reconnect.lock().await = false;
367            let _ = event_tx.send(MercuryEvent::Error(format!("Mercury permanent failure (code {code})")));
368            let _ = event_tx.send(MercuryEvent::Disconnected("permanent-failure".into()));
369            return;
370        }
371
372        if *should_reconnect.lock().await {
373            let mut attempts = reconnect_attempts.lock().await;
374            if *attempts >= max_reconnect {
375                error!("Max reconnection attempts ({max_reconnect}) exceeded");
376                *should_reconnect.lock().await = false;
377                let _ = event_tx.send(MercuryEvent::Disconnected("max-attempts-exceeded".into()));
378                return;
379            }
380            *attempts += 1;
381            let attempt = *attempts;
382            let delay_secs = (2.0f64.powi(attempt as i32 - 1)).min(backoff_max.as_secs_f64());
383            drop(attempts);
384
385            info!("Reconnecting (attempt {attempt}/{max_reconnect}) in {delay_secs}s");
386            let _ = event_tx.send(MercuryEvent::Reconnecting(attempt));
387
388            time::sleep(Duration::from_secs_f64(delay_secs)).await;
389
390            // Signal that reconnection should happen (handler will re-connect)
391            let _ = event_tx.send(MercuryEvent::Disconnected("reconnect-needed".into()));
392        } else {
393            let _ = event_tx.send(MercuryEvent::Disconnected("manual".into()));
394        }
395    }
396
397    /// Disconnect from Mercury.
398    pub async fn disconnect(&self) {
399        info!("Disconnecting from Mercury");
400        *self.should_reconnect.lock().await = false;
401        *self.connected.lock().await = false;
402        self.shutdown.notify_waiters();
403        let _ = self.event_tx.send(MercuryEvent::Disconnected("client".into()));
404    }
405
406    /// Whether the WebSocket is currently connected.
407    pub async fn connected(&self) -> bool {
408        *self.connected.lock().await
409    }
410
411    /// Current reconnection attempt count.
412    pub async fn current_reconnect_attempts(&self) -> u32 {
413        *self.reconnect_attempts.lock().await
414    }
415}
webex_message_handler/mercury_socket.rs

webex_message_handler/
mercury_socket.rs
