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 first = match self.exec().perform(&token, request) {
63            Ok(r) => r,
64            Err(err) if HttpExec::is_stale_pool_error(&err) => {
65                // Dead keep-alive / tunnel flap: drop the pool and retry once.
66                let _ = self.with_exec(|e| e.reopen());
67                *self.push.lock().unwrap_or_else(|e| e.into_inner()) = None;
68                self.exec().perform(&token, request)?
69            }
70            Err(err) => return Err(err),
71        };
72        if !unknown_token(&first) {
73            return Ok(first);
74        }
75        let refresher = self.refresher.lock().unwrap_or_else(|e| e.into_inner()).clone();
76        let Some(refresh) = refresher else { return Ok(first) };
77        match refresh() {
78            Ok(fresh) => {
79                self.set_token(&fresh);
80                self.exec().perform(&fresh, request)
81            }
82            Err(_) => Ok(first),
83        }
84    }
85
86    /// Who this token is: `(user_id, device_id)`.
87    pub fn whoami(&self) -> Result<(String, String)> {
88        let token = self.token()?;
89        let (st, v) = self.exec().get_json_as(&token, "/_matrix/client/v3/account/whoami")?;
90        if st != 200 {
91            return Err(AgentError::Protocol(format!("whoami: status {st}")));
92        }
93        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}")));
94        Ok((get("user_id")?, get("device_id")?))
95    }
96
97    /// Blocks until the server pushes news for this session or `timeout` passes. A server without
98    /// a push channel answers `Idle` at once; the caller's sync long poll is then the wait.
99    pub fn wait_for_news(&self, timeout: Duration) -> Result<News> {
100        {
101            let mut slot = self.push.lock().unwrap_or_else(|e| e.into_inner());
102            // A socket that died is dropped and opened again below (it never revives by itself).
103            if slot.as_ref().is_some_and(|link| !link.is_alive()) {
104                *slot = None;
105                *self.push_unavailable.lock().unwrap_or_else(|e| e.into_inner()) = None;
106            }
107            if slot.is_none() {
108                // A failed attempt is not repeated for a while.
109                if self.push_unavailable.lock().unwrap_or_else(|e| e.into_inner()).is_some_and(|t| t.elapsed() < Duration::from_secs(60)) {
110                    return Ok(News::Idle);
111                }
112                let token = self.token()?;
113                let (base, keep) = {
114                    let e = self.exec.read().unwrap_or_else(|e| e.into_inner());
115                    (e.base().as_str().to_string(), e.keep_prefix())
116                };
117                match PushLink::open(&base, keep, vec![token.to_string()], true) {
118                    Ok(link) => *slot = Some(link),
119                    Err(_) => {
120                        *self.push_unavailable.lock().unwrap_or_else(|e| e.into_inner()) = Some(Instant::now());
121                        return Ok(News::Idle);
122                    }
123                }
124            }
125        }
126        let deadline = Instant::now() + timeout;
127        loop {
128            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();
129            if !got.is_empty() {
130                return Ok(News::Pushed(got));
131            }
132            if Instant::now() >= deadline {
133                return Ok(News::Idle);
134            }
135            std::thread::sleep(Duration::from_millis(40));
136        }
137    }
138}
m4a_agent/backend/live.rs

m4a_agent/backend/
live.rs
