atagen

macltok

public
Star0
Watch0
2656c1d26d6b12a7d19e2c8902c8c88ffe20f982
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/engine/context.rsRust257 lines7.5 KB
1
use std::sync::Arc;
2
3
use matrix_sdk::{
4
    Room, deserialized_responses::TimelineEvent as SdkTimelineEvent, room::MessagesOptions,
5
};
6
use ruma::{OwnedEventId, OwnedUserId, RoomId, UInt};
7
8
use crate::{
9
    Error, ErrorKind, EventContext, EventContextCommand, EventContextDirection, EventContextPage,
10
    TimelineItem,
11
    engine::MediaService,
12
    normalise::{MediaRegistration, normalise_observed},
13
};
14
15
const MAX_CONTEXT_LIMIT: u16 = 50;
16
17
pub(crate) async fn execute(
18
    room: Room,
19
    current_user_id: Option<OwnedUserId>,
20
    media: Arc<MediaService>,
21
    command: EventContextCommand,
22
) -> Result<crate::CommandOutput, Error> {
23
    if let EventContextCommand::Before { token, limit, .. }
24
    | EventContextCommand::After { token, limit, .. } = &command
25
    {
26
        validate_limit(*limit)?;
27
        if token.is_empty() {
28
            return Err(Error::new(
29
                ErrorKind::Invalid,
30
                "context pagination token is empty",
31
            ));
32
        }
33
        return execute_page(
34
            room,
35
            current_user_id,
36
            media,
37
            token,
38
            *limit,
39
            matches!(command, EventContextCommand::Before { .. }),
40
        )
41
        .await;
42
    }
43
44
    let (event_id, before_limit, after_limit) = match &command {
45
        EventContextCommand::Fetch {
46
            event_id,
47
            before_limit,
48
            after_limit,
49
            ..
50
        } => {
51
            validate_limit(*before_limit)?;
52
            validate_limit(*after_limit)?;
53
            (event_id.clone(), *before_limit, *after_limit)
54
        }
55
        EventContextCommand::Before { .. } | EventContextCommand::After { .. } => unreachable!(),
56
    };
57
    let (context, registrations, _) =
58
        capture(room, current_user_id, event_id, before_limit, after_limit).await?;
59
    media.register(registrations).await;
60
    Ok(crate::CommandOutput::EventContext(context))
61
}
62
63
pub(crate) async fn capture(
64
    room: Room,
65
    current_user_id: Option<OwnedUserId>,
66
    event_id: OwnedEventId,
67
    before_limit: u16,
68
    after_limit: u16,
69
) -> Result<(EventContext, Vec<MediaRegistration>, Vec<serde_json::Value>), Error> {
70
    validate_limit(before_limit)?;
71
    validate_limit(after_limit)?;
72
    let raw = room
73
        .event_with_context(&event_id, false, UInt::from(0_u16), None)
74
        .await
75
        .map_err(|source| {
76
            let error = Error::from_sdk("fetching event context failed", source);
77
            if error
78
                .server
79
                .as_ref()
80
                .is_some_and(|server| server.status == 404)
81
            {
82
                Error::new(ErrorKind::EventNotFound, "event was not found")
83
            } else {
84
                error
85
            }
86
        })?;
87
    if raw.event.is_none() {
88
        return Err(Error::new(ErrorKind::EventNotFound, "event was not found"));
89
    }
90
    let before = fetch_side(
91
        &room,
92
        raw.prev_batch_token.as_deref(),
93
        before_limit,
94
        true,
95
        current_user_id.as_deref(),
96
    );
97
    let after = fetch_side(
98
        &room,
99
        raw.next_batch_token.as_deref(),
100
        after_limit,
101
        false,
102
        current_user_id.as_deref(),
103
    );
104
    let (events_before, mut before_media, start_token, mut observed) = before.await?;
105
    let (events_after, after_media, end_token, following) = after.await?;
106
    observed.extend(following);
107
    if let Some(value) = raw
108
        .event
109
        .as_ref()
110
        .and_then(|event| serde_json::from_str(event.raw().json().get()).ok())
111
    {
112
        observed.push(value);
113
    }
114
    let event = raw.event.as_ref().and_then(|event| {
115
        normalise_sdk_event(event, current_user_id.as_deref(), room.room_id()).map(
116
            |(item, media)| {
117
                before_media.extend(media);
118
                item
119
            },
120
        )
121
    });
122
    before_media.extend(after_media);
123
    Ok((
124
        EventContext {
125
            events_before,
126
            event,
127
            events_after,
128
            start_token,
129
            end_token,
130
        },
131
        before_media,
132
        observed,
133
    ))
134
}
135
136
async fn fetch_side(
137
    room: &Room,
138
    token: Option<&str>,
139
    limit: u16,
140
    backwards: bool,
141
    current_user_id: Option<&ruma::UserId>,
142
) -> Result<
143
    (
144
        Vec<TimelineItem>,
145
        Vec<MediaRegistration>,
146
        Option<String>,
147
        Vec<serde_json::Value>,
148
    ),
149
    Error,
150
> {
151
    if limit == 0 {
152
        return Ok((Vec::new(), Vec::new(), token.map(str::to_owned), Vec::new()));
153
    }
154
    let Some(token) = token else {
155
        return Ok((Vec::new(), Vec::new(), None, Vec::new()));
156
    };
157
    let mut options = if backwards {
158
        MessagesOptions::backward()
159
    } else {
160
        MessagesOptions::forward()
161
    };
162
    options = options.from(token);
163
    options.limit = UInt::from(limit);
164
    let messages = room
165
        .messages(options)
166
        .await
167
        .map_err(|source| Error::from_sdk("fetching event context events failed", source))?;
168
    let mut events = messages.chunk;
169
    events.truncate(usize::from(limit));
170
    if backwards {
171
        events.reverse();
172
    }
173
    let next_token = messages.end;
174
    let observed = events
175
        .iter()
176
        .filter_map(|event| serde_json::from_str(event.raw().json().get()).ok())
177
        .collect();
178
    let (items, media) = normalise_events(events, current_user_id, room.room_id());
179
    Ok((items, media, next_token, observed))
180
}
181
182
async fn execute_page(
183
    room: Room,
184
    current_user_id: Option<OwnedUserId>,
185
    media: Arc<MediaService>,
186
    token: &str,
187
    limit: u16,
188
    backwards: bool,
189
) -> Result<crate::CommandOutput, Error> {
190
    let mut options = if backwards {
191
        MessagesOptions::backward()
192
    } else {
193
        MessagesOptions::forward()
194
    };
195
    options = options.from(token);
196
    options.limit = UInt::from(limit);
197
    let messages = room
198
        .messages(options)
199
        .await
200
        .map_err(|source| Error::from_sdk("fetching event context page failed", source))?;
201
    let next_token = messages.end.clone();
202
    let mut events = messages.chunk;
203
    events.truncate(usize::from(limit));
204
    if backwards {
205
        events.reverse();
206
    }
207
    let (events, registrations) =
208
        normalise_events(events, current_user_id.as_deref(), room.room_id());
209
    media.register(registrations).await;
210
    let end_reached = next_token.is_none();
211
    let direction = if backwards {
212
        EventContextDirection::Before
213
    } else {
214
        EventContextDirection::After
215
    };
216
    Ok(crate::CommandOutput::EventContextPage(EventContextPage {
217
        direction,
218
        events,
219
        next_token,
220
        end_reached,
221
    }))
222
}
223
224
fn normalise_events(
225
    events: Vec<SdkTimelineEvent>,
226
    current_user_id: Option<&ruma::UserId>,
227
    room_id: &RoomId,
228
) -> (Vec<TimelineItem>, Vec<MediaRegistration>) {
229
    let mut items = Vec::with_capacity(events.len());
230
    let mut media = Vec::new();
231
    for event in events {
232
        if let Some((item, registrations)) = normalise_sdk_event(&event, current_user_id, room_id) {
233
            items.push(item);
234
            media.extend(registrations);
235
        }
236
    }
237
    (items, media)
238
}
239
240
fn normalise_sdk_event(
241
    event: &SdkTimelineEvent,
242
    current_user_id: Option<&ruma::UserId>,
243
    room_id: &RoomId,
244
) -> Option<(TimelineItem, Vec<MediaRegistration>)> {
245
    let value: serde_json::Value = serde_json::from_str(event.raw().json().get()).ok()?;
246
    let event_id = value.get("event_id")?.as_str()?;
247
    let identity = format!("{room_id}\n{event_id}");
248
    normalise_observed(&value, current_user_id, &identity)
249
}
250
251
fn validate_limit(limit: u16) -> Result<(), Error> {
252
    if limit > MAX_CONTEXT_LIMIT {
253
        Err(Error::new(ErrorKind::Invalid, "context limit exceeds 50"))
254
    } else {
255
        Ok(())
256
    }
257
}