Skip to main content

m4a_agent/backend/
live.rs

1//! What every tier shares once it has a token: executing the engine's requests with the bearer,
2//! one retry after a refresh when the server forgot the token, and the push socket as the news
3//! channel.
4
5use std::sync::{Arc, Mutex, RwLock};
6use std::time::{Duration, Instant};
7
8use mail4agent_messenger::{HttpResponseDescriptor, OutgoingRequest};
9use zeroize::Zeroizing;
10
11use crate::engine::{HttpExec, News, PushLink};
12use crate::error::{AgentError, Result};
13
14/// Logs in again (by signature) and returns the new token.
15pub type Refresher = Arc<dyn Fn() -> Result<String> + Send + Sync>;
16
17pub struct Live {
18    exec: RwLock<HttpExec>,
19    token: Mutex<Option<Zeroizing<String>>>,
20    push: Mutex<Option<PushLink>>,
21    push_unavailable: Mutex<Option<Instant>>,
22    refresher: Mutex<Option<Refresher>>,
23}
24
25fn unknown_token(r: &HttpResponseDescriptor) -> bool {
26    r.status == 401 && serde_json::from_slice::<serde_json::Value>(&r.body).ok().and_then(|v| v.get("errcode").and_then(|c| c.as_str()).map(|c| c == "M_UNKNOWN_TOKEN")).unwrap_or(false)
27}
28
29impl Live {
30    pub fn new(exec: HttpExec) -> Self {
31        Self { exec: RwLock::new(exec), token: Mutex::new(None), push: Mutex::new(None), push_unavailable: Mutex::new(None), refresher: Mutex::new(None) }
32    }
33
34    pub fn with_exec<T>(&self, f: impl FnOnce(&mut HttpExec) -> T) -> T {
35        f(&mut self.exec.write().unwrap_or_else(|e| e.into_inner()))
36    }
37
38    pub fn exec(&self) -> HttpExec {
39        self.exec.read().unwrap_or_else(|e| e.into_inner()).clone()
40    }
41
42    pub fn keep_prefix(&self) -> bool {
43        self.exec.read().unwrap_or_else(|e| e.into_inner()).keep_prefix()
44    }
45
46    pub fn set_token(&self, token: &str) {
47        *self.token.lock().unwrap_or_else(|e| e.into_inner()) = Some(Zeroizing::new(token.to_string()));
48        // A push socket registered with the old token is of no use any more.
49        *self.push.lock().unwrap_or_else(|e| e.into_inner()) = None;
50    }
51
52    pub fn set_refresher(&self, r: Refresher) {
53        *self.refresher.lock().unwrap_or_else(|e| e.into_inner()) = Some(r);
54    }
55
56    fn token(&self) -> Result<Zeroizing<String>> {
57        self.token.lock().unwrap_or_else(|e| e.into_inner()).clone().ok_or_else(|| AgentError::Identity("no session yet: ensure_session first".into()))
58    }
59
60    pub fn execute(&self, request: &OutgoingRequest) -> Result<HttpResponseDescriptor> {
61        let token = self.token()?;
62        let exec = self.exec();
63        let first = exec.perform(&token, request)?;
64        if !unknown_token(&first) {
65            return Ok(first);
66        }
67        let refresher = self.refresher.lock().unwrap_or_else(|e| e.into_inner()).clone();
68        let Some(refresh) = refresher else { return Ok(first) };
69        match refresh() {
70            Ok(fresh) => {
71                self.set_token(&fresh);
72                exec.perform(&fresh, request)
73            }
74            Err(_) => Ok(first),
75        }
76    }
77
78    /// Who this token is: `(user_id, device_id)`.
79    pub fn whoami(&self) -> Result<(String, String)> {
80        let token = self.token()?;
81        let (st, v) = self.exec().get_json_as(&token, "/_matrix/client/v3/account/whoami")?;
82        if st != 200 {
83            return Err(AgentError::Protocol(format!("whoami: status {st}")));
84        }
85        let get = |k: &str| v.get(k).and_then(|x| x.as_str()).map(str::to_string).ok_or_else(|| AgentError::Protocol(format!("whoami without {k}")));
86        Ok((get("user_id")?, get("device_id")?))
87    }
88
89    /// Blocks until the server pushes news for this session or `timeout` passes. A server without
90    /// a push channel answers `Idle` at once; the caller's sync long poll is then the wait.
91    pub fn wait_for_news(&self, timeout: Duration) -> Result<News> {
92        {
93            let mut slot = self.push.lock().unwrap_or_else(|e| e.into_inner());
94            // A socket that died is dropped and opened again below (it never revives by itself).
95            if slot.as_ref().is_some_and(|link| !link.is_alive()) {
96                *slot = None;
97                *self.push_unavailable.lock().unwrap_or_else(|e| e.into_inner()) = None;
98            }
99            if slot.is_none() {
100                // A failed attempt is not repeated for a while.
101                if self.push_unavailable.lock().unwrap_or_else(|e| e.into_inner()).is_some_and(|t| t.elapsed() < Duration::from_secs(60)) {
102                    return Ok(News::Idle);
103                }
104                let token = self.token()?;
105                let (base, keep) = {
106                    let e = self.exec.read().unwrap_or_else(|e| e.into_inner());
107                    (e.base().as_str().to_string(), e.keep_prefix())
108                };
109                match PushLink::open(&base, keep, vec![token.to_string()], true) {
110                    Ok(link) => *slot = Some(link),
111                    Err(_) => {
112                        *self.push_unavailable.lock().unwrap_or_else(|e| e.into_inner()) = Some(Instant::now());
113                        return Ok(News::Idle);
114                    }
115                }
116            }
117        }
118        let deadline = Instant::now() + timeout;
119        loop {
120            let got: Vec<_> = self.push.lock().unwrap_or_else(|e| e.into_inner()).as_ref().map(|p| p.drain().into_iter().map(|(_, e)| e).collect()).unwrap_or_default();
121            if !got.is_empty() {
122                return Ok(News::Pushed(got));
123            }
124            if Instant::now() >= deadline {
125                return Ok(News::Idle);
126            }
127            std::thread::sleep(Duration::from_millis(40));
128        }
129    }
130}