mail4agent_messenger_shell/
push.rs1use std::io::ErrorKind;
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::sync::{Arc, Mutex};
11use std::thread::{self, JoinHandle};
12use std::time::{Duration, Instant};
13
14use tungstenite::{stream::MaybeTlsStream, Message};
15
16use crate::{clip_public, ShellError};
17
18#[derive(Clone, Debug, PartialEq, Eq)]
20pub struct PushedRoomEvent {
21 pub room: String,
23 pub sender: String,
25 pub body: String,
29 pub event_id: String,
31 pub wire_type: String,
33}
34
35struct Incoming {
36 recipient: String,
37 event: PushedRoomEvent,
38}
39
40pub(crate) struct PushLink {
41 inbox: Arc<Mutex<Vec<Incoming>>>,
42 stop: Arc<AtomicBool>,
43 workers: Vec<JoinHandle<()>>,
44}
45
46impl Drop for PushLink {
47 fn drop(&mut self) {
48 self.stop.store(true, Ordering::Release);
49 for worker in self.workers.drain(..) {
50 let _ = worker.join();
51 }
52 }
53}
54
55impl PushLink {
56 pub(crate) fn open(base_url: &str, tokens: Vec<String>, separate: bool) -> Result<Self, ShellError> {
59 if tokens.is_empty() {
60 return Err(ShellError::Http("push socket has no session".into()));
61 }
62 let url = push_ws_url(base_url)?;
63 let inbox = Arc::new(Mutex::new(Vec::new()));
64 let stop = Arc::new(AtomicBool::new(false));
65 let groups: Vec<Vec<String>> = if separate { tokens.into_iter().map(|t| vec![t]).collect() } else { vec![tokens] };
66 let mut link = Self { inbox: Arc::clone(&inbox), stop: Arc::clone(&stop), workers: Vec::new() };
67 let mut readies = Vec::new();
68 let mut faileds = Vec::new();
69 for group in groups {
70 let ready = Arc::new(AtomicBool::new(false));
71 let failed = Arc::new(Mutex::new(None));
72 let worker = thread::spawn({
73 let (url, inbox, ready, failed, stop) = (url.clone(), Arc::clone(&inbox), Arc::clone(&ready), Arc::clone(&failed), Arc::clone(&stop));
74 move || worker_main(url, group, inbox, ready, failed, stop)
75 });
76 link.workers.push(worker);
77 readies.push(ready);
78 faileds.push(failed);
79 }
80 let start = Instant::now();
81 loop {
82 if readies.iter().all(|r| r.load(Ordering::Acquire)) {
83 return Ok(link);
84 }
85 for failed in &faileds {
86 if let Some(err) = failed.lock().unwrap_or_else(|err| err.into_inner()).clone() {
87 return Err(ShellError::Http(err)); }
89 }
90 if start.elapsed() > Duration::from_secs(5) {
91 return Err(ShellError::Http("push socket did not register".into()));
92 }
93 thread::sleep(Duration::from_millis(20));
94 }
95 }
96
97 pub(crate) fn drain(&self) -> Vec<(String, PushedRoomEvent)> {
98 self.inbox
99 .lock()
100 .unwrap_or_else(|err| err.into_inner())
101 .drain(..)
102 .map(|incoming| (incoming.recipient, incoming.event))
103 .collect()
104 }
105}
106
107fn push_ws_url(base: &str) -> Result<String, ShellError> {
108 let url = reqwest::Url::parse(base).map_err(|_| ShellError::BaseUrl)?;
109 let scheme = match url.scheme() {
110 "http" => "ws",
111 "https" => "wss",
112 _ => return Err(ShellError::BaseUrl),
113 };
114 let host = url
115 .host_str()
116 .filter(|host| !host.is_empty())
117 .ok_or(ShellError::BaseUrl)?;
118 let mut raw = format!("{scheme}://{host}");
119 if let Some(port) = url.port() {
120 raw.push(':');
121 raw.push_str(&port.to_string());
122 }
123 raw.push_str("/client/v3/push");
124 Ok(raw)
125}
126
127fn worker_main(
128 url: String,
129 mut tokens: Vec<String>,
130 inbox: Arc<Mutex<Vec<Incoming>>>,
131 ready: Arc<AtomicBool>,
132 failed: Arc<Mutex<Option<String>>>,
133 stop: Arc<AtomicBool>,
134) {
135 let result = run_socket(&url, &mut tokens, &inbox, &ready, &stop);
136 for token in &mut tokens {
137 zeroize::Zeroize::zeroize(token);
138 }
139 if let Err(err) = result {
140 if !stop.load(Ordering::Acquire) {
141 *failed.lock().unwrap_or_else(|err| err.into_inner()) = Some(err);
142 }
143 }
144}
145
146fn run_socket(
147 url: &str,
148 tokens: &mut [String],
149 inbox: &Mutex<Vec<Incoming>>,
150 ready: &AtomicBool,
151 stop: &AtomicBool,
152) -> Result<(), String> {
153 let mut request = tungstenite::client::IntoClientRequest::into_client_request(url)
156 .map_err(|err| clip_public(err.to_string()))?;
157 if let Some(first) = tokens.first() {
158 if let Ok(value) = tungstenite::http::HeaderValue::from_str(&format!("Bearer {first}")) {
159 request.headers_mut().insert(tungstenite::http::header::AUTHORIZATION, value);
160 }
161 }
162 let (mut socket, _response) =
163 tungstenite::connect(request).map_err(|err| clip_public(err.to_string()))?;
164 match socket.get_mut() {
165 MaybeTlsStream::Plain(tcp) => {
166 let _ = tcp.set_read_timeout(Some(Duration::from_millis(200)));
167 }
168 MaybeTlsStream::Rustls(tls) => {
169 let _ = tls.sock.set_read_timeout(Some(Duration::from_millis(200)));
170 }
171 _ => {}
172 }
173 let register = serde_json::json!({ "type": "register", "tokens": tokens }).to_string();
174 for token in tokens.iter_mut() {
175 zeroize::Zeroize::zeroize(token);
176 }
177 socket
178 .send(Message::text(register))
179 .map_err(|err| clip_public(err.to_string()))?;
180 loop {
181 if stop.load(Ordering::Acquire) {
182 return Ok(());
183 }
184 let message = match socket.read() {
185 Ok(message) => message,
186 Err(tungstenite::Error::Io(err))
187 if matches!(err.kind(), ErrorKind::TimedOut | ErrorKind::WouldBlock) =>
188 {
189 continue;
190 }
191 Err(err) => return Err(clip_public(err.to_string())),
192 };
193 let text = match message {
194 Message::Text(text) => text,
195 Message::Close(_) => return Err("push socket closed".into()),
196 _ => continue,
197 };
198 let Ok(value) = serde_json::from_str::<serde_json::Value>(text.as_str()) else {
199 continue;
200 };
201 match value.get("type").and_then(|item| item.as_str()) {
202 Some("registered") => ready.store(true, Ordering::Release),
203 Some("event") => {
204 if let Some(envelope_id) = value.get("envelope_id").and_then(|item| item.as_str()) {
205 let ack = serde_json::json!({ "type": "ack", "envelope_id": envelope_id })
206 .to_string();
207 socket
208 .send(Message::text(ack))
209 .map_err(|err| clip_public(err.to_string()))?;
210 }
211 if let Some(incoming) = parse_event(&value) {
212 inbox
213 .lock()
214 .unwrap_or_else(|err| err.into_inner())
215 .push(incoming);
216 }
217 }
218 _ => {}
219 }
220 }
221}
222
223fn parse_event(value: &serde_json::Value) -> Option<Incoming> {
224 if value.get("access_token").is_some() || value.get("bearer").is_some() {
225 return None;
226 }
227 let event = value.get("event")?;
228 if event.get("access_token").is_some() || event.get("bearer").is_some() {
229 return None;
230 }
231 let room = event.get("room")?.as_str()?.to_string();
232 let sender = event.get("sender")?.as_str()?.to_string();
233 let event_id = event.get("event_id")?.as_str()?.to_string();
234 let recipient = event.get("recipient")?.as_str()?.to_string();
235 let wire_type = event
236 .get("wire_type")
237 .and_then(|item| item.as_str())
238 .unwrap_or("m.room.message")
239 .to_string();
240 let body = String::new();
242 if room.is_empty() || sender.is_empty() || event_id.is_empty() || recipient.is_empty() || wire_type.is_empty()
243 {
244 return None;
245 }
246 Some(Incoming {
247 recipient,
248 event: PushedRoomEvent {
249 room,
250 sender,
251 body,
252 event_id,
253 wire_type,
254 },
255 })
256}
257
258#[cfg(test)]
259mod tests {
260 use super::*;
261
262 #[test]
263 fn https_homeserver_opens_the_push_socket_as_wss() {
264 let url = push_ws_url("https://example.test/ignored").expect("url");
265 assert_eq!(url, "wss://example.test/client/v3/push");
266 let local = push_ws_url("http://127.0.0.1:9").expect("local");
267 assert_eq!(local, "ws://127.0.0.1:9/client/v3/push");
268 }
269
270 #[test]
271 fn encrypted_push_keeps_the_event_id_and_drops_any_body() {
272 let value = serde_json::json!({
273 "type": "event",
274 "v": 1,
275 "envelope_id": "p1",
276 "event": {
277 "room": "!room:example",
278 "sender": "@a:example",
279 "event_id": "$evt",
280 "recipient": "@b:example",
281 "wire_type": "m.room.encrypted",
282 "body": "not-plaintext",
283 }
284 });
285 let incoming = parse_event(&value).expect("parsed");
286 assert_eq!(incoming.recipient, "@b:example");
287 assert_eq!(incoming.event.event_id, "$evt");
288 assert_eq!(incoming.event.wire_type, "m.room.encrypted");
289 assert!(incoming.event.body.is_empty());
290 assert_eq!(incoming.event.room, "!room:example");
291 assert_eq!(incoming.event.sender, "@a:example");
292 }
293
294 #[test]
295 fn v1_plaintext_wire_type_parses_without_body() {
296 let value = serde_json::json!({
297 "type": "event",
298 "v": 1,
299 "envelope_id": "p2",
300 "event": {
301 "room": "!room:example",
302 "sender": "@a:example",
303 "event_id": "$evt2",
304 "recipient": "@b:example",
305 "wire_type": "m.room.message",
306 }
307 });
308 let incoming = parse_event(&value).expect("v1 without body must parse");
309 assert!(incoming.event.body.is_empty());
310 assert_eq!(incoming.event.wire_type, "m.room.message");
311 assert_eq!(incoming.event.event_id, "$evt2");
312 }
313}