palpo_data/user/
key_backup.rs

1use std::collections::BTreeMap;
2
3use diesel::prelude::*;
4
5use crate::core::UnixMillis;
6use crate::core::client::backup::{BackupAlgorithm, KeyBackupData};
7use crate::core::identifiers::*;
8use crate::core::serde::{JsonValue, RawJson};
9use crate::schema::*;
10use crate::{DataResult,  connect};
11
12#[derive(Identifiable, Queryable, Debug, Clone)]
13#[diesel(table_name = e2e_room_keys)]
14pub struct DbRoomKey {
15    pub id: i64,
16
17    pub user_id: OwnedUserId,
18    pub room_id: OwnedRoomId,
19    pub session_id: OwnedSessionId,
20
21    pub version: i64,
22
23    pub first_message_index: Option<i64>,
24    pub forwarded_count: Option<i64>,
25    pub is_verified: bool,
26    pub session_data: JsonValue,
27    pub created_at: UnixMillis,
28}
29#[derive(Insertable, AsChangeset, Debug, Clone)]
30#[diesel(table_name = e2e_room_keys)]
31pub struct NewDbRoomKey {
32    pub user_id: OwnedUserId,
33    pub room_id: OwnedRoomId,
34    pub session_id: OwnedSessionId,
35
36    pub version: i64,
37
38    pub first_message_index: Option<i64>,
39    pub forwarded_count: Option<i64>,
40    pub is_verified: bool,
41    pub session_data: JsonValue,
42    pub created_at: UnixMillis,
43}
44
45impl Into<KeyBackupData> for DbRoomKey {
46    fn into(self) -> KeyBackupData {
47        KeyBackupData {
48            first_message_index: self.first_message_index.unwrap_or(0) as u64,
49            forwarded_count: self.forwarded_count.unwrap_or(0) as u64,
50            is_verified: self.is_verified,
51            session_data: serde_json::from_value(self.session_data).unwrap(),
52        }
53    }
54}
55
56#[derive(Identifiable, Queryable, Debug, Clone)]
57#[diesel(table_name = e2e_room_keys_versions)]
58pub struct DbRoomKeysVersion {
59    pub id: i64,
60
61    pub user_id: OwnedUserId,
62    pub version: i64,
63    pub algorithm: JsonValue,
64    pub auth_data: JsonValue,
65    pub is_trashed: bool,
66    pub etag: i64,
67    pub created_at: UnixMillis,
68}
69#[derive(Insertable, Debug, Clone)]
70#[diesel(table_name = e2e_room_keys_versions)]
71pub struct NewDbRoomKeysVersion {
72    pub user_id: OwnedUserId,
73    pub version: i64,
74    pub algorithm: JsonValue,
75    pub auth_data: JsonValue,
76    pub created_at: UnixMillis,
77}
78
79pub fn create_backup(user_id: &UserId, algorithm: &RawJson<BackupAlgorithm>) -> DataResult<DbRoomKeysVersion> {
80    let version = UnixMillis::now().get() as i64;
81    let new_keys_version = NewDbRoomKeysVersion {
82        user_id: user_id.to_owned(),
83        version,
84        algorithm: serde_json::to_value(algorithm)?,
85        auth_data: serde_json::to_value(BTreeMap::<String, JsonValue>::new())?,
86        created_at: UnixMillis::now(),
87    };
88    diesel::insert_into(e2e_room_keys_versions::table)
89        .values(&new_keys_version)
90        .get_result(&mut connect()?)
91        .map_err(Into::into)
92}
93
94pub fn update_backup(user_id: &UserId, version: i64, algorithm: &BackupAlgorithm) -> DataResult<()> {
95    diesel::update(
96        e2e_room_keys_versions::table
97            .filter(e2e_room_keys_versions::user_id.eq(user_id))
98            .filter(e2e_room_keys_versions::version.eq(version)),
99    )
100    .set((
101        e2e_room_keys_versions::algorithm.eq(serde_json::to_value(algorithm)?),
102        e2e_room_keys_versions::etag.eq(UnixMillis::now().get() as i64),
103    ))
104    .execute(&mut connect()?)?;
105    Ok(())
106}
107
108pub fn get_latest_room_key(user_id: &UserId) -> DataResult<Option<DbRoomKey>> {
109    e2e_room_keys::table
110        .filter(e2e_room_keys::user_id.eq(user_id))
111        .order(e2e_room_keys::version.desc())
112        .first::<DbRoomKey>(&mut connect()?)
113        .optional()
114        .map_err(Into::into)
115}
116
117pub fn get_room_key(user_id: &UserId, room_id: &RoomId, version: i64) -> DataResult<Option<DbRoomKey>> {
118    e2e_room_keys::table
119        .filter(e2e_room_keys::user_id.eq(user_id))
120        .filter(e2e_room_keys::room_id.eq(room_id))
121        .filter(e2e_room_keys::version.eq(version))
122        .first::<DbRoomKey>(&mut connect()?)
123        .optional()
124        .map_err(Into::into)
125}
126
127pub fn get_latest_room_keys_version(user_id: &UserId) -> DataResult<Option<DbRoomKeysVersion>> {
128    e2e_room_keys_versions::table
129        .filter(e2e_room_keys_versions::user_id.eq(user_id))
130        .order(e2e_room_keys_versions::version.desc())
131        .first::<DbRoomKeysVersion>(&mut connect()?)
132        .optional()
133        .map_err(Into::into)
134}
135pub fn get_room_keys_version(user_id: &UserId, version: i64) -> DataResult<Option<DbRoomKeysVersion>> {
136    e2e_room_keys_versions::table
137        .filter(e2e_room_keys_versions::user_id.eq(user_id))
138        .filter(e2e_room_keys_versions::version.eq(version))
139        .first::<DbRoomKeysVersion>(&mut connect()?)
140        .optional()
141        .map_err(Into::into)
142}
143
144pub fn add_key(
145    user_id: &UserId,
146    version: i64,
147    room_id: &RoomId,
148    session_id: &SessionId,
149    key_data: &KeyBackupData,
150) -> DataResult<()> {
151    let new_key = NewDbRoomKey {
152        user_id: user_id.to_owned(),
153        room_id: room_id.to_owned(),
154        session_id: session_id.to_owned(),
155        version: version.to_owned(),
156        first_message_index: Some(key_data.first_message_index as i64),
157        forwarded_count: Some(key_data.forwarded_count as i64),
158        is_verified: key_data.is_verified,
159        session_data: serde_json::to_value(&key_data.session_data)?,
160        created_at: UnixMillis::now(),
161    };
162
163    let exist_key = get_key_for_session(user_id, version, room_id, session_id)?;
164    let replace = if let Some(exist_key) = exist_key {
165        if new_key.is_verified && !exist_key.is_verified {
166            true
167        } else if new_key.first_message_index < exist_key.first_message_index {
168            true
169        } else if new_key.first_message_index == exist_key.first_message_index {
170            new_key.forwarded_count < exist_key.forwarded_count
171        } else {
172            false
173        }
174    } else {
175        true
176    };
177    if replace {
178        diesel::insert_into(e2e_room_keys::table)
179            .values(&new_key)
180            .on_conflict((
181                e2e_room_keys::user_id,
182                e2e_room_keys::room_id,
183                e2e_room_keys::session_id,
184                e2e_room_keys::version,
185            ))
186            .do_update()
187            .set(&new_key)
188            .execute(&mut *connect()?)?;
189    }
190    Ok(())
191}
192
193pub fn count_keys(user_id: &UserId, version: i64) -> DataResult<i64> {
194    e2e_room_keys::table
195        .filter(e2e_room_keys::user_id.eq(user_id))
196        .filter(e2e_room_keys::version.eq(version))
197        .count()
198        .get_result(&mut connect()?)
199        .map_err(Into::into)
200}
201
202pub fn get_etag(user_id: &UserId, version: i64) -> DataResult<String> {
203    e2e_room_keys_versions::table
204        .filter(e2e_room_keys_versions::user_id.eq(user_id))
205        .filter(e2e_room_keys_versions::version.eq(version))
206        .select(e2e_room_keys_versions::etag)
207        .first(&mut connect()?)
208        .map(|etag: i64| etag.to_string())
209        .map_err(Into::into)
210}
211
212pub fn get_key_for_session(
213    user_id: &UserId,
214    version: i64,
215    room_id: &RoomId,
216    session_id: &SessionId,
217) -> DataResult<Option<DbRoomKey>> {
218    e2e_room_keys::table
219        .filter(e2e_room_keys::user_id.eq(user_id))
220        .filter(e2e_room_keys::version.eq(version))
221        .filter(e2e_room_keys::room_id.eq(room_id))
222        .filter(e2e_room_keys::session_id.eq(session_id))
223        .first::<DbRoomKey>(&mut connect()?)
224        .optional()
225        .map_err(Into::into)
226}
227
228pub fn delete_backup(user_id: &UserId, version: i64) -> DataResult<()> {
229    delete_all_keys(user_id, version)?;
230    diesel::update(
231        e2e_room_keys_versions::table
232            .filter(e2e_room_keys_versions::user_id.eq(user_id))
233            .filter(e2e_room_keys_versions::version.eq(version)),
234    )
235    .set(e2e_room_keys_versions::is_trashed.eq(true))
236    .execute(&mut connect()?)?;
237    Ok(())
238}
239
240pub fn delete_all_keys(user_id: &UserId, version: i64) -> DataResult<()> {
241    diesel::delete(
242        e2e_room_keys::table
243            .filter(e2e_room_keys::user_id.eq(user_id))
244            .filter(e2e_room_keys::version.eq(version)),
245    )
246    .execute(&mut connect()?)?;
247    Ok(())
248}
249
250pub fn delete_room_keys(user_id: &UserId, version: i64, room_id: &RoomId) -> DataResult<()> {
251    diesel::delete(
252        e2e_room_keys::table
253            .filter(e2e_room_keys::user_id.eq(user_id))
254            .filter(e2e_room_keys::version.eq(version))
255            .filter(e2e_room_keys::room_id.eq(room_id)),
256    )
257    .execute(&mut connect()?)?;
258    Ok(())
259}
260
261pub fn delete_room_key(user_id: &UserId, version: i64, room_id: &RoomId, session_id: &SessionId) -> DataResult<()> {
262    diesel::delete(
263        e2e_room_keys::table
264            .filter(e2e_room_keys::user_id.eq(user_id))
265            .filter(e2e_room_keys::version.eq(version))
266            .filter(e2e_room_keys::room_id.eq(room_id))
267            .filter(e2e_room_keys::session_id.eq(session_id)),
268    )
269    .execute(&mut connect()?)?;
270    Ok(())
271}
palpo_data/user/key_backup.rs

palpo_data/user/
key_backup.rs
