mutiny_rs/
websocket.rs

1use std::{sync::Arc, time::{SystemTime, UNIX_EPOCH}};
2use std::time::Duration;
3use futures_util::{stream::{SplitSink, SplitStream}, SinkExt, StreamExt};
4use serde_json::json;
5use tokio::{net::TcpStream, spawn, sync::Mutex, task::JoinHandle};
6use tokio_tungstenite::{connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream};
7use tokio_tungstenite::tungstenite::Utf8Bytes;
8
9// Internal crate imports
10use crate::{client::EventHandler, context::Context, model::user::User};
11use crate::model::ready::Ready;
12use crate::client::ClientCache;
13use crate::model::message::Message as ChatMessage;
14use crate::model::user::RelationshipStatus;
15
16// Type aliases for readability
17type WsWriter = SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>;
18type WsReader = SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>;
19
20pub struct WebSocket {
21    pub writer: Arc<Mutex<WsWriter>>,
22    handler: Arc<Box<dyn EventHandler>>,
23}
24
25impl WebSocket {
26    /// Connects to the Gateway, sends Auth, and spawns the background loop.
27    /// Returns the Client (for sending messages) and the background Task Handle (to keep main alive).
28    pub async fn connect(handler: Box<dyn EventHandler>, token: String) -> (Arc<WebSocket>, JoinHandle<()>) {
29
30        let (ws_stream, _) = connect_async("wss://ws.revolt.chat").await.expect("Failed to connect to WebSocket");
31        let (writer, reader) = ws_stream.split();
32
33        // Wrap writer in Mutex for shared access
34        let writer_arc = Arc::new(Mutex::new(writer));
35
36        writer_arc.lock().await.send(Message::Text(Utf8Bytes::from(json!({
37            "type": "Authenticate",
38            "token": token
39        }).to_string()))).await.expect("Failed to send Authenticate packet");
40
41        let ws_client = Arc::new(WebSocket {
42            writer: writer_arc.clone(),
43            handler: Arc::from(handler),
44        });
45
46        let handler_token = token.clone();
47        let handler_event = ws_client.handler.clone();
48
49        let handle = spawn(async move {
50            Self::handler(reader, writer_arc, handler_token, handler_event).await;
51        });
52
53        (ws_client, handle)
54    }
55
56    /// The main Event Loop. Running in the background.
57    async fn handler(
58        mut reader: WsReader,
59        writer: Arc<Mutex<WsWriter>>,
60        token: String,
61        event: Arc<Box<dyn EventHandler>>,
62    ) {
63        let mut bot: Option<User> = None;
64
65        // Since we are in a loop, this cache persists as long as the connection is alive.
66        let cache = ClientCache::new();
67
68        let hb_writer = writer.clone();
69        spawn(async move {
70            loop {
71                tokio::time::sleep(Duration::from_secs(30)).await;
72                let ping = json!({
73                    "type": "Ping",
74                    "timestamp": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64
75                });
76                if hb_writer.lock().await.send(Message::Text(ping.to_string().into())).await.is_err() {
77                    break; // Stop heartbeat if writer is dead
78                }
79            }
80        });
81
82        // --- Main Read Loop ---
83        while let Some(msg_result) = reader.next().await {
84            let message = match msg_result {
85                Ok(m) => m,
86                Err(e) => {
87                    eprintln!("WebSocket Error: {}", e);
88                    break;
89                }
90            };
91
92            if message.is_text() {
93                let raw_text = message.to_text().unwrap_or("");
94
95                let json_value: serde_json::Value = match serde_json::from_str(raw_text) {
96                    Ok(v) => v,
97                    Err(e) => {
98                        eprintln!("JSON Parse Error: {}", e);
99                        continue;
100                    }
101                };
102
103                if let Some(event_type) = json_value["type"].as_str() {
104                    match event_type {
105                        "Ready" => {
106                            match serde_json::from_value::<Ready>(json_value.clone()) {
107                                Ok(ready) => {
108                                    bot = ready.users.iter()
109                                        .find(|u| u.relationship == RelationshipStatus::User)
110                                        .cloned();
111
112                                    if bot.is_none() {
113                                        eprintln!("Warning: Could not find own Bot User in Ready payload!");
114                                    }
115                                    cache.hydrate(&ready).await;
116
117                                    if let Some(ref bot_user) = bot {
118                                        let ctx = Context::new(
119                                            &token,
120                                            json_value,
121                                            writer.clone(),
122                                            bot_user.clone(),
123                                            cache.clone()
124                                        );
125                                        event.ready(ctx, ready).await;
126                                    }
127                                }
128                                Err(e) => eprintln!("Failed to parse Ready struct: {}", e),
129                            }
130                        }
131
132                        "Message" => {
133                            if let Some(ref bot_user) = bot {
134                                let ctx = Context::new(
135                                    &token,
136                                    json_value.clone(),
137                                    writer.clone(),
138                                    bot_user.clone(),
139                                    cache.clone()
140                                );
141
142                                match serde_json::from_value::<ChatMessage>(json_value) {
143                                    Ok(msg) => {
144                                        ctx.cache.messages.insert(msg.id.clone(), msg.clone()).await;
145
146                                        event.message(ctx, msg).await;
147                                    }
148                                    Err(e) => eprintln!("Failed to parse Message: {}", e),
149                                }
150                            }
151                        }
152
153                        _ => {
154                            // Ignore other events for now
155                        }
156                    }
157                }
158            } else if message.is_close() {
159                println!("Gateway closed connection.");
160                break;
161            }
162        }
163    }
164}
mutiny_rs/websocket.rs

mutiny_rs/
websocket.rs
