atagen

macltok

public
Star0
Watch0
5c749146259c5ffd2d87a86858afdde32013b6a0
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/engine/session.rsRust132 lines4.4 KB
1
use std::sync::Arc;
2
3
use matrix_sdk::{BoxFuture, Client, SessionChange};
4
use ruma::OwnedDeviceId;
5
use tokio::sync::{
6
    broadcast::{self, error::RecvError},
7
    mpsc, watch,
8
};
9
10
use super::Logouts;
11
use crate::{Error, Event, Termination, transport::EventSender};
12
13
#[cfg(not(target_family = "wasm"))]
14
pub(crate) type SaveRecord = dyn Fn(Client) -> BoxFuture<'static, Result<(), Error>> + Send + Sync;
15
#[cfg(target_family = "wasm")]
16
pub(crate) type SaveRecord = dyn Fn(Client) -> BoxFuture<'static, Result<(), Error>>;
17
18
/// how the engine keeps the saved record current, handed over by `Session::build`
19
pub(crate) struct Persistence {
20
    pub(crate) save: Box<SaveRecord>,
21
    /// what the sdk's own save callback made of each refresh, when the store could write
22
    /// there and then
23
    pub(crate) inline: mpsc::UnboundedReceiver<Result<(), Error>>,
24
}
25
26
impl Persistence {
27
    async fn persist(&mut self, client: &Client) -> Result<(), Error> {
28
        let mut inline = None;
29
        while let Ok(result) = self.inline.try_recv() {
30
            inline = Some(result);
31
        }
32
        if let Some(Ok(())) = inline {
33
            return Ok(());
34
        }
35
        // an inline save for a newer refresh can land while this one writes, so go again
36
        // until what's on disk is what the client holds
37
        loop {
38
            let tokens = client.session_tokens();
39
            (self.save)(client.clone()).await?;
40
            if client.session_tokens() == tokens {
41
                return Ok(());
42
            }
43
        }
44
    }
45
}
46
47
pub(super) enum Notice {
48
    Event(Event),
49
    Rejected(Option<OwnedDeviceId>),
50
}
51
52
/// saves and notices are split, so a full event stream never holds a save back
53
pub(super) async fn session(
54
    client: Client,
55
    mut changes: broadcast::Receiver<SessionChange>,
56
    mut persistence: Option<Persistence>,
57
    notices: mpsc::UnboundedSender<Notice>,
58
    device_id: OwnedDeviceId,
59
    logouts: Arc<watch::Sender<Logouts>>,
60
) {
61
    loop {
62
        let event = match changes.recv().await {
63
            // a missed refresh costs one needless save, a missed rejection comes round again
64
            Ok(SessionChange::TokensRefreshed) | Err(RecvError::Lagged(_)) => {
65
                match persistence.as_mut() {
66
                    None => Event::SessionTokensChanged,
67
                    Some(persistence) => match persistence.persist(&client).await {
68
                        Ok(()) => Event::SessionTokensChanged,
69
                        Err(error) => Event::SessionSaveFailed {
70
                            message: error.message,
71
                        },
72
                    },
73
                }
74
            }
75
            Ok(SessionChange::UnknownToken(data)) => {
76
                let soft_logout = data.soft_logout.then_some(device_id);
77
                let mut held = false;
78
                // our own logout revokes the token, so a rejection meanwhile is its doing
79
                logouts.send_if_modified(|logouts| {
80
                    held = logouts.in_flight > 0;
81
                    if held {
82
                        logouts.rejected = Some(soft_logout.clone());
83
                    }
84
                    held
85
                });
86
                if !held {
87
                    let _ = notices.send(Notice::Rejected(soft_logout));
88
                }
89
                return;
90
            }
91
            Err(RecvError::Closed) => return,
92
        };
93
        if notices.send(Notice::Event(event)).is_err() {
94
            return;
95
        }
96
    }
97
}
98
99
pub(super) async fn session_notices(
100
    mut notices: mpsc::UnboundedReceiver<Notice>,
101
    events: EventSender,
102
    stop: mpsc::UnboundedSender<Termination>,
103
) {
104
    while let Some(notice) = notices.recv().await {
105
        match notice {
106
            Notice::Event(event) => {
107
                if events.send(event).await.is_err() {
108
                    return;
109
                }
110
            }
111
            Notice::Rejected(soft_logout) => {
112
                return reject_session(&events, &stop, soft_logout).await;
113
            }
114
        }
115
    }
116
}
117
118
pub(super) async fn reject_session(
119
    events: &EventSender,
120
    stop: &mpsc::UnboundedSender<Termination>,
121
    soft_logout: Option<OwnedDeviceId>,
122
) {
123
    let _ = events
124
        .send(Event::SessionAccessRejected {
125
            soft_logout: soft_logout.clone(),
126
        })
127
        .await;
128
    let _ = stop.send(match soft_logout {
129
        Some(device_id) => Termination::SoftLoggedOut { device_id },
130
        None => Termination::AccessRevoked,
131
    });
132
}