atagen

macltok

public
Star0
Watch0
5a52012c9b1d7ec241e293fb63039a10884d6fdc
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/engine/mutations.rsRust378 lines13.2 KB
1
use std::{
2
    collections::{HashMap, HashSet, VecDeque},
3
    sync::Arc,
4
};
5
6
use matrix_sdk::{
7
    Client, Room,
8
    ruma::{
9
        Int,
10
        api::client::room::{
11
            Visibility as SdkVisibility,
12
            create_room::v3::{Request as CreateRoomRequest, RoomPreset as SdkPreset},
13
        },
14
        events::{
15
            AnyGlobalAccountDataEventContent, GlobalAccountDataEventType, MessageLikeEventType,
16
            StateEventType, room::power_levels::RoomPowerLevelsEventContent,
17
        },
18
        serde::Raw,
19
    },
20
};
21
use ruma::{OwnedRoomId, OwnedUserId};
22
use tokio::{sync::mpsc, task::JoinSet};
23
24
use super::commands::{Job, guarded};
25
use crate::{
26
    CommandOutput, CreateRoomOptions, Error, ErrorKind, PowerLevelChanges, RoomCommand, RoomPreset,
27
    RoomVisibility,
28
};
29
30
/// state changes that must not race each other. same room, same alias, same dm
31
/// partner or same account data type run in order, everything else runs at once
32
pub(super) enum Mutation {
33
    Room(RoomCommand),
34
    GlobalAccountData {
35
        event_type: GlobalAccountDataEventType,
36
        content: Raw<AnyGlobalAccountDataEventContent>,
37
    },
38
}
39
40
pub(super) type MutationJob = Job<Mutation>;
41
42
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
43
enum Key {
44
    Room(OwnedRoomId),
45
    Alias(String),
46
    Direct(OwnedUserId),
47
    AccountData(String),
48
    Independent(u64),
49
}
50
51
pub(super) async fn run_scheduler(
52
    client: Client,
53
    mut jobs: mpsc::Receiver<MutationJob>,
54
    queue_limit: usize,
55
) {
56
    let fresh_dms = FreshDms::default();
57
    let mut waiting = HashMap::<Key, VecDeque<MutationJob>>::new();
58
    let mut running = HashSet::<Key>::new();
59
    let mut tasks = JoinSet::<Key>::new();
60
    let mut queued = 0_usize;
61
    let mut independent = 0_u64;
62
    let mut receiving = true;
63
    while receiving || !tasks.is_empty() {
64
        tokio::select! {
65
            biased;
66
            Some(finished) = tasks.join_next(), if !tasks.is_empty() => {
67
                // a job's own panic is caught inside it, so this is only ever a key
68
                let Ok(key) = finished else { continue };
69
                let next = waiting.get_mut(&key).and_then(VecDeque::pop_front);
70
                if waiting.get(&key).is_some_and(VecDeque::is_empty) {
71
                    waiting.remove(&key);
72
                }
73
                match next {
74
                    Some(job) => {
75
                        queued -= 1;
76
                        spawn(&mut tasks, &client, &fresh_dms, key, job);
77
                    }
78
                    None => {
79
                        running.remove(&key);
80
                    }
81
                }
82
            }
83
            job = jobs.recv(), if receiving => {
84
                let Some(job) = job else {
85
                    receiving = false;
86
                    continue;
87
                };
88
                let key = key(&job.command, client.user_id().map(ruma::UserId::server_name)).unwrap_or_else(|| {
89
                    independent += 1;
90
                    Key::Independent(independent)
91
                });
92
                if running.insert(key.clone()) {
93
                    spawn(&mut tasks, &client, &fresh_dms, key, job);
94
                } else if queued < queue_limit {
95
                    waiting.entry(key).or_default().push_back(job);
96
                    queued += 1;
97
                } else {
98
                    job.reply.done(Err(Error::new(
99
                        ErrorKind::QueueFull,
100
                        "room command queue is full",
101
                    )));
102
                }
103
            }
104
        }
105
    }
106
}
107
108
/// `server` is ours, since a new alias lives there
109
fn key(mutation: &Mutation, server: Option<&ruma::ServerName>) -> Option<Key> {
110
    let room_id = match mutation {
111
        Mutation::GlobalAccountData { event_type, .. } => {
112
            return Some(Key::AccountData(event_type.to_string()));
113
        }
114
        Mutation::Room(RoomCommand::Join { room, .. }) => {
115
            return Some(match OwnedRoomId::try_from(room.clone()) {
116
                Ok(room_id) => Key::Room(room_id),
117
                Err(alias) => Key::Alias(alias.to_string()),
118
            });
119
        }
120
        Mutation::Room(RoomCommand::CreateDirect(user_id)) => {
121
            return Some(Key::Direct(user_id.clone()));
122
        }
123
        Mutation::Room(RoomCommand::Create(options)) => {
124
            return match (&options.alias_localpart, options.invite.as_slice()) {
125
                (Some(localpart), _) => Some(Key::Alias(format!("#{localpart}:{}", server?))),
126
                (None, [user_id]) if options.direct => Some(Key::Direct(user_id.clone())),
127
                _ => None,
128
            };
129
        }
130
        Mutation::Room(
131
            RoomCommand::Leave(room_id)
132
            | RoomCommand::Forget(room_id)
133
            | RoomCommand::AcceptInvite(room_id)
134
            | RoomCommand::RejectInvite(room_id)
135
            | RoomCommand::InviteUser { room_id, .. }
136
            | RoomCommand::KickUser { room_id, .. }
137
            | RoomCommand::BanUser { room_id, .. }
138
            | RoomCommand::UnbanUser { room_id, .. }
139
            | RoomCommand::SetPowerLevels { room_id, .. }
140
            | RoomCommand::SetUserPowerLevel { room_id, .. }
141
            | RoomCommand::SetAccountData { room_id, .. },
142
        ) => room_id,
143
        Mutation::Room(_) => return None,
144
    };
145
    Some(Key::Room(room_id.clone()))
146
}
147
148
/// dms made here that the next sync hasn't told the sdk about yet, since `get_dm_room`
149
/// only knows `m.direct` once it comes back round
150
type FreshDms = Arc<std::sync::Mutex<HashMap<OwnedUserId, OwnedRoomId>>>;
151
152
fn spawn(
153
    tasks: &mut JoinSet<Key>,
154
    client: &Client,
155
    fresh_dms: &FreshDms,
156
    key: Key,
157
    job: MutationJob,
158
) {
159
    let client = client.clone();
160
    let fresh_dms = fresh_dms.clone();
161
    tasks.spawn(async move {
162
        let Job { command, reply } = job;
163
        reply.send(guarded("room command", execute(&client, &fresh_dms, command)).await);
164
        key
165
    });
166
}
167
168
async fn execute(
169
    client: &Client,
170
    fresh_dms: &FreshDms,
171
    mutation: Mutation,
172
) -> Result<CommandOutput, Error> {
173
    let command = match mutation {
174
        Mutation::GlobalAccountData {
175
            event_type,
176
            content,
177
        } => {
178
            client
179
                .account()
180
                .set_account_data_raw(event_type, content)
181
                .await
182
                .map_err(|source| Error::from_sdk("writing account data failed", source))?;
183
            return Ok(CommandOutput::Done);
184
        }
185
        Mutation::Room(command) => command,
186
    };
187
    let room = |room_id: &OwnedRoomId| client.get_room(room_id).ok_or_else(Error::room_unknown);
188
    let joined = |room: Room| CommandOutput::Room(room.room_id().to_owned());
189
    let sdk = |message: &'static str| move |source| Error::from_sdk(message, source);
190
    match command {
191
        RoomCommand::Join { room, via } => client
192
            .join_room_by_id_or_alias(&room, &via)
193
            .await
194
            .map(joined)
195
            .map_err(sdk("joining room failed")),
196
        RoomCommand::Leave(room_id) => room(&room_id)?
197
            .leave()
198
            .await
199
            .map(done)
200
            .map_err(sdk("leaving room failed")),
201
        RoomCommand::Forget(room_id) => room(&room_id)?
202
            .forget()
203
            .await
204
            .map(done)
205
            .map_err(sdk("forgetting room failed")),
206
        RoomCommand::AcceptInvite(room_id) => room(&room_id)?
207
            .join()
208
            .await
209
            .map(done)
210
            .map_err(sdk("accepting invite failed")),
211
        RoomCommand::RejectInvite(room_id) => room(&room_id)?
212
            .leave()
213
            .await
214
            .map(done)
215
            .map_err(sdk("rejecting invite failed")),
216
        RoomCommand::InviteUser { room_id, user_id } => room(&room_id)?
217
            .invite_user_by_id(&user_id)
218
            .await
219
            .map(done)
220
            .map_err(sdk("inviting user failed")),
221
        RoomCommand::KickUser {
222
            room_id,
223
            user_id,
224
            reason,
225
        } => room(&room_id)?
226
            .kick_user(&user_id, reason.as_deref())
227
            .await
228
            .map(done)
229
            .map_err(sdk("kicking user failed")),
230
        RoomCommand::BanUser {
231
            room_id,
232
            user_id,
233
            reason,
234
        } => room(&room_id)?
235
            .ban_user(&user_id, reason.as_deref())
236
            .await
237
            .map(done)
238
            .map_err(sdk("banning user failed")),
239
        RoomCommand::UnbanUser {
240
            room_id,
241
            user_id,
242
            reason,
243
        } => room(&room_id)?
244
            .unban_user(&user_id, reason.as_deref())
245
            .await
246
            .map(done)
247
            .map_err(sdk("unbanning user failed")),
248
        RoomCommand::SetPowerLevels { room_id, changes } => {
249
            apply_power_level_changes(&room(&room_id)?, changes)
250
                .await
251
                .map(done)
252
        }
253
        RoomCommand::SetUserPowerLevel {
254
            room_id,
255
            user_id,
256
            level,
257
        } => room(&room_id)?
258
            .update_power_levels(vec![(&user_id, power_level(level)?)])
259
            .await
260
            .map(done)
261
            .map_err(sdk("updating user power level failed")),
262
        RoomCommand::SetAccountData {
263
            room_id,
264
            event_type,
265
            content,
266
        } => room(&room_id)?
267
            .set_account_data_raw(event_type, content)
268
            .await
269
            .map(done)
270
            .map_err(sdk("writing room account data failed")),
271
        RoomCommand::Create(options) => client
272
            .create_room(create_room_request(options))
273
            .await
274
            .map(joined)
275
            .map_err(sdk("creating room failed")),
276
        RoomCommand::CreateDirect(user_id) => {
277
            if let Some(existing) = client.get_dm_room(&user_id) {
278
                return Ok(joined(existing));
279
            }
280
            let fresh = fresh_dms
281
                .lock()
282
                .expect("nothing panics holding this")
283
                .get(&user_id)
284
                .cloned();
285
            if let Some(room_id) = fresh.filter(|room_id| client.get_room(room_id).is_some()) {
286
                return Ok(CommandOutput::Room(room_id));
287
            }
288
            let room = client
289
                .create_dm(&user_id)
290
                .await
291
                .map_err(sdk("creating direct room failed"))?;
292
            fresh_dms
293
                .lock()
294
                .expect("nothing panics holding this")
295
                .insert(user_id, room.room_id().to_owned());
296
            Ok(joined(room))
297
        }
298
        _ => Err(Error::new(ErrorKind::Other, "room command was misrouted")),
299
    }
300
}
301
302
/// one state event carrying every requested change, since the sdk's own patcher is private
303
async fn apply_power_level_changes(room: &Room, changes: PowerLevelChanges) -> Result<(), Error> {
304
    let mut levels = room
305
        .power_levels()
306
        .await
307
        .map_err(|source| Error::from_sdk("reading power levels failed", source.into()))?;
308
    for (target, value) in [
309
        (&mut levels.ban, changes.ban),
310
        (&mut levels.invite, changes.invite),
311
        (&mut levels.kick, changes.kick),
312
        (&mut levels.redact, changes.redact),
313
        (&mut levels.events_default, changes.events_default),
314
        (&mut levels.state_default, changes.state_default),
315
        (&mut levels.users_default, changes.users_default),
316
        (&mut levels.notifications.room, changes.room_notification),
317
    ] {
318
        if let Some(value) = value {
319
            *target = power_level(value)?;
320
        }
321
    }
322
    for (event_type, value) in [
323
        (StateEventType::RoomName.into(), changes.room_name),
324
        (StateEventType::RoomAvatar.into(), changes.room_avatar),
325
        (StateEventType::RoomTopic.into(), changes.room_topic),
326
        (StateEventType::SpaceChild.into(), changes.space_child),
327
        (MessageLikeEventType::Beacon.into(), changes.beacon),
328
        (StateEventType::BeaconInfo.into(), changes.beacon_info),
329
    ] {
330
        if let Some(value) = value {
331
            levels.events.insert(event_type, power_level(value)?);
332
        }
333
    }
334
    let content = RoomPowerLevelsEventContent::try_from(levels).map_err(|source| {
335
        Error::with_source(
336
            ErrorKind::Invalid,
337
            "normalising power levels failed",
338
            source,
339
        )
340
    })?;
341
    room.send_state_event(content)
342
        .await
343
        .map(drop)
344
        .map_err(|source| Error::from_sdk("updating power levels failed", source))
345
}
346
347
fn done<T>(_: T) -> CommandOutput {
348
    CommandOutput::Done
349
}
350
351
fn power_level(value: i64) -> Result<Int, Error> {
352
    value.try_into().map_err(|source| {
353
        Error::with_source(
354
            ErrorKind::Invalid,
355
            "power level is outside the Matrix integer range",
356
            source,
357
        )
358
    })
359
}
360
361
fn create_room_request(options: CreateRoomOptions) -> CreateRoomRequest {
362
    let mut request = CreateRoomRequest::new();
363
    request.name = options.name;
364
    request.topic = options.topic;
365
    request.room_alias_name = options.alias_localpart;
366
    request.invite = options.invite;
367
    request.is_direct = options.direct;
368
    request.visibility = match options.visibility {
369
        RoomVisibility::Private => SdkVisibility::Private,
370
        RoomVisibility::Public => SdkVisibility::Public,
371
    };
372
    request.preset = options.preset.map(|preset| match preset {
373
        RoomPreset::PrivateChat => SdkPreset::PrivateChat,
374
        RoomPreset::TrustedPrivateChat => SdkPreset::TrustedPrivateChat,
375
        RoomPreset::PublicChat => SdkPreset::PublicChat,
376
    });
377
    request
378
}