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::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 super::http::clip_public;
17use crate::error::AgentError;
18
19/// One room event the homeserver pushed for a single session.
20#[derive(Clone, Debug, PartialEq, Eq)]
21pub struct PushedRoomEvent {
22    /// Room id.
23    pub room: String,
24    /// Sender mxid.
25    pub sender: String,
26    /// Always empty on push v1 (server never sends plaintext). Kept so
27    /// call sites compiling against this struct stay stable; wake uses
28    /// `/sync` + decrypt, not this field.
29    pub body: String,
30    /// Matrix event id.
31    pub event_id: String,
32    /// `m.room.message` or `m.room.encrypted`.
33    pub wire_type: String,
34}
35
36struct Incoming {
37    recipient: String,
38    event: PushedRoomEvent,
39}
40
41pub struct PushLink {
42    inbox: Arc<Mutex<Vec<Incoming>>>,
43    stop: Arc<AtomicBool>,
44    workers: Vec<JoinHandle<()>>,
45}
46
47impl Drop for PushLink {
48    fn drop(&mut self) {
49        self.stop.store(true, Ordering::Release);
50        for worker in self.workers.drain(..) {
51            let _ = worker.join();
52        }
53    }
54}
55
56impl PushLink {
57    /// One socket for all `tokens`, or with `separate` (product mode: a handshake
58    /// is vouched for one identity) one socket per token sharing one inbox.
59    pub fn open(base_url: &str, keep_prefix: bool, tokens: Vec<String>, separate: bool) -> Result<Self, AgentError> {
60        if tokens.is_empty() {
61            return Err(AgentError::Transport("push socket has no session".into()));
62        }
63        let url = push_ws_url(base_url, keep_prefix)?;
64        let inbox = Arc::new(Mutex::new(Vec::new()));
65        let stop = Arc::new(AtomicBool::new(false));
66        let groups: Vec<Vec<String>> = if separate { tokens.into_iter().map(|t| vec![t]).collect() } else { vec![tokens] };
67        let mut link = Self { inbox: Arc::clone(&inbox), stop: Arc::clone(&stop), workers: Vec::new() };
68        let mut readies = Vec::new();
69        let mut faileds = Vec::new();
70        for group in groups {
71            let ready = Arc::new(AtomicBool::new(false));
72            let failed = Arc::new(Mutex::new(None));
73            let worker = thread::spawn({
74                let (url, inbox, ready, failed, stop) = (url.clone(), Arc::clone(&inbox), Arc::clone(&ready), Arc::clone(&failed), Arc::clone(&stop));
75                move || worker_main(url, group, inbox, ready, failed, stop)
76            });
77            link.workers.push(worker);
78            readies.push(ready);
79            faileds.push(failed);
80        }
81        let start = Instant::now();
82        loop {
83            if readies.iter().all(|r| r.load(Ordering::Acquire)) {
84                return Ok(link);
85            }
86            for failed in &faileds {
87                if let Some(err) = failed.lock().unwrap_or_else(|err| err.into_inner()).clone() {
88                    return Err(AgentError::Transport(err)); // Drop stops and joins the workers
89                }
90            }
91            if start.elapsed() > Duration::from_secs(5) {
92                return Err(AgentError::Transport("push socket did not register".into()));
93            }
94            thread::sleep(Duration::from_millis(20));
95        }
96    }
97
98    /// `false` once a socket worker has ended (the server closed it, the network dropped, the
99    /// server restarted). A dead link never revives by itself: the owner opens a new one.
100    pub fn is_alive(&self) -> bool {
101        self.workers.iter().all(|worker| !worker.is_finished())
102    }
103
104    pub fn drain(&self) -> Vec<(String, PushedRoomEvent)> {
105        self.inbox
106            .lock()
107            .unwrap_or_else(|err| err.into_inner())
108            .drain(..)
109            .map(|incoming| (incoming.recipient, incoming.event))
110            .collect()
111    }
112}
113
114fn push_ws_url(base: &str, keep_prefix: bool) -> Result<String, AgentError> {
115    let url = reqwest::Url::parse(base).map_err(|_| AgentError::Protocol("not a valid server URL".into()))?;
116    let scheme = match url.scheme() {
117        "http" => "ws",
118        "https" => "wss",
119        _ => return Err(AgentError::Protocol("not a valid server URL".into())),
120    };
121    let host = url
122        .host_str()
123        .filter(|host| !host.is_empty())
124        .ok_or(AgentError::Protocol("not a valid server URL".into()))?;
125    let mut raw = format!("{scheme}://{host}");
126    if let Some(port) = url.port() {
127        raw.push(':');
128        raw.push_str(&port.to_string());
129    }
130    raw.push_str(if keep_prefix { "/_matrix/client/v3/push" } else { "/client/v3/push" });
131    Ok(raw)
132}
133
134fn worker_main(
135    url: String,
136    mut tokens: Vec<String>,
137    inbox: Arc<Mutex<Vec<Incoming>>>,
138    ready: Arc<AtomicBool>,
139    failed: Arc<Mutex<Option<String>>>,
140    stop: Arc<AtomicBool>,
141) {
142    let result = run_socket(&url, &mut tokens, &inbox, &ready, &stop);
143    for token in &mut tokens {
144        zeroize::Zeroize::zeroize(token);
145    }
146    if let Err(err) = result {
147        if !stop.load(Ordering::Acquire) {
148            *failed.lock().unwrap_or_else(|err| err.into_inner()) = Some(err);
149        }
150    }
151}
152
153fn run_socket(
154    url: &str,
155    tokens: &mut [String],
156    inbox: &Mutex<Vec<Incoming>>,
157    ready: &AtomicBool,
158    stop: &AtomicBool,
159) -> Result<(), String> {
160    // The first token also rides the handshake as a bearer: a product proxy
161    // authenticates the handshake with it and signs it; a plain server ignores it.
162    let mut request = tungstenite::client::IntoClientRequest::into_client_request(url)
163        .map_err(|err| clip_public(err.to_string()))?;
164    if let Some(first) = tokens.first() {
165        if let Ok(value) = tungstenite::http::HeaderValue::from_str(&format!("Bearer {first}")) {
166            request.headers_mut().insert(tungstenite::http::header::AUTHORIZATION, value);
167        }
168    }
169    let (mut socket, _response) =
170        tungstenite::connect(request).map_err(|err| clip_public(err.to_string()))?;
171    match socket.get_mut() {
172        MaybeTlsStream::Plain(tcp) => {
173            let _ = tcp.set_read_timeout(Some(Duration::from_millis(200)));
174        }
175        MaybeTlsStream::Rustls(tls) => {
176            let _ = tls.sock.set_read_timeout(Some(Duration::from_millis(200)));
177        }
178        _ => {}
179    }
180    let register = serde_json::json!({ "type": "register", "tokens": tokens }).to_string();
181    for token in tokens.iter_mut() {
182        zeroize::Zeroize::zeroize(token);
183    }
184    socket
185        .send(Message::text(register))
186        .map_err(|err| clip_public(err.to_string()))?;
187    loop {
188        if stop.load(Ordering::Acquire) {
189            return Ok(());
190        }
191        let message = match socket.read() {
192            Ok(message) => message,
193            Err(tungstenite::Error::Io(err))
194                if matches!(err.kind(), ErrorKind::TimedOut | ErrorKind::WouldBlock) =>
195            {
196                continue;
197            }
198            Err(err) => return Err(clip_public(err.to_string())),
199        };
200        let text = match message {
201            Message::Text(text) => text,
202            Message::Close(_) => return Err("push socket closed".into()),
203            _ => continue,
204        };
205        let Ok(value) = serde_json::from_str::<serde_json::Value>(text.as_str()) else {
206            continue;
207        };
208        match value.get("type").and_then(|item| item.as_str()) {
209            Some("registered") => ready.store(true, Ordering::Release),
210            Some("event") => {
211                if let Some(envelope_id) = value.get("envelope_id").and_then(|item| item.as_str()) {
212                    let ack = serde_json::json!({ "type": "ack", "envelope_id": envelope_id })
213                        .to_string();
214                    socket
215                        .send(Message::text(ack))
216                        .map_err(|err| clip_public(err.to_string()))?;
217                }
218                if let Some(incoming) = parse_event(&value) {
219                    inbox
220                        .lock()
221                        .unwrap_or_else(|err| err.into_inner())
222                        .push(incoming);
223                }
224            }
225            _ => {}
226        }
227    }
228}
229
230fn parse_event(value: &serde_json::Value) -> Option<Incoming> {
231    if value.get("access_token").is_some() || value.get("bearer").is_some() {
232        return None;
233    }
234    let event = value.get("event")?;
235    if event.get("access_token").is_some() || event.get("bearer").is_some() {
236        return None;
237    }
238    let room = event.get("room")?.as_str()?.to_string();
239    let sender = event.get("sender")?.as_str()?.to_string();
240    let event_id = event.get("event_id")?.as_str()?.to_string();
241    let recipient = event.get("recipient")?.as_str()?.to_string();
242    let wire_type = event
243        .get("wire_type")
244        .and_then(|item| item.as_str())
245        .unwrap_or("m.room.message")
246        .to_string();
247    // Push v1: never trust/require body. Always empty locally.
248    let body = String::new();
249    if room.is_empty() || sender.is_empty() || event_id.is_empty() || recipient.is_empty() || wire_type.is_empty()
250    {
251        return None;
252    }
253    Some(Incoming {
254        recipient,
255        event: PushedRoomEvent {
256            room,
257            sender,
258            body,
259            event_id,
260            wire_type,
261        },
262    })
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    #[test]
270    fn a_link_whose_worker_ended_is_not_alive() {
271        let alive = PushLink { inbox: Arc::default(), stop: Arc::default(), workers: vec![thread::spawn(|| thread::sleep(Duration::from_millis(400)))] };
272        assert!(alive.is_alive());
273        let dead = PushLink { inbox: Arc::default(), stop: Arc::default(), workers: vec![thread::spawn(|| {})] };
274        thread::sleep(Duration::from_millis(100));
275        assert!(!dead.is_alive());
276        thread::sleep(Duration::from_millis(500));
277        assert!(!alive.is_alive());
278    }
279
280    #[test]
281    fn https_homeserver_opens_the_push_socket_as_wss() {
282        let url = push_ws_url("https://example.test/ignored", false).expect("url");
283        assert_eq!(url, "wss://example.test/client/v3/push");
284        let local = push_ws_url("http://127.0.0.1:9", false).expect("local");
285        assert_eq!(local, "ws://127.0.0.1:9/client/v3/push");
286    }
287
288    #[test]
289    fn encrypted_push_keeps_the_event_id_and_drops_any_body() {
290        let value = serde_json::json!({
291            "type": "event",
292            "v": 1,
293            "envelope_id": "p1",
294            "event": {
295                "room": "!room:example",
296                "sender": "@a:example",
297                "event_id": "$evt",
298                "recipient": "@b:example",
299                "wire_type": "m.room.encrypted",
300                "body": "not-plaintext",
301            }
302        });
303        let incoming = parse_event(&value).expect("parsed");
304        assert_eq!(incoming.recipient, "@b:example");
305        assert_eq!(incoming.event.event_id, "$evt");
306        assert_eq!(incoming.event.wire_type, "m.room.encrypted");
307        assert!(incoming.event.body.is_empty());
308        assert_eq!(incoming.event.room, "!room:example");
309        assert_eq!(incoming.event.sender, "@a:example");
310    }
311
312    #[test]
313    fn v1_plaintext_wire_type_parses_without_body() {
314        let value = serde_json::json!({
315            "type": "event",
316            "v": 1,
317            "envelope_id": "p2",
318            "event": {
319                "room": "!room:example",
320                "sender": "@a:example",
321                "event_id": "$evt2",
322                "recipient": "@b:example",
323                "wire_type": "m.room.message",
324            }
325        });
326        let incoming = parse_event(&value).expect("v1 without body must parse");
327        assert!(incoming.event.body.is_empty());
328        assert_eq!(incoming.event.wire_type, "m.room.message");
329        assert_eq!(incoming.event.event_id, "$evt2");
330    }
331}
m4a_agent/engine/push.rs

m4a_agent/engine/
push.rs
