Skip to main content

m4a_agent/engine/
push.rs

1//! The client's push 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::net::{TcpStream, ToSocketAddrs};
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::sync::{Arc, Mutex};
12use std::thread::{self, JoinHandle};
13use std::time::{Duration, Instant};
14
15use tungstenite::{stream::MaybeTlsStream, Message};
16
17use super::http::clip_public;
18use crate::error::AgentError;
19
20/// One room event the homeserver pushed for a single session.
21#[derive(Clone, Debug, PartialEq, Eq)]
22pub struct PushedRoomEvent {
23    /// Room id.
24    pub room: String,
25    /// Sender mxid.
26    pub sender: String,
27    /// Always empty on push v1 (server never sends plaintext). Kept so
28    /// call sites compiling against this struct stay stable; wake uses
29    /// `/sync` + decrypt, not this field.
30    pub body: String,
31    /// Matrix event id.
32    pub event_id: String,
33    /// `m.room.message` or `m.room.encrypted`.
34    pub wire_type: String,
35}
36
37/// Every wait on the push socket is bounded. A network flap (or a local proxy that accepts the
38/// TCP connection and then never answers) used to leave the connect, the TLS/WebSocket handshake
39/// or the read loop blocked forever, and the listener stopped waking sessions until the socket
40/// was killed by hand.
41#[derive(Clone, Copy, Debug, PartialEq, Eq)]
42pub struct PushTimeouts {
43    /// TCP connect, per resolved address.
44    pub connect: Duration,
45    /// TLS + WebSocket upgrade, and every write after it.
46    pub handshake: Duration,
47    /// How long [`PushLink::open`] waits for the server's `registered`.
48    pub register: Duration,
49    /// A WebSocket ping goes out this often; its pong (or any frame) proves the path is alive.
50    pub ping_every: Duration,
51    /// No frame at all for this long ends the socket: [`PushLink::is_alive`] turns false and the
52    /// owner opens a new one. Several pings fit in it, so one lost pong is not a reconnect.
53    pub silent_after: Duration,
54    /// TCP keepalive idle time (the probes after it are a few seconds apart).
55    pub keepalive: Duration,
56}
57
58impl Default for PushTimeouts {
59    fn default() -> Self {
60        Self {
61            connect: Duration::from_secs(10),
62            handshake: Duration::from_secs(15),
63            register: Duration::from_secs(5),
64            ping_every: Duration::from_secs(20),
65            silent_after: Duration::from_secs(70),
66            keepalive: Duration::from_secs(15),
67        }
68    }
69}
70
71struct Incoming {
72    recipient: String,
73    event: PushedRoomEvent,
74}
75
76pub struct PushLink {
77    inbox: Arc<Mutex<Vec<Incoming>>>,
78    stop: Arc<AtomicBool>,
79    workers: Vec<JoinHandle<()>>,
80    /// Why each worker ended, when it ended on an error (clipped public text).
81    faileds: Vec<Arc<Mutex<Option<String>>>>,
82}
83
84impl Drop for PushLink {
85    /// Stops the workers. A worker notices within a read tick; one still inside a connect or a
86    /// handshake ends at its own deadline. The caller is never held past a short grace: a
87    /// worker that has not ended by then is left to finish on its own (it holds no lock).
88    fn drop(&mut self) {
89        self.stop.store(true, Ordering::Release);
90        let grace = Instant::now() + Duration::from_millis(600);
91        while Instant::now() < grace && !self.workers.iter().all(|w| w.is_finished()) {
92            thread::sleep(Duration::from_millis(10));
93        }
94        for worker in self.workers.drain(..) {
95            if worker.is_finished() {
96                let _ = worker.join();
97            }
98        }
99    }
100}
101
102impl PushLink {
103    /// One socket for all `tokens`, or with `separate` (product mode: a handshake
104    /// is vouched for one identity) one socket per token sharing one inbox.
105    pub fn open(base_url: &str, keep_prefix: bool, tokens: Vec<String>, separate: bool) -> Result<Self, AgentError> {
106        Self::open_with(base_url, keep_prefix, tokens, separate, PushTimeouts::default())
107    }
108
109    /// [`Self::open`] with explicit [`PushTimeouts`].
110    pub fn open_with(base_url: &str, keep_prefix: bool, tokens: Vec<String>, separate: bool, timeouts: PushTimeouts) -> Result<Self, AgentError> {
111        if tokens.is_empty() {
112            return Err(AgentError::Transport("push socket has no session".into()));
113        }
114        let url = push_ws_url(base_url, keep_prefix)?;
115        let inbox = Arc::new(Mutex::new(Vec::new()));
116        let stop = Arc::new(AtomicBool::new(false));
117        let groups: Vec<Vec<String>> = if separate { tokens.into_iter().map(|t| vec![t]).collect() } else { vec![tokens] };
118        let mut link = Self { inbox: Arc::clone(&inbox), stop: Arc::clone(&stop), workers: Vec::new(), faileds: Vec::new() };
119        let mut readies = Vec::new();
120        for group in groups {
121            let ready = Arc::new(AtomicBool::new(false));
122            let failed = Arc::new(Mutex::new(None));
123            let worker = thread::spawn({
124                let (url, inbox, ready, failed, stop) = (url.clone(), Arc::clone(&inbox), Arc::clone(&ready), Arc::clone(&failed), Arc::clone(&stop));
125                move || worker_main(url, group, inbox, ready, failed, stop, timeouts)
126            });
127            link.workers.push(worker);
128            readies.push(ready);
129            link.faileds.push(failed);
130        }
131        let start = Instant::now();
132        loop {
133            if readies.iter().all(|r| r.load(Ordering::Acquire)) {
134                return Ok(link);
135            }
136            if let Some(err) = link.last_error() {
137                return Err(AgentError::Transport(err)); // Drop stops the workers
138            }
139            if start.elapsed() > timeouts.register {
140                return Err(AgentError::Transport("push socket did not register in time".into()));
141            }
142            thread::sleep(Duration::from_millis(20));
143        }
144    }
145
146    /// `false` once a socket worker has ended (the server closed it, the network dropped, the
147    /// server restarted). A dead link never revives by itself: the owner opens a new one.
148    pub fn is_alive(&self) -> bool {
149        self.workers.iter().all(|worker| !worker.is_finished())
150    }
151
152    /// Why a worker ended, when it ended on an error: connect/handshake timeout, a silent
153    /// socket, the server closing it. Public text, no URL or token.
154    pub fn last_error(&self) -> Option<String> {
155        self.faileds.iter().find_map(|f| f.lock().unwrap_or_else(|err| err.into_inner()).clone())
156    }
157
158    pub fn drain(&self) -> Vec<(String, PushedRoomEvent)> {
159        self.inbox
160            .lock()
161            .unwrap_or_else(|err| err.into_inner())
162            .drain(..)
163            .map(|incoming| (incoming.recipient, incoming.event))
164            .collect()
165    }
166}
167
168fn push_ws_url(base: &str, keep_prefix: bool) -> Result<String, AgentError> {
169    let url = reqwest::Url::parse(base).map_err(|_| AgentError::Protocol("not a valid server URL".into()))?;
170    let scheme = match url.scheme() {
171        "http" => "ws",
172        "https" => "wss",
173        _ => return Err(AgentError::Protocol("not a valid server URL".into())),
174    };
175    let host = url
176        .host_str()
177        .filter(|host| !host.is_empty())
178        .ok_or(AgentError::Protocol("not a valid server URL".into()))?;
179    let mut raw = format!("{scheme}://{host}");
180    if let Some(port) = url.port() {
181        raw.push(':');
182        raw.push_str(&port.to_string());
183    }
184    raw.push_str(if keep_prefix { "/_matrix/client/v3/push" } else { "/client/v3/push" });
185    Ok(raw)
186}
187
188fn worker_main(
189    url: String,
190    mut tokens: Vec<String>,
191    inbox: Arc<Mutex<Vec<Incoming>>>,
192    ready: Arc<AtomicBool>,
193    failed: Arc<Mutex<Option<String>>>,
194    stop: Arc<AtomicBool>,
195    timeouts: PushTimeouts,
196) {
197    let result = run_socket(&url, &mut tokens, &inbox, &ready, &stop, &timeouts);
198    for token in &mut tokens {
199        zeroize::Zeroize::zeroize(token);
200    }
201    if let Err(err) = result {
202        if !stop.load(Ordering::Acquire) {
203            *failed.lock().unwrap_or_else(|err| err.into_inner()) = Some(err);
204        }
205    }
206}
207
208fn run_socket(
209    url: &str,
210    tokens: &mut [String],
211    inbox: &Mutex<Vec<Incoming>>,
212    ready: &AtomicBool,
213    stop: &AtomicBool,
214    timeouts: &PushTimeouts,
215) -> Result<(), String> {
216    // The first token also rides the handshake as a bearer: a product proxy
217    // authenticates the handshake with it and signs it; a plain server ignores it.
218    let mut request = tungstenite::client::IntoClientRequest::into_client_request(url)
219        .map_err(|err| clip_public(err.to_string()))?;
220    if let Some(first) = tokens.first() {
221        if let Ok(value) = tungstenite::http::HeaderValue::from_str(&format!("Bearer {first}")) {
222            request.headers_mut().insert(tungstenite::http::header::AUTHORIZATION, value);
223        }
224    }
225    let tcp = connect_tcp(request.uri(), timeouts)?;
226    // The handshake reads and writes under the handshake deadline: a peer that accepted the
227    // connection and then went quiet fails the handshake instead of holding it forever.
228    tcp.set_read_timeout(Some(timeouts.handshake)).map_err(|err| clip_public(err.to_string()))?;
229    tcp.set_write_timeout(Some(timeouts.handshake)).map_err(|err| clip_public(err.to_string()))?;
230    let (mut socket, _response) = match tungstenite::client_tls_with_config(request, tcp, None, None) {
231        Ok(done) => done,
232        Err(tungstenite::HandshakeError::Interrupted(_)) => return Err("push handshake timed out".into()),
233        Err(tungstenite::HandshakeError::Failure(tungstenite::Error::Io(err)))
234            if matches!(err.kind(), ErrorKind::TimedOut | ErrorKind::WouldBlock) =>
235        {
236            return Err("push handshake timed out".into())
237        }
238        Err(tungstenite::HandshakeError::Failure(err)) => return Err(clip_public(err.to_string())),
239    };
240    let tcp = match socket.get_mut() {
241        MaybeTlsStream::Plain(tcp) => Some(&*tcp),
242        MaybeTlsStream::Rustls(tls) => Some(&tls.sock),
243        _ => None,
244    };
245    if let Some(tcp) = tcp {
246        let _ = tcp.set_read_timeout(Some(Duration::from_millis(200)));
247        let _ = tcp.set_write_timeout(Some(timeouts.handshake));
248    }
249    let register = serde_json::json!({ "type": "register", "tokens": tokens }).to_string();
250    for token in tokens.iter_mut() {
251        zeroize::Zeroize::zeroize(token);
252    }
253    socket
254        .send(Message::text(register))
255        .map_err(|err| clip_public(err.to_string()))?;
256    let mut last_frame = Instant::now();
257    let mut last_ping = Instant::now();
258    loop {
259        if stop.load(Ordering::Acquire) {
260            return Ok(());
261        }
262        // A half-open connection never errors on read: only the silence shows it.
263        if last_frame.elapsed() >= timeouts.silent_after {
264            return Err(format!("push socket silent for {}s", last_frame.elapsed().as_secs()));
265        }
266        if last_ping.elapsed() >= timeouts.ping_every {
267            last_ping = Instant::now();
268            match socket.send(Message::Ping(Vec::new().into())) {
269                Ok(()) => {}
270                Err(tungstenite::Error::Io(err)) if matches!(err.kind(), ErrorKind::TimedOut | ErrorKind::WouldBlock) => {
271                    return Err("push socket write timed out".into())
272                }
273                Err(err) => return Err(clip_public(err.to_string())),
274            }
275        }
276        let message = match socket.read() {
277            Ok(message) => message,
278            Err(tungstenite::Error::Io(err))
279                if matches!(err.kind(), ErrorKind::TimedOut | ErrorKind::WouldBlock) =>
280            {
281                continue;
282            }
283            Err(err) => return Err(clip_public(err.to_string())),
284        };
285        last_frame = Instant::now();
286        let text = match message {
287            Message::Text(text) => text,
288            Message::Close(_) => return Err("push socket closed".into()),
289            _ => continue,
290        };
291        let Ok(value) = serde_json::from_str::<serde_json::Value>(text.as_str()) else {
292            continue;
293        };
294        match value.get("type").and_then(|item| item.as_str()) {
295            Some("registered") => ready.store(true, Ordering::Release),
296            Some("event") => {
297                if let Some(envelope_id) = value.get("envelope_id").and_then(|item| item.as_str()) {
298                    let ack = serde_json::json!({ "type": "ack", "envelope_id": envelope_id })
299                        .to_string();
300                    socket
301                        .send(Message::text(ack))
302                        .map_err(|err| clip_public(err.to_string()))?;
303                }
304                if let Some(incoming) = parse_event(&value) {
305                    inbox
306                        .lock()
307                        .unwrap_or_else(|err| err.into_inner())
308                        .push(incoming);
309                }
310            }
311            _ => {}
312        }
313    }
314}
315
316/// TCP to the push URL's host with a deadline per address, keepalive on, and (Linux) a user
317/// timeout so unacknowledged writes end the connection instead of queueing forever.
318fn connect_tcp(uri: &tungstenite::http::Uri, timeouts: &PushTimeouts) -> Result<TcpStream, String> {
319    let host = uri.host().ok_or("push URL has no host")?;
320    let host = host.trim_start_matches('[').trim_end_matches(']');
321    let port = uri.port_u16().unwrap_or(if uri.scheme_str() == Some("wss") { 443 } else { 80 });
322    let addrs = (host, port).to_socket_addrs().map_err(|err| format!("push host not resolved: {}", err.kind()))?;
323    let mut last = String::from("push host has no address");
324    for addr in addrs {
325        match TcpStream::connect_timeout(&addr, timeouts.connect) {
326            Ok(tcp) => {
327                let _ = tcp.set_nodelay(true);
328                tune_keepalive(&tcp, timeouts);
329                return Ok(tcp);
330            }
331            Err(err) if err.kind() == ErrorKind::TimedOut => last = "push connect timed out".into(),
332            Err(err) => last = format!("push connect failed: {err}"),
333        }
334    }
335    Err(last)
336}
337
338fn tune_keepalive(tcp: &TcpStream, timeouts: &PushTimeouts) {
339    let sock = socket2::SockRef::from(tcp);
340    let keepalive = socket2::TcpKeepalive::new().with_time(timeouts.keepalive);
341    #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
342    let keepalive = keepalive.with_interval(Duration::from_secs(5));
343    #[cfg(any(target_os = "linux", target_os = "macos"))]
344    let keepalive = keepalive.with_retries(3);
345    let _ = sock.set_tcp_keepalive(&keepalive);
346    #[cfg(target_os = "linux")]
347    let _ = sock.set_tcp_user_timeout(Some(timeouts.keepalive * 2));
348}
349
350fn parse_event(value: &serde_json::Value) -> Option<Incoming> {
351    if value.get("access_token").is_some() || value.get("bearer").is_some() {
352        return None;
353    }
354    let event = value.get("event")?;
355    if event.get("access_token").is_some() || event.get("bearer").is_some() {
356        return None;
357    }
358    let room = event.get("room")?.as_str()?.to_string();
359    let sender = event.get("sender")?.as_str()?.to_string();
360    let event_id = event.get("event_id")?.as_str()?.to_string();
361    let recipient = event.get("recipient")?.as_str()?.to_string();
362    let wire_type = event
363        .get("wire_type")
364        .and_then(|item| item.as_str())
365        .unwrap_or("m.room.message")
366        .to_string();
367    // Push v1: never trust/require body. Always empty locally.
368    let body = String::new();
369    if room.is_empty() || sender.is_empty() || event_id.is_empty() || recipient.is_empty() || wire_type.is_empty()
370    {
371        return None;
372    }
373    Some(Incoming {
374        recipient,
375        event: PushedRoomEvent {
376            room,
377            sender,
378            body,
379            event_id,
380            wire_type,
381        },
382    })
383}
384
385#[cfg(test)]
386mod tests {
387    use super::*;
388
389    #[test]
390    fn a_link_whose_worker_ended_is_not_alive() {
391        let alive = PushLink { inbox: Arc::default(), stop: Arc::default(), workers: vec![thread::spawn(|| thread::sleep(Duration::from_millis(400)))], faileds: Vec::new() };
392        assert!(alive.is_alive());
393        let dead = PushLink { inbox: Arc::default(), stop: Arc::default(), workers: vec![thread::spawn(|| {})], faileds: Vec::new() };
394        thread::sleep(Duration::from_millis(100));
395        assert!(!dead.is_alive());
396        thread::sleep(Duration::from_millis(500));
397        assert!(!alive.is_alive());
398    }
399
400    #[test]
401    fn https_homeserver_opens_the_push_socket_as_wss() {
402        let url = push_ws_url("https://example.test/ignored", false).expect("url");
403        assert_eq!(url, "wss://example.test/client/v3/push");
404        let local = push_ws_url("http://127.0.0.1:9", false).expect("local");
405        assert_eq!(local, "ws://127.0.0.1:9/client/v3/push");
406    }
407
408    #[test]
409    fn encrypted_push_keeps_the_event_id_and_drops_any_body() {
410        let value = serde_json::json!({
411            "type": "event",
412            "v": 1,
413            "envelope_id": "p1",
414            "event": {
415                "room": "!room:example",
416                "sender": "@a:example",
417                "event_id": "$evt",
418                "recipient": "@b:example",
419                "wire_type": "m.room.encrypted",
420                "body": "not-plaintext",
421            }
422        });
423        let incoming = parse_event(&value).expect("parsed");
424        assert_eq!(incoming.recipient, "@b:example");
425        assert_eq!(incoming.event.event_id, "$evt");
426        assert_eq!(incoming.event.wire_type, "m.room.encrypted");
427        assert!(incoming.event.body.is_empty());
428        assert_eq!(incoming.event.room, "!room:example");
429        assert_eq!(incoming.event.sender, "@a:example");
430    }
431
432    #[test]
433    fn v1_plaintext_wire_type_parses_without_body() {
434        let value = serde_json::json!({
435            "type": "event",
436            "v": 1,
437            "envelope_id": "p2",
438            "event": {
439                "room": "!room:example",
440                "sender": "@a:example",
441                "event_id": "$evt2",
442                "recipient": "@b:example",
443                "wire_type": "m.room.message",
444            }
445        });
446        let incoming = parse_event(&value).expect("v1 without body must parse");
447        assert!(incoming.event.body.is_empty());
448        assert_eq!(incoming.event.wire_type, "m.room.message");
449        assert_eq!(incoming.event.event_id, "$evt2");
450    }
451}