Skip to main content

mail4agent_messenger_shell/
push.rs

1//! The machine client's one socket. The client opens it against the
2//! homeserver it already has. Push contract **v1**: metadata only
3//! (`room`, `sender`, `event_id`, `recipient`, `wire_type`) — never a
4//! plaintext `body`. The ack goes out before the event is handed to that
5//! session. This path does not POST. The client drives `/sync` and
6//! decrypts after the push.
7
8use 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/// One room event the homeserver pushed for a single session.
19#[derive(Clone, Debug, PartialEq, Eq)]
20pub struct PushedRoomEvent {
21    /// Room id.
22    pub room: String,
23    /// Sender mxid.
24    pub sender: String,
25    /// Always empty on push v1 (server never sends plaintext). Kept so
26    /// call sites compiling against this struct stay stable; wake uses
27    /// `/sync` + decrypt, not this field.
28    pub body: String,
29    /// Matrix event id.
30    pub event_id: String,
31    /// `m.room.message` or `m.room.encrypted`.
32    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    /// One socket for all `tokens`, or with `separate` (product mode: a handshake
57    /// is vouched for one identity) one socket per token sharing one inbox.
58    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)); // Drop stops and joins the workers
88                }
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    // The first token also rides the handshake as a bearer: a product proxy
154    // authenticates the handshake with it and signs it; a plain server ignores it.
155    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    // Push v1: never trust/require body. Always empty locally.
241    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}