From d0ee82325449bc935fd6b39eee50191cc638ec77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6sters?= Date: Tue, 8 Jun 2021 18:10:00 +0200 Subject: [PATCH 1/8] feat: swappable database backend --- Cargo.lock | 106 ++++++ Cargo.toml | 1 + src/client_server/account.rs | 12 +- src/client_server/alias.rs | 11 +- src/client_server/backup.rs | 32 +- src/client_server/config.rs | 10 +- src/client_server/context.rs | 4 +- src/client_server/device.rs | 12 +- src/client_server/directory.rs | 10 +- src/client_server/keys.rs | 17 +- src/client_server/media.rs | 10 +- src/client_server/membership.rs | 22 +- src/client_server/message.rs | 5 +- src/client_server/presence.rs | 8 +- src/client_server/profile.rs | 12 +- src/client_server/push.rs | 22 +- src/client_server/read_marker.rs | 6 +- src/client_server/redact.rs | 3 +- src/client_server/room.rs | 8 +- src/client_server/search.rs | 3 +- src/client_server/session.rs | 8 +- src/client_server/state.rs | 12 +- src/client_server/sync.rs | 34 +- src/client_server/tag.rs | 8 +- src/client_server/to_device.rs | 4 +- src/client_server/typing.rs | 4 +- src/client_server/user_directory.rs | 4 +- src/database.rs | 185 ++++----- src/database/abstraction.rs | 309 +++++++++++++++ src/database/account_data.rs | 57 ++- src/database/admin.rs | 30 +- src/database/appservice.rs | 31 +- src/database/globals.rs | 81 ++-- src/database/key_backups.rs | 141 ++++--- src/database/media.rs | 23 +- src/database/pusher.rs | 34 +- src/database/rooms.rs | 565 ++++++++++++++-------------- src/database/rooms/edus.rs | 103 +++-- src/database/sending.rs | 122 +++--- src/database/transaction_ids.rs | 14 +- src/database/uiaa.rs | 13 +- src/database/users.rs | 234 ++++++------ src/error.rs | 7 +- src/main.rs | 6 +- src/ruma_wrapper.rs | 7 +- src/server_server.rs | 64 ++-- src/utils.rs | 27 +- 47 files changed, 1447 insertions(+), 994 deletions(-) create mode 100644 src/database/abstraction.rs diff --git a/Cargo.lock b/Cargo.lock index 36a66590..630d4142 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -103,6 +103,25 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "383d29d513d8764dcdc42ea295d979eb99c3c9f00607b3692cf68a431f7dca72" +[[package]] +name = "bindgen" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd4865004a46a0aafb2a0a5eb19d3c9fc46ee5f063a6cfc605c69ac9ecf5263d" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "lazy_static", + "lazycell", + "peeking_take_while", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", +] + [[package]] name = "bitflags" version = "1.2.1" @@ -162,6 +181,15 @@ dependencies = [ "jobserver", ] +[[package]] +name = "cexpr" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4aedb84272dbe89af497cf81375129abda4fc0a9e7c5d317498c15cc30c0d27" +dependencies = [ + "nom", +] + [[package]] name = "cfg-if" version = "0.1.10" @@ -187,6 +215,17 @@ dependencies = [ "winapi", ] +[[package]] +name = "clang-sys" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "853eda514c284c2287f4bf20ae614f8781f40a81d32ecda6e91449304dfe077c" +dependencies = [ + "glob", + "libc", + "libloading", +] + [[package]] name = "color_quant" version = "1.1.0" @@ -212,6 +251,7 @@ dependencies = [ "reqwest", "ring", "rocket", + "rocksdb", "ruma", "rust-argon2", "rustls", @@ -1008,12 +1048,40 @@ version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + [[package]] name = "libc" version = "0.2.95" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "789da6d93f1b866ffe175afc5322a4d76c038605a1c3319bb57b06967ca98a36" +[[package]] +name = "libloading" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f84d96438c15fcd6c3f244c8fce01d1e2b9c6b5623e9c711dc9286d8fc92d6a" +dependencies = [ + "cfg-if 1.0.0", + "winapi", +] + +[[package]] +name = "librocksdb-sys" +version = "6.17.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5da125e1c0f22c7cae785982115523a0738728498547f415c9054cb17c7e89f9" +dependencies = [ + "bindgen", + "cc", + "glob", + "libc", +] + [[package]] name = "linked-hash-map" version = "0.5.4" @@ -1158,6 +1226,16 @@ dependencies = [ "version_check", ] +[[package]] +name = "nom" +version = "5.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffb4262d26ed83a1c0a33a38fe2bb15797329c85770da05e6b828ddb782627af" +dependencies = [ + "memchr", + "version_check", +] + [[package]] name = "ntapi" version = "0.3.6" @@ -1339,6 +1417,12 @@ dependencies = [ "syn", ] +[[package]] +name = "peeking_take_while" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" + [[package]] name = "pem" version = "0.8.3" @@ -1777,6 +1861,16 @@ dependencies = [ "uncased", ] +[[package]] +name = "rocksdb" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c749134fda8bfc90d0de643d59bfc841dcb3ac8a1062e12b6754bd60235c48b3" +dependencies = [ + "libc", + "librocksdb-sys", +] + [[package]] name = "ruma" version = "0.1.2" @@ -2046,6 +2140,12 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc_version" version = "0.2.3" @@ -2245,6 +2345,12 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "shlex" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fdf1b9db47230893d76faad238fd6097fd6d6a9245cd7a4d90dbd639536bbd2" + [[package]] name = "signal-hook-registry" version = "1.3.0" diff --git a/Cargo.toml b/Cargo.toml index f36d8382..eb43da5a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,6 +25,7 @@ ruma = { git = "https://github.com/ruma/ruma", rev = "b39537812c12caafcbf8b7bd74 tokio = "1.2.0" # Used for storing data permanently sled = { version = "0.34.6", features = ["compression", "no_metrics"] } +rocksdb = { version = "0.16.0", features = ["multi-threaded-cf"] } #sled = { git = "https://github.com/spacejam/sled.git", rev = "e4640e0773595229f398438886f19bca6f7326a2", features = ["compression"] } # Used for the http request / response body type for Ruma endpoints used with reqwest diff --git a/src/client_server/account.rs b/src/client_server/account.rs index 0cf30a07..56de5fc9 100644 --- a/src/client_server/account.rs +++ b/src/client_server/account.rs @@ -1,4 +1,4 @@ -use std::{collections::BTreeMap, convert::TryInto}; +use std::{collections::BTreeMap, convert::TryInto, sync::Arc}; use super::{State, DEVICE_ID_LENGTH, SESSION_ID_LENGTH, TOKEN_LENGTH}; use crate::{pdu::PduBuilder, utils, ConduitResult, Database, Error, Ruma}; @@ -42,7 +42,7 @@ const GUEST_NAME_LENGTH: usize = 10; )] #[tracing::instrument(skip(db, body))] pub async fn get_register_available_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { // Validate user id @@ -85,7 +85,7 @@ pub async fn get_register_available_route( )] #[tracing::instrument(skip(db, body))] pub async fn register_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_registration() { @@ -227,7 +227,7 @@ pub async fn register_route( )?; // If this is the first user on this server, create the admins room - if db.users.count() == 1 { + if db.users.count()? == 1 { // Create a user for the server let conduit_user = UserId::parse_with_server_name("conduit", db.globals.server_name()) .expect("@conduit:server_name is valid"); @@ -506,7 +506,7 @@ pub async fn register_route( )] #[tracing::instrument(skip(db, body))] pub async fn change_password_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -598,7 +598,7 @@ pub async fn whoami_route(body: Ruma) -> ConduitResult, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/alias.rs b/src/client_server/alias.rs index 07b49773..40252af2 100644 --- a/src/client_server/alias.rs +++ b/src/client_server/alias.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{ConduitResult, Database, Error, Ruma}; use regex::Regex; @@ -22,7 +24,7 @@ use rocket::{delete, get, put}; )] #[tracing::instrument(skip(db, body))] pub async fn create_alias_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if db.rooms.id_from_alias(&body.room_alias)?.is_some() { @@ -43,7 +45,7 @@ pub async fn create_alias_route( )] #[tracing::instrument(skip(db, body))] pub async fn delete_alias_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { db.rooms.set_alias(&body.room_alias, None, &db.globals)?; @@ -59,7 +61,7 @@ pub async fn delete_alias_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_alias_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { get_alias_helper(&db, &body.room_alias).await @@ -86,7 +88,8 @@ pub async fn get_alias_helper( match db.rooms.id_from_alias(&room_alias)? { Some(r) => room_id = Some(r), None => { - for (_id, registration) in db.appservice.iter_all().filter_map(|r| r.ok()) { + let iter = db.appservice.iter_all()?; + for (_id, registration) in iter.filter_map(|r| r.ok()) { let aliases = registration .get("namespaces") .and_then(|ns| ns.get("aliases")) diff --git a/src/client_server/backup.rs b/src/client_server/backup.rs index 12f3bfd3..fcca676f 100644 --- a/src/client_server/backup.rs +++ b/src/client_server/backup.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{ConduitResult, Database, Error, Ruma}; use ruma::api::client::{ @@ -19,7 +21,7 @@ use rocket::{delete, get, post, put}; )] #[tracing::instrument(skip(db, body))] pub async fn create_backup_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -38,7 +40,7 @@ pub async fn create_backup_route( )] #[tracing::instrument(skip(db, body))] pub async fn update_backup_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -56,7 +58,7 @@ pub async fn update_backup_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_latest_backup_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -84,7 +86,7 @@ pub async fn get_latest_backup_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_backup_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -111,7 +113,7 @@ pub async fn get_backup_route( )] #[tracing::instrument(skip(db, body))] pub async fn delete_backup_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -130,7 +132,7 @@ pub async fn delete_backup_route( )] #[tracing::instrument(skip(db, body))] pub async fn add_backup_keys_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -164,7 +166,7 @@ pub async fn add_backup_keys_route( )] #[tracing::instrument(skip(db, body))] pub async fn add_backup_key_sessions_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -196,7 +198,7 @@ pub async fn add_backup_key_sessions_route( )] #[tracing::instrument(skip(db, body))] pub async fn add_backup_key_session_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -225,7 +227,7 @@ pub async fn add_backup_key_session_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_backup_keys_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -241,14 +243,14 @@ pub async fn get_backup_keys_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_backup_key_sessions_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); let sessions = db .key_backups - .get_room(&sender_user, &body.version, &body.room_id); + .get_room(&sender_user, &body.version, &body.room_id)?; Ok(get_backup_key_sessions::Response { sessions }.into()) } @@ -259,7 +261,7 @@ pub async fn get_backup_key_sessions_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_backup_key_session_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -281,7 +283,7 @@ pub async fn get_backup_key_session_route( )] #[tracing::instrument(skip(db, body))] pub async fn delete_backup_keys_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -304,7 +306,7 @@ pub async fn delete_backup_keys_route( )] #[tracing::instrument(skip(db, body))] pub async fn delete_backup_key_sessions_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -327,7 +329,7 @@ pub async fn delete_backup_key_sessions_route( )] #[tracing::instrument(skip(db, body))] pub async fn delete_backup_key_session_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/config.rs b/src/client_server/config.rs index ce437efd..829bf94a 100644 --- a/src/client_server/config.rs +++ b/src/client_server/config.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{ConduitResult, Database, Error, Ruma}; use ruma::{ @@ -23,7 +25,7 @@ use rocket::{get, put}; )] #[tracing::instrument(skip(db, body))] pub async fn set_global_account_data_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -58,7 +60,7 @@ pub async fn set_global_account_data_route( )] #[tracing::instrument(skip(db, body))] pub async fn set_room_account_data_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -90,7 +92,7 @@ pub async fn set_room_account_data_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_global_account_data_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -117,7 +119,7 @@ pub async fn get_global_account_data_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_room_account_data_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/context.rs b/src/client_server/context.rs index 1fee2f26..b86fd0bf 100644 --- a/src/client_server/context.rs +++ b/src/client_server/context.rs @@ -1,7 +1,7 @@ use super::State; use crate::{ConduitResult, Database, Error, Ruma}; use ruma::api::client::{error::ErrorKind, r0::context::get_context}; -use std::convert::TryFrom; +use std::{convert::TryFrom, sync::Arc}; #[cfg(feature = "conduit_bin")] use rocket::get; @@ -12,7 +12,7 @@ use rocket::get; )] #[tracing::instrument(skip(db, body))] pub async fn get_context_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/device.rs b/src/client_server/device.rs index 961ba97a..2441524d 100644 --- a/src/client_server/device.rs +++ b/src/client_server/device.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{utils, ConduitResult, Database, Error, Ruma}; use ruma::api::client::{ @@ -18,7 +20,7 @@ use rocket::{delete, get, post, put}; )] #[tracing::instrument(skip(db, body))] pub async fn get_devices_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -38,7 +40,7 @@ pub async fn get_devices_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_device_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -57,7 +59,7 @@ pub async fn get_device_route( )] #[tracing::instrument(skip(db, body))] pub async fn update_device_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -83,7 +85,7 @@ pub async fn update_device_route( )] #[tracing::instrument(skip(db, body))] pub async fn delete_device_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -137,7 +139,7 @@ pub async fn delete_device_route( )] #[tracing::instrument(skip(db, body))] pub async fn delete_devices_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/directory.rs b/src/client_server/directory.rs index 9864a5e3..ad609cd5 100644 --- a/src/client_server/directory.rs +++ b/src/client_server/directory.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{ConduitResult, Database, Error, Result, Ruma}; use log::info; @@ -33,7 +35,7 @@ use rocket::{get, post, put}; )] #[tracing::instrument(skip(db, body))] pub async fn get_public_rooms_filtered_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { get_public_rooms_filtered_helper( @@ -53,7 +55,7 @@ pub async fn get_public_rooms_filtered_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_public_rooms_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let response = get_public_rooms_filtered_helper( @@ -82,7 +84,7 @@ pub async fn get_public_rooms_route( )] #[tracing::instrument(skip(db, body))] pub async fn set_room_visibility_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -112,7 +114,7 @@ pub async fn set_room_visibility_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_room_visibility_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { Ok(get_room_visibility::Response { diff --git a/src/client_server/keys.rs b/src/client_server/keys.rs index d856bf31..f80a3294 100644 --- a/src/client_server/keys.rs +++ b/src/client_server/keys.rs @@ -14,7 +14,10 @@ use ruma::{ encryption::UnsignedDeviceInfo, DeviceId, DeviceKeyAlgorithm, UserId, }; -use std::collections::{BTreeMap, HashSet}; +use std::{ + collections::{BTreeMap, HashSet}, + sync::Arc, +}; #[cfg(feature = "conduit_bin")] use rocket::{get, post}; @@ -25,7 +28,7 @@ use rocket::{get, post}; )] #[tracing::instrument(skip(db, body))] pub async fn upload_keys_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -74,7 +77,7 @@ pub async fn upload_keys_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_keys_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -95,7 +98,7 @@ pub async fn get_keys_route( )] #[tracing::instrument(skip(db, body))] pub async fn claim_keys_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let response = claim_keys_helper(&body.one_time_keys, &db)?; @@ -111,7 +114,7 @@ pub async fn claim_keys_route( )] #[tracing::instrument(skip(db, body))] pub async fn upload_signing_keys_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -174,7 +177,7 @@ pub async fn upload_signing_keys_route( )] #[tracing::instrument(skip(db, body))] pub async fn upload_signatures_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -235,7 +238,7 @@ pub async fn upload_signatures_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_key_changes_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/media.rs b/src/client_server/media.rs index 74ca6c84..0673787c 100644 --- a/src/client_server/media.rs +++ b/src/client_server/media.rs @@ -7,14 +7,14 @@ use ruma::api::client::{ #[cfg(feature = "conduit_bin")] use rocket::{get, post}; -use std::convert::TryInto; +use std::{convert::TryInto, sync::Arc}; const MXC_LENGTH: usize = 32; #[cfg_attr(feature = "conduit_bin", get("/_matrix/media/r0/config"))] #[tracing::instrument(skip(db))] pub async fn get_media_config_route( - db: State<'_, Database>, + db: State<'_, Arc>, ) -> ConduitResult { Ok(get_media_config::Response { upload_size: db.globals.max_request_size().into(), @@ -28,7 +28,7 @@ pub async fn get_media_config_route( )] #[tracing::instrument(skip(db, body))] pub async fn create_content_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let mxc = format!( @@ -62,7 +62,7 @@ pub async fn create_content_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_content_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let mxc = format!("mxc://{}/{}", body.server_name, body.media_id); @@ -112,7 +112,7 @@ pub async fn get_content_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_content_thumbnail_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let mxc = format!("mxc://{}/{}", body.server_name, body.media_id); diff --git a/src/client_server/membership.rs b/src/client_server/membership.rs index 96fe8001..92d7aced 100644 --- a/src/client_server/membership.rs +++ b/src/client_server/membership.rs @@ -44,7 +44,7 @@ use rocket::{get, post}; )] #[tracing::instrument(skip(db, body))] pub async fn join_room_by_id_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -81,7 +81,7 @@ pub async fn join_room_by_id_route( )] #[tracing::instrument(skip(db, body))] pub async fn join_room_by_id_or_alias_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -135,7 +135,7 @@ pub async fn join_room_by_id_or_alias_route( )] #[tracing::instrument(skip(db, body))] pub async fn leave_room_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -153,7 +153,7 @@ pub async fn leave_room_route( )] #[tracing::instrument(skip(db, body))] pub async fn invite_user_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -173,7 +173,7 @@ pub async fn invite_user_route( )] #[tracing::instrument(skip(db, body))] pub async fn kick_user_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -222,7 +222,7 @@ pub async fn kick_user_route( )] #[tracing::instrument(skip(db, body))] pub async fn ban_user_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -279,7 +279,7 @@ pub async fn ban_user_route( )] #[tracing::instrument(skip(db, body))] pub async fn unban_user_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -327,7 +327,7 @@ pub async fn unban_user_route( )] #[tracing::instrument(skip(db, body))] pub async fn forget_room_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -345,7 +345,7 @@ pub async fn forget_room_route( )] #[tracing::instrument(skip(db, body))] pub async fn joined_rooms_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -366,7 +366,7 @@ pub async fn joined_rooms_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_member_events_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -396,7 +396,7 @@ pub async fn get_member_events_route( )] #[tracing::instrument(skip(db, body))] pub async fn joined_members_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/message.rs b/src/client_server/message.rs index 96de93dd..0d19f347 100644 --- a/src/client_server/message.rs +++ b/src/client_server/message.rs @@ -11,6 +11,7 @@ use ruma::{ use std::{ collections::BTreeMap, convert::{TryFrom, TryInto}, + sync::Arc, }; #[cfg(feature = "conduit_bin")] @@ -22,7 +23,7 @@ use rocket::{get, put}; )] #[tracing::instrument(skip(db, body))] pub async fn send_message_event_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -85,7 +86,7 @@ pub async fn send_message_event_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_message_events_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/presence.rs b/src/client_server/presence.rs index 9f4f7a39..ce80dfd7 100644 --- a/src/client_server/presence.rs +++ b/src/client_server/presence.rs @@ -1,7 +1,7 @@ use super::State; use crate::{utils, ConduitResult, Database, Ruma}; use ruma::api::client::r0::presence::{get_presence, set_presence}; -use std::{convert::TryInto, time::Duration}; +use std::{convert::TryInto, sync::Arc, time::Duration}; #[cfg(feature = "conduit_bin")] use rocket::{get, put}; @@ -12,7 +12,7 @@ use rocket::{get, put}; )] #[tracing::instrument(skip(db, body))] pub async fn set_presence_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -53,7 +53,7 @@ pub async fn set_presence_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_presence_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -62,7 +62,7 @@ pub async fn get_presence_route( for room_id in db .rooms - .get_shared_rooms(vec![sender_user.clone(), body.user_id.clone()]) + .get_shared_rooms(vec![sender_user.clone(), body.user_id.clone()])? { let room_id = room_id?; diff --git a/src/client_server/profile.rs b/src/client_server/profile.rs index 882b02e7..32bb6083 100644 --- a/src/client_server/profile.rs +++ b/src/client_server/profile.rs @@ -13,7 +13,7 @@ use ruma::{ #[cfg(feature = "conduit_bin")] use rocket::{get, put}; -use std::convert::TryInto; +use std::{convert::TryInto, sync::Arc}; #[cfg_attr( feature = "conduit_bin", @@ -21,7 +21,7 @@ use std::convert::TryInto; )] #[tracing::instrument(skip(db, body))] pub async fn set_displayname_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -107,7 +107,7 @@ pub async fn set_displayname_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_displayname_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { Ok(get_display_name::Response { @@ -122,7 +122,7 @@ pub async fn get_displayname_route( )] #[tracing::instrument(skip(db, body))] pub async fn set_avatar_url_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -208,7 +208,7 @@ pub async fn set_avatar_url_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_avatar_url_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { Ok(get_avatar_url::Response { @@ -223,7 +223,7 @@ pub async fn get_avatar_url_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_profile_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.users.exists(&body.user_id)? { diff --git a/src/client_server/push.rs b/src/client_server/push.rs index e37e660d..d6f62126 100644 --- a/src/client_server/push.rs +++ b/src/client_server/push.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{ConduitResult, Database, Error, Ruma}; use ruma::{ @@ -22,7 +24,7 @@ use rocket::{delete, get, post, put}; )] #[tracing::instrument(skip(db, body))] pub async fn get_pushrules_all_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -47,7 +49,7 @@ pub async fn get_pushrules_all_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_pushrule_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -101,7 +103,7 @@ pub async fn get_pushrule_route( )] #[tracing::instrument(skip(db, req))] pub async fn set_pushrule_route( - db: State<'_, Database>, + db: State<'_, Arc>, req: Ruma>, ) -> ConduitResult { let sender_user = req.sender_user.as_ref().expect("user is authenticated"); @@ -204,7 +206,7 @@ pub async fn set_pushrule_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_pushrule_actions_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -263,7 +265,7 @@ pub async fn get_pushrule_actions_route( )] #[tracing::instrument(skip(db, body))] pub async fn set_pushrule_actions_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -337,7 +339,7 @@ pub async fn set_pushrule_actions_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_pushrule_enabled_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -398,7 +400,7 @@ pub async fn get_pushrule_enabled_route( )] #[tracing::instrument(skip(db, body))] pub async fn set_pushrule_enabled_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -477,7 +479,7 @@ pub async fn set_pushrule_enabled_route( )] #[tracing::instrument(skip(db, body))] pub async fn delete_pushrule_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -546,7 +548,7 @@ pub async fn delete_pushrule_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_pushers_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -563,7 +565,7 @@ pub async fn get_pushers_route( )] #[tracing::instrument(skip(db, body))] pub async fn set_pushers_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/read_marker.rs b/src/client_server/read_marker.rs index 1b7ea0b5..837170ff 100644 --- a/src/client_server/read_marker.rs +++ b/src/client_server/read_marker.rs @@ -12,7 +12,7 @@ use ruma::{ #[cfg(feature = "conduit_bin")] use rocket::post; -use std::collections::BTreeMap; +use std::{collections::BTreeMap, sync::Arc}; #[cfg_attr( feature = "conduit_bin", @@ -20,7 +20,7 @@ use std::collections::BTreeMap; )] #[tracing::instrument(skip(db, body))] pub async fn set_read_marker_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -87,7 +87,7 @@ pub async fn set_read_marker_route( )] #[tracing::instrument(skip(db, body))] pub async fn create_receipt_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/redact.rs b/src/client_server/redact.rs index be5d3b11..e1930823 100644 --- a/src/client_server/redact.rs +++ b/src/client_server/redact.rs @@ -4,6 +4,7 @@ use ruma::{ api::client::r0::redact::redact_event, events::{room::redaction, EventType}, }; +use std::sync::Arc; #[cfg(feature = "conduit_bin")] use rocket::put; @@ -14,7 +15,7 @@ use rocket::put; )] #[tracing::instrument(skip(db, body))] pub async fn redact_event_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/room.rs b/src/client_server/room.rs index 0bc67d4d..7c507750 100644 --- a/src/client_server/room.rs +++ b/src/client_server/room.rs @@ -13,7 +13,7 @@ use ruma::{ serde::Raw, RoomAliasId, RoomId, RoomVersionId, }; -use std::{cmp::max, collections::BTreeMap, convert::TryFrom}; +use std::{cmp::max, collections::BTreeMap, convert::TryFrom, sync::Arc}; #[cfg(feature = "conduit_bin")] use rocket::{get, post}; @@ -24,7 +24,7 @@ use rocket::{get, post}; )] #[tracing::instrument(skip(db, body))] pub async fn create_room_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -304,7 +304,7 @@ pub async fn create_room_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_room_event_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -332,7 +332,7 @@ pub async fn get_room_event_route( )] #[tracing::instrument(skip(db, body))] pub async fn upgrade_room_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, _room_id: String, ) -> ConduitResult { diff --git a/src/client_server/search.rs b/src/client_server/search.rs index a668a0d0..ef5ddc2f 100644 --- a/src/client_server/search.rs +++ b/src/client_server/search.rs @@ -1,6 +1,7 @@ use super::State; use crate::{ConduitResult, Database, Error, Ruma}; use ruma::api::client::{error::ErrorKind, r0::search::search_events}; +use std::sync::Arc; #[cfg(feature = "conduit_bin")] use rocket::post; @@ -13,7 +14,7 @@ use std::collections::BTreeMap; )] #[tracing::instrument(skip(db, body))] pub async fn search_events_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/session.rs b/src/client_server/session.rs index 3718003a..9a75ae28 100644 --- a/src/client_server/session.rs +++ b/src/client_server/session.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::{State, DEVICE_ID_LENGTH, TOKEN_LENGTH}; use crate::{utils, ConduitResult, Database, Error, Ruma}; use log::info; @@ -50,7 +52,7 @@ pub async fn get_login_types_route() -> ConduitResult )] #[tracing::instrument(skip(db, body))] pub async fn login_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { // Validate login method @@ -167,7 +169,7 @@ pub async fn login_route( )] #[tracing::instrument(skip(db, body))] pub async fn logout_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -195,7 +197,7 @@ pub async fn logout_route( )] #[tracing::instrument(skip(db, body))] pub async fn logout_all_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/state.rs b/src/client_server/state.rs index 68e0c7f1..c431ac0d 100644 --- a/src/client_server/state.rs +++ b/src/client_server/state.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{pdu::PduBuilder, ConduitResult, Database, Error, Result, Ruma}; use ruma::{ @@ -25,7 +27,7 @@ use rocket::{get, put}; )] #[tracing::instrument(skip(db, body))] pub async fn send_state_event_for_key_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -51,7 +53,7 @@ pub async fn send_state_event_for_key_route( )] #[tracing::instrument(skip(db, body))] pub async fn send_state_event_for_empty_key_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -77,7 +79,7 @@ pub async fn send_state_event_for_empty_key_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_state_events_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -124,7 +126,7 @@ pub async fn get_state_events_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_state_events_for_key_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -175,7 +177,7 @@ pub async fn get_state_events_for_key_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_state_events_for_empty_key_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/sync.rs b/src/client_server/sync.rs index 63ad590d..2f28706a 100644 --- a/src/client_server/sync.rs +++ b/src/client_server/sync.rs @@ -1,5 +1,5 @@ use super::State; -use crate::{ConduitResult, Database, Error, Ruma}; +use crate::{ConduitResult, Database, Error, Result, Ruma}; use log::error; use ruma::{ api::client::r0::sync::sync_events, @@ -13,6 +13,7 @@ use rocket::{get, tokio}; use std::{ collections::{hash_map, BTreeMap, HashMap, HashSet}, convert::{TryFrom, TryInto}, + sync::Arc, time::Duration, }; @@ -33,7 +34,7 @@ use std::{ )] #[tracing::instrument(skip(db, body))] pub async fn sync_events_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -71,18 +72,23 @@ pub async fn sync_events_route( let mut non_timeline_pdus = db .rooms - .pdus_since(&sender_user, &room_id, since)? + .pdus_until(&sender_user, &room_id, u64::MAX) .filter_map(|r| { + // Filter out buggy events if r.is_err() { error!("Bad pdu in pdus_since: {:?}", r); } r.ok() - }); // Filter out buggy events + }) + .take_while(|(pduid, _)| { + db.rooms + .pdu_count(pduid) + .map_or(false, |count| count > since) + }); // Take the last 10 events for the timeline let timeline_pdus = non_timeline_pdus .by_ref() - .rev() .take(10) .collect::>() .into_iter() @@ -226,7 +232,7 @@ pub async fn sync_events_route( match (since_membership, current_membership) { (MembershipState::Leave, MembershipState::Join) => { // A new user joined an encrypted room - if !share_encrypted_room(&db, &sender_user, &user_id, &room_id) { + if !share_encrypted_room(&db, &sender_user, &user_id, &room_id)? { device_list_updates.insert(user_id); } } @@ -257,6 +263,7 @@ pub async fn sync_events_route( .filter(|user_id| { // Only send keys if the sender doesn't share an encrypted room with the target already !share_encrypted_room(&db, sender_user, user_id, &room_id) + .unwrap_or(false) }), ); } @@ -274,7 +281,7 @@ pub async fn sync_events_route( for hero in db .rooms - .all_pdus(&sender_user, &room_id)? + .all_pdus(&sender_user, &room_id) .filter_map(|pdu| pdu.ok()) // Ignore all broken pdus .filter(|(_, pdu)| pdu.kind == EventType::RoomMember) .map(|(_, pdu)| { @@ -411,7 +418,7 @@ pub async fn sync_events_route( let mut edus = db .rooms .edus - .readreceipts_since(&room_id, since)? + .readreceipts_since(&room_id, since) .filter_map(|r| r.ok()) // Filter out buggy events .map(|(_, _, v)| v) .collect::>(); @@ -549,7 +556,7 @@ pub async fn sync_events_route( for user_id in left_encrypted_users { let still_share_encrypted_room = db .rooms - .get_shared_rooms(vec![sender_user.clone(), user_id.clone()]) + .get_shared_rooms(vec![sender_user.clone(), user_id.clone()])? .filter_map(|r| r.ok()) .filter_map(|other_room_id| { Some( @@ -639,9 +646,10 @@ fn share_encrypted_room( sender_user: &UserId, user_id: &UserId, ignore_room: &RoomId, -) -> bool { - db.rooms - .get_shared_rooms(vec![sender_user.clone(), user_id.clone()]) +) -> Result { + Ok(db + .rooms + .get_shared_rooms(vec![sender_user.clone(), user_id.clone()])? .filter_map(|r| r.ok()) .filter(|room_id| room_id != ignore_room) .filter_map(|other_room_id| { @@ -652,5 +660,5 @@ fn share_encrypted_room( .is_some(), ) }) - .any(|encrypted| encrypted) + .any(|encrypted| encrypted)) } diff --git a/src/client_server/tag.rs b/src/client_server/tag.rs index 63e70ffb..2382fe0a 100644 --- a/src/client_server/tag.rs +++ b/src/client_server/tag.rs @@ -4,7 +4,7 @@ use ruma::{ api::client::r0::tag::{create_tag, delete_tag, get_tags}, events::EventType, }; -use std::collections::BTreeMap; +use std::{collections::BTreeMap, sync::Arc}; #[cfg(feature = "conduit_bin")] use rocket::{delete, get, put}; @@ -15,7 +15,7 @@ use rocket::{delete, get, put}; )] #[tracing::instrument(skip(db, body))] pub async fn update_tag_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -52,7 +52,7 @@ pub async fn update_tag_route( )] #[tracing::instrument(skip(db, body))] pub async fn delete_tag_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); @@ -86,7 +86,7 @@ pub async fn delete_tag_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_tags_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/to_device.rs b/src/client_server/to_device.rs index 460bd057..f2a97abd 100644 --- a/src/client_server/to_device.rs +++ b/src/client_server/to_device.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{ConduitResult, Database, Error, Ruma}; use ruma::api::client::{ @@ -14,7 +16,7 @@ use rocket::put; )] #[tracing::instrument(skip(db, body))] pub async fn send_event_to_device_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/typing.rs b/src/client_server/typing.rs index 4b7feb7f..a0a5d430 100644 --- a/src/client_server/typing.rs +++ b/src/client_server/typing.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{utils, ConduitResult, Database, Ruma}; use create_typing_event::Typing; @@ -12,7 +14,7 @@ use rocket::put; )] #[tracing::instrument(skip(db, body))] pub fn create_typing_event_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let sender_user = body.sender_user.as_ref().expect("user is authenticated"); diff --git a/src/client_server/user_directory.rs b/src/client_server/user_directory.rs index b3582746..0ddc7e85 100644 --- a/src/client_server/user_directory.rs +++ b/src/client_server/user_directory.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::State; use crate::{ConduitResult, Database, Ruma}; use ruma::api::client::r0::user_directory::search_users; @@ -11,7 +13,7 @@ use rocket::post; )] #[tracing::instrument(skip(db, body))] pub async fn search_users_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { let limit = u64::from(body.limit) as usize; diff --git a/src/database.rs b/src/database.rs index 7a55b030..e3b954ef 100644 --- a/src/database.rs +++ b/src/database.rs @@ -1,3 +1,5 @@ +pub mod abstraction; + pub mod account_data; pub mod admin; pub mod appservice; @@ -12,10 +14,10 @@ pub mod uiaa; pub mod users; use crate::{utils, Error, Result}; +use abstraction::DatabaseEngine; use directories::ProjectDirs; -use futures::StreamExt; use log::{error, info}; -use rocket::futures::{self, channel::mpsc}; +use rocket::futures::{channel::mpsc, stream::FuturesUnordered, StreamExt}; use ruma::{DeviceId, ServerName, UserId}; use serde::Deserialize; use std::{ @@ -74,7 +76,8 @@ fn default_log() -> String { "info,state_res=warn,rocket=off,_=off,sled=off".to_owned() } -#[derive(Clone)] +pub type Engine = abstraction::SledEngine; + pub struct Database { pub globals: globals::Globals, pub users: users::Users, @@ -88,7 +91,6 @@ pub struct Database { pub admin: admin::Admin, pub appservice: appservice::Appservice, pub pusher: pusher::PushData, - pub _db: sled::Db, } impl Database { @@ -105,126 +107,126 @@ impl Database { } /// Load an existing database or create a new one. - pub async fn load_or_create(config: Config) -> Result { - let db = sled::Config::default() - .path(&config.database_path) - .cache_capacity(config.cache_capacity as u64) - .use_compression(true) - .open()?; + pub async fn load_or_create(config: Config) -> Result> { + let builder = Engine::open(&config)?; if config.max_request_size < 1024 { eprintln!("ERROR: Max request size is less than 1KB. Please increase it."); } let (admin_sender, admin_receiver) = mpsc::unbounded(); + let (sending_sender, sending_receiver) = mpsc::unbounded(); - let db = Self { + let db = Arc::new(Self { users: users::Users { - userid_password: db.open_tree("userid_password")?, - userid_displayname: db.open_tree("userid_displayname")?, - userid_avatarurl: db.open_tree("userid_avatarurl")?, - userdeviceid_token: db.open_tree("userdeviceid_token")?, - userdeviceid_metadata: db.open_tree("userdeviceid_metadata")?, - userid_devicelistversion: db.open_tree("userid_devicelistversion")?, - token_userdeviceid: db.open_tree("token_userdeviceid")?, - onetimekeyid_onetimekeys: db.open_tree("onetimekeyid_onetimekeys")?, - userid_lastonetimekeyupdate: db.open_tree("userid_lastonetimekeyupdate")?, - keychangeid_userid: db.open_tree("keychangeid_userid")?, - keyid_key: db.open_tree("keyid_key")?, - userid_masterkeyid: db.open_tree("userid_masterkeyid")?, - userid_selfsigningkeyid: db.open_tree("userid_selfsigningkeyid")?, - userid_usersigningkeyid: db.open_tree("userid_usersigningkeyid")?, - todeviceid_events: db.open_tree("todeviceid_events")?, + userid_password: builder.open_tree("userid_password")?, + userid_displayname: builder.open_tree("userid_displayname")?, + userid_avatarurl: builder.open_tree("userid_avatarurl")?, + userdeviceid_token: builder.open_tree("userdeviceid_token")?, + userdeviceid_metadata: builder.open_tree("userdeviceid_metadata")?, + userid_devicelistversion: builder.open_tree("userid_devicelistversion")?, + token_userdeviceid: builder.open_tree("token_userdeviceid")?, + onetimekeyid_onetimekeys: builder.open_tree("onetimekeyid_onetimekeys")?, + userid_lastonetimekeyupdate: builder.open_tree("userid_lastonetimekeyupdate")?, + keychangeid_userid: builder.open_tree("keychangeid_userid")?, + keyid_key: builder.open_tree("keyid_key")?, + userid_masterkeyid: builder.open_tree("userid_masterkeyid")?, + userid_selfsigningkeyid: builder.open_tree("userid_selfsigningkeyid")?, + userid_usersigningkeyid: builder.open_tree("userid_usersigningkeyid")?, + todeviceid_events: builder.open_tree("todeviceid_events")?, }, uiaa: uiaa::Uiaa { - userdevicesessionid_uiaainfo: db.open_tree("userdevicesessionid_uiaainfo")?, - userdevicesessionid_uiaarequest: db.open_tree("userdevicesessionid_uiaarequest")?, + userdevicesessionid_uiaainfo: builder.open_tree("userdevicesessionid_uiaainfo")?, + userdevicesessionid_uiaarequest: builder + .open_tree("userdevicesessionid_uiaarequest")?, }, rooms: rooms::Rooms { edus: rooms::RoomEdus { - readreceiptid_readreceipt: db.open_tree("readreceiptid_readreceipt")?, - roomuserid_privateread: db.open_tree("roomuserid_privateread")?, // "Private" read receipt - roomuserid_lastprivatereadupdate: db + readreceiptid_readreceipt: builder.open_tree("readreceiptid_readreceipt")?, + roomuserid_privateread: builder.open_tree("roomuserid_privateread")?, // "Private" read receipt + roomuserid_lastprivatereadupdate: builder .open_tree("roomuserid_lastprivatereadupdate")?, - typingid_userid: db.open_tree("typingid_userid")?, - roomid_lasttypingupdate: db.open_tree("roomid_lasttypingupdate")?, - presenceid_presence: db.open_tree("presenceid_presence")?, - userid_lastpresenceupdate: db.open_tree("userid_lastpresenceupdate")?, + typingid_userid: builder.open_tree("typingid_userid")?, + roomid_lasttypingupdate: builder.open_tree("roomid_lasttypingupdate")?, + presenceid_presence: builder.open_tree("presenceid_presence")?, + userid_lastpresenceupdate: builder.open_tree("userid_lastpresenceupdate")?, }, - pduid_pdu: db.open_tree("pduid_pdu")?, - eventid_pduid: db.open_tree("eventid_pduid")?, - roomid_pduleaves: db.open_tree("roomid_pduleaves")?, - - alias_roomid: db.open_tree("alias_roomid")?, - aliasid_alias: db.open_tree("aliasid_alias")?, - publicroomids: db.open_tree("publicroomids")?, - - tokenids: db.open_tree("tokenids")?, - - roomserverids: db.open_tree("roomserverids")?, - serverroomids: db.open_tree("serverroomids")?, - userroomid_joined: db.open_tree("userroomid_joined")?, - roomuserid_joined: db.open_tree("roomuserid_joined")?, - roomuseroncejoinedids: db.open_tree("roomuseroncejoinedids")?, - userroomid_invitestate: db.open_tree("userroomid_invitestate")?, - roomuserid_invitecount: db.open_tree("roomuserid_invitecount")?, - userroomid_leftstate: db.open_tree("userroomid_leftstate")?, - roomuserid_leftcount: db.open_tree("roomuserid_leftcount")?, - - userroomid_notificationcount: db.open_tree("userroomid_notificationcount")?, - userroomid_highlightcount: db.open_tree("userroomid_highlightcount")?, - - statekey_shortstatekey: db.open_tree("statekey_shortstatekey")?, - stateid_shorteventid: db.open_tree("stateid_shorteventid")?, - eventid_shorteventid: db.open_tree("eventid_shorteventid")?, - shorteventid_eventid: db.open_tree("shorteventid_eventid")?, - shorteventid_shortstatehash: db.open_tree("shorteventid_shortstatehash")?, - roomid_shortstatehash: db.open_tree("roomid_shortstatehash")?, - statehash_shortstatehash: db.open_tree("statehash_shortstatehash")?, - - eventid_outlierpdu: db.open_tree("eventid_outlierpdu")?, - prevevent_parent: db.open_tree("prevevent_parent")?, + pduid_pdu: builder.open_tree("pduid_pdu")?, + eventid_pduid: builder.open_tree("eventid_pduid")?, + roomid_pduleaves: builder.open_tree("roomid_pduleaves")?, + + alias_roomid: builder.open_tree("alias_roomid")?, + aliasid_alias: builder.open_tree("aliasid_alias")?, + publicroomids: builder.open_tree("publicroomids")?, + + tokenids: builder.open_tree("tokenids")?, + + roomserverids: builder.open_tree("roomserverids")?, + serverroomids: builder.open_tree("serverroomids")?, + userroomid_joined: builder.open_tree("userroomid_joined")?, + roomuserid_joined: builder.open_tree("roomuserid_joined")?, + roomuseroncejoinedids: builder.open_tree("roomuseroncejoinedids")?, + userroomid_invitestate: builder.open_tree("userroomid_invitestate")?, + roomuserid_invitecount: builder.open_tree("roomuserid_invitecount")?, + userroomid_leftstate: builder.open_tree("userroomid_leftstate")?, + roomuserid_leftcount: builder.open_tree("roomuserid_leftcount")?, + + userroomid_notificationcount: builder.open_tree("userroomid_notificationcount")?, + userroomid_highlightcount: builder.open_tree("userroomid_highlightcount")?, + + statekey_shortstatekey: builder.open_tree("statekey_shortstatekey")?, + stateid_shorteventid: builder.open_tree("stateid_shorteventid")?, + eventid_shorteventid: builder.open_tree("eventid_shorteventid")?, + shorteventid_eventid: builder.open_tree("shorteventid_eventid")?, + shorteventid_shortstatehash: builder.open_tree("shorteventid_shortstatehash")?, + roomid_shortstatehash: builder.open_tree("roomid_shortstatehash")?, + statehash_shortstatehash: builder.open_tree("statehash_shortstatehash")?, + + eventid_outlierpdu: builder.open_tree("eventid_outlierpdu")?, + prevevent_parent: builder.open_tree("prevevent_parent")?, }, account_data: account_data::AccountData { - roomuserdataid_accountdata: db.open_tree("roomuserdataid_accountdata")?, + roomuserdataid_accountdata: builder.open_tree("roomuserdataid_accountdata")?, }, media: media::Media { - mediaid_file: db.open_tree("mediaid_file")?, + mediaid_file: builder.open_tree("mediaid_file")?, }, key_backups: key_backups::KeyBackups { - backupid_algorithm: db.open_tree("backupid_algorithm")?, - backupid_etag: db.open_tree("backupid_etag")?, - backupkeyid_backup: db.open_tree("backupkeyid_backup")?, + backupid_algorithm: builder.open_tree("backupid_algorithm")?, + backupid_etag: builder.open_tree("backupid_etag")?, + backupkeyid_backup: builder.open_tree("backupkeyid_backup")?, }, transaction_ids: transaction_ids::TransactionIds { - userdevicetxnid_response: db.open_tree("userdevicetxnid_response")?, + userdevicetxnid_response: builder.open_tree("userdevicetxnid_response")?, }, sending: sending::Sending { - servername_educount: db.open_tree("servername_educount")?, - servernamepduids: db.open_tree("servernamepduids")?, - servercurrentevents: db.open_tree("servercurrentevents")?, + servername_educount: builder.open_tree("servername_educount")?, + servernamepduids: builder.open_tree("servernamepduids")?, + servercurrentevents: builder.open_tree("servercurrentevents")?, maximum_requests: Arc::new(Semaphore::new(config.max_concurrent_requests as usize)), + sender: sending_sender, }, admin: admin::Admin { sender: admin_sender, }, appservice: appservice::Appservice { cached_registrations: Arc::new(RwLock::new(HashMap::new())), - id_appserviceregistrations: db.open_tree("id_appserviceregistrations")?, + id_appserviceregistrations: builder.open_tree("id_appserviceregistrations")?, + }, + pusher: pusher::PushData { + senderkey_pusher: builder.open_tree("senderkey_pusher")?, }, - pusher: pusher::PushData::new(&db)?, globals: globals::Globals::load( - db.open_tree("global")?, - db.open_tree("server_signingkeys")?, + builder.open_tree("global")?, + builder.open_tree("server_signingkeys")?, config, )?, - _db: db, - }; + }); // MIGRATIONS + // TODO: database versions of new dbs should probably not be 0 if db.globals.database_version()? < 1 { - for roomserverid in db.rooms.roomserverids.iter().keys() { - let roomserverid = roomserverid?; + for (roomserverid, _) in db.rooms.roomserverids.iter() { let mut parts = roomserverid.split(|&b| b == 0xff); let room_id = parts.next().expect("split always returns one element"); let servername = match parts.next() { @@ -238,7 +240,7 @@ impl Database { serverroomid.push(0xff); serverroomid.extend_from_slice(room_id); - db.rooms.serverroomids.insert(serverroomid, &[])?; + db.rooms.serverroomids.insert(&serverroomid, &[])?; } db.globals.bump_database_version(1)?; @@ -248,15 +250,13 @@ impl Database { if db.globals.database_version()? < 2 { // We accidentally inserted hashed versions of "" into the db instead of just "" - for userid_password in db.users.userid_password.iter() { - let (userid, password) = userid_password?; - + for (userid, password) in db.users.userid_password.iter() { let password = utils::string_from_bytes(&password); if password.map_or(false, |password| { argon2::verify_encoded(&password, b"").unwrap_or(false) }) { - db.users.userid_password.insert(userid, b"")?; + db.users.userid_password.insert(&userid, b"")?; } } @@ -268,7 +268,8 @@ impl Database { // This data is probably outdated db.rooms.edus.presenceid_presence.clear()?; - db.admin.start_handler(db.clone(), admin_receiver); + db.admin.start_handler(Arc::clone(&db), admin_receiver); + db.sending.start_handler(Arc::clone(&db), sending_receiver); Ok(db) } @@ -282,7 +283,7 @@ impl Database { userdeviceid_prefix.extend_from_slice(device_id.as_bytes()); userdeviceid_prefix.push(0xff); - let mut futures = futures::stream::FuturesUnordered::new(); + let mut futures = FuturesUnordered::new(); // Return when *any* user changed his key // TODO: only send for user they share a room with diff --git a/src/database/abstraction.rs b/src/database/abstraction.rs new file mode 100644 index 00000000..5a2afd5a --- /dev/null +++ b/src/database/abstraction.rs @@ -0,0 +1,309 @@ +use std::{ + collections::BTreeMap, + future::Future, + pin::Pin, + sync::{Arc, RwLock}, +}; + +use log::warn; +use rocksdb::{ + BoundColumnFamily, ColumnFamilyDescriptor, DBWithThreadMode, Direction, MultiThreaded, Options, +}; + +use super::Config; +use crate::{utils, Result}; + +pub struct SledEngine(sled::Db); +pub struct SledEngineTree(sled::Tree); +pub struct RocksDbEngine(rocksdb::DBWithThreadMode); +pub struct RocksDbEngineTree<'a> { + db: Arc, + name: &'a str, + watchers: RwLock, Vec>>>, +} + +pub trait DatabaseEngine: Sized { + fn open(config: &Config) -> Result>; + fn open_tree(self: &Arc, name: &'static str) -> Result>; +} + +pub trait Tree: Send + Sync { + fn get(&self, key: &[u8]) -> Result>>; + + fn insert(&self, key: &[u8], value: &[u8]) -> Result<()>; + + fn remove(&self, key: &[u8]) -> Result<()>; + + fn iter<'a>(&'a self) -> Box, Box<[u8]>)> + Send + Sync + 'a>; + + fn iter_from<'a>( + &'a self, + from: &[u8], + backwards: bool, + ) -> Box, Box<[u8]>)> + 'a>; + + fn increment(&self, key: &[u8]) -> Result>; + + fn scan_prefix<'a>( + &'a self, + prefix: Vec, + ) -> Box, Box<[u8]>)> + 'a>; + + fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin + Send + 'a>>; + + fn clear(&self) -> Result<()> { + for (key, _) in self.iter() { + self.remove(&key)?; + } + + Ok(()) + } +} + +impl DatabaseEngine for SledEngine { + fn open(config: &Config) -> Result> { + Ok(Arc::new(SledEngine( + sled::Config::default() + .path(&config.database_path) + .cache_capacity(config.cache_capacity as u64) + .use_compression(true) + .open()?, + ))) + } + + fn open_tree(self: &Arc, name: &'static str) -> Result> { + Ok(Arc::new(SledEngineTree(self.0.open_tree(name)?))) + } +} + +impl Tree for SledEngineTree { + fn get(&self, key: &[u8]) -> Result>> { + Ok(self.0.get(key)?.map(|v| v.to_vec())) + } + + fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> { + self.0.insert(key, value)?; + Ok(()) + } + + fn remove(&self, key: &[u8]) -> Result<()> { + self.0.remove(key)?; + Ok(()) + } + + fn iter<'a>(&'a self) -> Box, Box<[u8]>)> + Send + Sync + 'a> { + Box::new( + self.0 + .iter() + .filter_map(|r| { + if let Err(e) = &r { + warn!("Error: {}", e); + } + r.ok() + }) + .map(|(k, v)| (k.to_vec().into(), v.to_vec().into())), + ) + } + + fn iter_from( + &self, + from: &[u8], + backwards: bool, + ) -> Box, Box<[u8]>)>> { + let iter = if backwards { + self.0.range(..from) + } else { + self.0.range(from..) + }; + + let iter = iter + .filter_map(|r| { + if let Err(e) = &r { + warn!("Error: {}", e); + } + r.ok() + }) + .map(|(k, v)| (k.to_vec().into(), v.to_vec().into())); + + if backwards { + Box::new(iter.rev()) + } else { + Box::new(iter) + } + } + + fn increment(&self, key: &[u8]) -> Result> { + Ok(self + .0 + .update_and_fetch(key, utils::increment) + .map(|o| o.expect("increment always sets a value").to_vec())?) + } + + fn scan_prefix<'a>( + &'a self, + prefix: Vec, + ) -> Box, Box<[u8]>)> + 'a> { + let iter = self + .0 + .scan_prefix(prefix) + .filter_map(|r| { + if let Err(e) = &r { + warn!("Error: {}", e); + } + r.ok() + }) + .map(|(k, v)| (k.to_vec().into(), v.to_vec().into())); + + Box::new(iter) + } + + fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin + Send + 'a>> { + let prefix = prefix.to_vec(); + Box::pin(async move { + self.0.watch_prefix(prefix).await; + }) + } +} + +impl DatabaseEngine for RocksDbEngine { + fn open(config: &Config) -> Result> { + let mut db_opts = Options::default(); + db_opts.create_if_missing(true); + + let cfs = DBWithThreadMode::::list_cf(&db_opts, &config.database_path) + .unwrap_or_default(); + + let mut options = Options::default(); + options.set_merge_operator_associative("increment", utils::increment_rocksdb); + + let db = DBWithThreadMode::::open_cf_descriptors( + &db_opts, + &config.database_path, + cfs.iter() + .map(|name| ColumnFamilyDescriptor::new(name, options.clone())), + )?; + + Ok(Arc::new(RocksDbEngine(db))) + } + + fn open_tree(self: &Arc, name: &'static str) -> Result> { + let mut options = Options::default(); + options.set_merge_operator_associative("increment", utils::increment_rocksdb); + + // Create if it doesn't exist + let _ = self.0.create_cf(name, &options); + + Ok(Arc::new(RocksDbEngineTree { + name, + db: Arc::clone(self), + watchers: RwLock::new(BTreeMap::new()), + })) + } +} + +impl RocksDbEngineTree<'_> { + fn cf(&self) -> BoundColumnFamily<'_> { + self.db.0.cf_handle(self.name).unwrap() + } +} + +impl Tree for RocksDbEngineTree<'_> { + fn get(&self, key: &[u8]) -> Result>> { + Ok(self.db.0.get_cf(self.cf(), key)?) + } + + fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> { + let watchers = self.watchers.read().unwrap(); + let mut triggered = Vec::new(); + + for length in 0..=key.len() { + if watchers.contains_key(&key[..length]) { + triggered.push(&key[..length]); + } + } + + drop(watchers); + + if !triggered.is_empty() { + let mut watchers = self.watchers.write().unwrap(); + for prefix in triggered { + if let Some(txs) = watchers.remove(prefix) { + for tx in txs { + let _ = tx.send(()); + } + } + } + } + + Ok(self.db.0.put_cf(self.cf(), key, value)?) + } + + fn remove(&self, key: &[u8]) -> Result<()> { + Ok(self.db.0.delete_cf(self.cf(), key)?) + } + + fn iter<'a>(&'a self) -> Box, Box<[u8]>)> + Send + Sync + 'a> { + Box::new( + self.db + .0 + .iterator_cf(self.cf(), rocksdb::IteratorMode::Start), + ) + } + + fn iter_from<'a>( + &'a self, + from: &[u8], + backwards: bool, + ) -> Box, Box<[u8]>)> + 'a> { + Box::new(self.db.0.iterator_cf( + self.cf(), + rocksdb::IteratorMode::From( + from, + if backwards { + Direction::Reverse + } else { + Direction::Forward + }, + ), + )) + } + + fn increment(&self, key: &[u8]) -> Result> { + // TODO: atomic? + let old = self.get(key)?; + let new = utils::increment(old.as_deref()).unwrap(); + self.insert(key, &new)?; + Ok(new) + } + + fn scan_prefix<'a>( + &'a self, + prefix: Vec, + ) -> Box, Box<[u8]>)> + 'a> { + Box::new( + self.db + .0 + .iterator_cf( + self.cf(), + rocksdb::IteratorMode::From(&prefix, Direction::Forward), + ) + .take_while(move |(k, _)| k.starts_with(&prefix)), + ) + } + + fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin + Send + 'a>> { + let (tx, rx) = tokio::sync::oneshot::channel(); + + self.watchers + .write() + .unwrap() + .entry(prefix.to_vec()) + .or_default() + .push(tx); + + Box::pin(async move { + // Tx is never destroyed + rx.await.unwrap(); + }) + } +} diff --git a/src/database/account_data.rs b/src/database/account_data.rs index bb970c34..2ba7bc3d 100644 --- a/src/database/account_data.rs +++ b/src/database/account_data.rs @@ -6,12 +6,12 @@ use ruma::{ RoomId, UserId, }; use serde::{de::DeserializeOwned, Serialize}; -use sled::IVec; -use std::{collections::HashMap, convert::TryFrom}; +use std::{collections::HashMap, convert::TryFrom, sync::Arc}; + +use super::abstraction::Tree; -#[derive(Clone)] pub struct AccountData { - pub(super) roomuserdataid_accountdata: sled::Tree, // RoomUserDataId = Room + User + Count + Type + pub(super) roomuserdataid_accountdata: Arc, // RoomUserDataId = Room + User + Count + Type } impl AccountData { @@ -34,9 +34,8 @@ impl AccountData { prefix.push(0xff); // Remove old entry - if let Some(previous) = self.find_event(room_id, user_id, &event_type) { - let (old_key, _) = previous?; - self.roomuserdataid_accountdata.remove(old_key)?; + if let Some((old_key, _)) = self.find_event(room_id, user_id, &event_type)? { + self.roomuserdataid_accountdata.remove(&old_key)?; } let mut key = prefix; @@ -52,8 +51,10 @@ impl AccountData { )); } - self.roomuserdataid_accountdata - .insert(key, &*json.to_string())?; + self.roomuserdataid_accountdata.insert( + &key, + &serde_json::to_vec(&json).expect("to_vec always works on json values"), + )?; Ok(()) } @@ -65,9 +66,8 @@ impl AccountData { user_id: &UserId, kind: EventType, ) -> Result> { - self.find_event(room_id, user_id, &kind) - .map(|r| { - let (_, v) = r?; + self.find_event(room_id, user_id, &kind)? + .map(|(_, v)| { serde_json::from_slice(&v).map_err(|_| Error::bad_database("could not deserialize")) }) .transpose() @@ -98,8 +98,7 @@ impl AccountData { for r in self .roomuserdataid_accountdata - .range(&*first_possible..) - .filter_map(|r| r.ok()) + .iter_from(&first_possible, false) .take_while(move |(k, _)| k.starts_with(&prefix)) .map(|(k, v)| { Ok::<_, Error>(( @@ -128,7 +127,7 @@ impl AccountData { room_id: Option<&RoomId>, user_id: &UserId, kind: &EventType, - ) -> Option> { + ) -> Result, Box<[u8]>)>> { let mut prefix = room_id .map(|r| r.to_string()) .unwrap_or_default() @@ -137,23 +136,21 @@ impl AccountData { prefix.push(0xff); prefix.extend_from_slice(&user_id.as_bytes()); prefix.push(0xff); + + let mut last_possible_key = prefix.clone(); + last_possible_key.extend_from_slice(&u64::MAX.to_be_bytes()); + let kind = kind.clone(); - self.roomuserdataid_accountdata - .scan_prefix(prefix) - .rev() - .find(move |r| { - r.as_ref() - .map(|(k, _)| { - k.rsplit(|&b| b == 0xff) - .next() - .map(|current_event_type| { - current_event_type == kind.as_ref().as_bytes() - }) - .unwrap_or(false) - }) + Ok(self + .roomuserdataid_accountdata + .iter_from(&last_possible_key, true) + .take_while(move |(k, _)| k.starts_with(&prefix)) + .find(move |(k, _)| { + k.rsplit(|&b| b == 0xff) + .next() + .map(|current_event_type| current_event_type == kind.as_ref().as_bytes()) .unwrap_or(false) - }) - .map(|r| Ok(r?)) + })) } } diff --git a/src/database/admin.rs b/src/database/admin.rs index 30143859..7826cfea 100644 --- a/src/database/admin.rs +++ b/src/database/admin.rs @@ -1,6 +1,9 @@ -use std::convert::{TryFrom, TryInto}; +use std::{ + convert::{TryFrom, TryInto}, + sync::Arc, +}; -use crate::pdu::PduBuilder; +use crate::{pdu::PduBuilder, Database}; use log::warn; use rocket::futures::{channel::mpsc, stream::StreamExt}; use ruma::{ @@ -22,7 +25,7 @@ pub struct Admin { impl Admin { pub fn start_handler( &self, - db: super::Database, + db: Arc, mut receiver: mpsc::UnboundedReceiver, ) { tokio::spawn(async move { @@ -73,14 +76,17 @@ impl Admin { db.appservice.register_appservice(yaml).unwrap(); // TODO handle error } AdminCommand::ListAppservices => { - let appservices = db.appservice.iter_ids().collect::>(); - let count = appservices.len(); - let output = format!( - "Appservices ({}): {}", - count, - appservices.into_iter().filter_map(|r| r.ok()).collect::>().join(", ") - ); - send_message(message::MessageEventContent::text_plain(output)); + if let Ok(appservices) = db.appservice.iter_ids().map(|ids| ids.collect::>()) { + let count = appservices.len(); + let output = format!( + "Appservices ({}): {}", + count, + appservices.into_iter().filter_map(|r| r.ok()).collect::>().join(", ") + ); + send_message(message::MessageEventContent::text_plain(output)); + } else { + send_message(message::MessageEventContent::text_plain("Failed to get appservices.")); + } } AdminCommand::SendMessage(message) => { send_message(message); @@ -93,6 +99,6 @@ impl Admin { } pub fn send(&self, command: AdminCommand) { - self.sender.unbounded_send(command).unwrap() + self.sender.unbounded_send(command).unwrap(); } } diff --git a/src/database/appservice.rs b/src/database/appservice.rs index 222eb182..21b18a76 100644 --- a/src/database/appservice.rs +++ b/src/database/appservice.rs @@ -4,18 +4,21 @@ use std::{ sync::{Arc, RwLock}, }; -#[derive(Clone)] +use super::abstraction::Tree; + pub struct Appservice { pub(super) cached_registrations: Arc>>, - pub(super) id_appserviceregistrations: sled::Tree, + pub(super) id_appserviceregistrations: Arc, } impl Appservice { pub fn register_appservice(&self, yaml: serde_yaml::Value) -> Result<()> { // TODO: Rumaify let id = yaml.get("id").unwrap().as_str().unwrap(); - self.id_appserviceregistrations - .insert(id, serde_yaml::to_string(&yaml).unwrap().as_bytes())?; + self.id_appserviceregistrations.insert( + id.as_bytes(), + serde_yaml::to_string(&yaml).unwrap().as_bytes(), + )?; self.cached_registrations .write() .unwrap() @@ -33,7 +36,7 @@ impl Appservice { || { Ok(self .id_appserviceregistrations - .get(id)? + .get(id.as_bytes())? .map(|bytes| { Ok::<_, Error>(serde_yaml::from_slice(&bytes).map_err(|_| { Error::bad_database( @@ -47,21 +50,25 @@ impl Appservice { ) } - pub fn iter_ids(&self) -> impl Iterator> { - self.id_appserviceregistrations.iter().keys().map(|id| { - Ok(utils::string_from_bytes(&id?).map_err(|_| { + pub fn iter_ids<'a>( + &'a self, + ) -> Result> + Send + Sync + 'a> { + Ok(self.id_appserviceregistrations.iter().map(|(id, _)| { + Ok(utils::string_from_bytes(&id).map_err(|_| { Error::bad_database("Invalid id bytes in id_appserviceregistrations.") })?) - }) + })) } - pub fn iter_all(&self) -> impl Iterator> + '_ { - self.iter_ids().filter_map(|id| id.ok()).map(move |id| { + pub fn iter_all( + &self, + ) -> Result> + '_ + Send + Sync> { + Ok(self.iter_ids()?.filter_map(|id| id.ok()).map(move |id| { Ok(( id.clone(), self.get_registration(&id)? .expect("iter_ids only returns appservices that exist"), )) - }) + })) } } diff --git a/src/database/globals.rs b/src/database/globals.rs index 5d91d374..37ebf136 100644 --- a/src/database/globals.rs +++ b/src/database/globals.rs @@ -13,22 +13,23 @@ use std::{ use tokio::sync::Semaphore; use trust_dns_resolver::TokioAsyncResolver; -pub const COUNTER: &str = "c"; +use super::abstraction::Tree; + +pub const COUNTER: &[u8] = b"c"; type WellKnownMap = HashMap, (String, String)>; type TlsNameMap = HashMap; type RateLimitState = (Instant, u32); // Time if last failed try, number of failed tries -#[derive(Clone)] pub struct Globals { pub actual_destination_cache: Arc>, // actual_destination, host pub tls_name_override: Arc>, - pub(super) globals: sled::Tree, + pub(super) globals: Arc, config: Config, keypair: Arc, reqwest_client: reqwest::Client, dns_resolver: TokioAsyncResolver, jwt_decoding_key: Option>, - pub(super) server_signingkeys: sled::Tree, + pub(super) server_signingkeys: Arc, pub bad_event_ratelimiter: Arc>>, pub bad_signature_ratelimiter: Arc, RateLimitState>>>, pub servername_ratelimiter: Arc, Arc>>>, @@ -69,15 +70,20 @@ impl ServerCertVerifier for MatrixServerVerifier { impl Globals { pub fn load( - globals: sled::Tree, - server_signingkeys: sled::Tree, + globals: Arc, + server_signingkeys: Arc, config: Config, ) -> Result { - let bytes = &*globals - .update_and_fetch("keypair", utils::generate_keypair)? - .expect("utils::generate_keypair always returns Some"); + let keypair_bytes = globals.get(b"keypair")?.map_or_else( + || { + let keypair = utils::generate_keypair(); + globals.insert(b"keypair", &keypair)?; + Ok::<_, Error>(keypair) + }, + |s| Ok(s.to_vec()), + )?; - let mut parts = bytes.splitn(2, |&b| b == 0xff); + let mut parts = keypair_bytes.splitn(2, |&b| b == 0xff); let keypair = utils::string_from_bytes( // 1. version @@ -102,7 +108,7 @@ impl Globals { Ok(k) => k, Err(e) => { error!("Keypair invalid. Deleting..."); - globals.remove("keypair")?; + globals.remove(b"keypair")?; return Err(e); } }; @@ -159,13 +165,8 @@ impl Globals { } pub fn next_count(&self) -> Result { - Ok(utils::u64_from_bytes( - &self - .globals - .update_and_fetch(COUNTER, utils::increment)? - .expect("utils::increment will always put in a value"), - ) - .map_err(|_| Error::bad_database("Count has invalid bytes."))?) + Ok(utils::u64_from_bytes(&self.globals.increment(COUNTER)?) + .map_err(|_| Error::bad_database("Count has invalid bytes."))?) } pub fn current_count(&self) -> Result { @@ -211,21 +212,30 @@ impl Globals { /// Remove the outdated keys and insert the new ones. /// /// This doesn't actually check that the keys provided are newer than the old set. - pub fn add_signing_key(&self, origin: &ServerName, new_keys: &ServerSigningKeys) -> Result<()> { - self.server_signingkeys - .update_and_fetch(origin.as_bytes(), |signingkeys| { - let mut keys = signingkeys - .and_then(|keys| serde_json::from_slice(keys).ok()) - .unwrap_or_else(|| { - // Just insert "now", it doesn't matter - ServerSigningKeys::new(origin.to_owned(), MilliSecondsSinceUnixEpoch::now()) - }); - keys.verify_keys - .extend(new_keys.verify_keys.clone().into_iter()); - keys.old_verify_keys - .extend(new_keys.old_verify_keys.clone().into_iter()); - Some(serde_json::to_vec(&keys).expect("serversigningkeys can be serialized")) - })?; + pub fn add_signing_key(&self, origin: &ServerName, new_keys: ServerSigningKeys) -> Result<()> { + // Not atomic, but this is not critical + let signingkeys = self.server_signingkeys.get(origin.as_bytes())?; + + let mut keys = signingkeys + .and_then(|keys| serde_json::from_slice(&keys).ok()) + .unwrap_or_else(|| { + // Just insert "now", it doesn't matter + ServerSigningKeys::new(origin.to_owned(), MilliSecondsSinceUnixEpoch::now()) + }); + + let ServerSigningKeys { + verify_keys, + old_verify_keys, + .. + } = new_keys; + + keys.verify_keys.extend(verify_keys.into_iter()); + keys.old_verify_keys.extend(old_verify_keys.into_iter()); + + self.server_signingkeys.insert( + origin.as_bytes(), + &serde_json::to_vec(&keys).expect("serversigningkeys can be serialized"), + )?; Ok(()) } @@ -254,14 +264,15 @@ impl Globals { } pub fn database_version(&self) -> Result { - self.globals.get("version")?.map_or(Ok(0), |version| { + self.globals.get(b"version")?.map_or(Ok(0), |version| { utils::u64_from_bytes(&version) .map_err(|_| Error::bad_database("Database version id is invalid.")) }) } pub fn bump_database_version(&self, new_version: u64) -> Result<()> { - self.globals.insert("version", &new_version.to_be_bytes())?; + self.globals + .insert(b"version", &new_version.to_be_bytes())?; Ok(()) } } diff --git a/src/database/key_backups.rs b/src/database/key_backups.rs index 0f9af2eb..0685c482 100644 --- a/src/database/key_backups.rs +++ b/src/database/key_backups.rs @@ -6,13 +6,14 @@ use ruma::{ }, RoomId, UserId, }; -use std::{collections::BTreeMap, convert::TryFrom}; +use std::{collections::BTreeMap, convert::TryFrom, sync::Arc}; + +use super::abstraction::Tree; -#[derive(Clone)] pub struct KeyBackups { - pub(super) backupid_algorithm: sled::Tree, // BackupId = UserId + Version(Count) - pub(super) backupid_etag: sled::Tree, // BackupId = UserId + Version(Count) - pub(super) backupkeyid_backup: sled::Tree, // BackupKeyId = UserId + Version + RoomId + SessionId + pub(super) backupid_algorithm: Arc, // BackupId = UserId + Version(Count) + pub(super) backupid_etag: Arc, // BackupId = UserId + Version(Count) + pub(super) backupkeyid_backup: Arc, // BackupKeyId = UserId + Version + RoomId + SessionId } impl KeyBackups { @@ -30,8 +31,7 @@ impl KeyBackups { self.backupid_algorithm.insert( &key, - &*serde_json::to_string(backup_metadata) - .expect("BackupAlgorithm::to_string always works"), + &serde_json::to_vec(backup_metadata).expect("BackupAlgorithm::to_vec always works"), )?; self.backupid_etag .insert(&key, &globals.next_count()?.to_be_bytes())?; @@ -48,13 +48,8 @@ impl KeyBackups { key.push(0xff); - for outdated_key in self - .backupkeyid_backup - .scan_prefix(&key) - .keys() - .filter_map(|r| r.ok()) - { - self.backupkeyid_backup.remove(outdated_key)?; + for (outdated_key, _) in self.backupkeyid_backup.scan_prefix(key) { + self.backupkeyid_backup.remove(&outdated_key)?; } Ok(()) @@ -80,8 +75,9 @@ impl KeyBackups { self.backupid_algorithm.insert( &key, - &*serde_json::to_string(backup_metadata) - .expect("BackupAlgorithm::to_string always works"), + &serde_json::to_string(backup_metadata) + .expect("BackupAlgorithm::to_string always works") + .as_bytes(), )?; self.backupid_etag .insert(&key, &globals.next_count()?.to_be_bytes())?; @@ -91,11 +87,14 @@ impl KeyBackups { pub fn get_latest_backup(&self, user_id: &UserId) -> Result> { let mut prefix = user_id.as_bytes().to_vec(); prefix.push(0xff); + let mut last_possible_key = prefix.clone(); + last_possible_key.extend_from_slice(&u64::MAX.to_be_bytes()); + self.backupid_algorithm - .scan_prefix(&prefix) - .last() - .map_or(Ok(None), |r| { - let (key, value) = r?; + .iter_from(&last_possible_key, true) + .take_while(move |(k, _)| k.starts_with(&prefix)) + .next() + .map_or(Ok(None), |(key, value)| { let version = utils::string_from_bytes( key.rsplit(|&b| b == 0xff) .next() @@ -117,10 +116,13 @@ impl KeyBackups { key.push(0xff); key.extend_from_slice(version.as_bytes()); - self.backupid_algorithm.get(key)?.map_or(Ok(None), |bytes| { - Ok(serde_json::from_slice(&bytes) - .map_err(|_| Error::bad_database("Algorithm in backupid_algorithm is invalid."))?) - }) + self.backupid_algorithm + .get(&key)? + .map_or(Ok(None), |bytes| { + Ok(serde_json::from_slice(&bytes).map_err(|_| { + Error::bad_database("Algorithm in backupid_algorithm is invalid.") + })?) + }) } pub fn add_key( @@ -153,7 +155,7 @@ impl KeyBackups { self.backupkeyid_backup.insert( &key, - &*serde_json::to_string(&key_data).expect("KeyBackupData::to_string always works"), + &serde_json::to_vec(&key_data).expect("KeyBackupData::to_vec always works"), )?; Ok(()) @@ -164,7 +166,7 @@ impl KeyBackups { prefix.push(0xff); prefix.extend_from_slice(version.as_bytes()); - Ok(self.backupkeyid_backup.scan_prefix(&prefix).count()) + Ok(self.backupkeyid_backup.scan_prefix(prefix).count()) } pub fn get_etag(&self, user_id: &UserId, version: &str) -> Result { @@ -194,33 +196,37 @@ impl KeyBackups { let mut rooms = BTreeMap::::new(); - for result in self.backupkeyid_backup.scan_prefix(&prefix).map(|r| { - let (key, value) = r?; - let mut parts = key.rsplit(|&b| b == 0xff); + for result in self + .backupkeyid_backup + .scan_prefix(prefix) + .map(|(key, value)| { + let mut parts = key.rsplit(|&b| b == 0xff); - let session_id = utils::string_from_bytes( - &parts - .next() - .ok_or_else(|| Error::bad_database("backupkeyid_backup key is invalid."))?, - ) - .map_err(|_| Error::bad_database("backupkeyid_backup session_id is invalid."))?; + let session_id = + utils::string_from_bytes(&parts.next().ok_or_else(|| { + Error::bad_database("backupkeyid_backup key is invalid.") + })?) + .map_err(|_| { + Error::bad_database("backupkeyid_backup session_id is invalid.") + })?; - let room_id = RoomId::try_from( - utils::string_from_bytes( - &parts - .next() - .ok_or_else(|| Error::bad_database("backupkeyid_backup key is invalid."))?, + let room_id = RoomId::try_from( + utils::string_from_bytes(&parts.next().ok_or_else(|| { + Error::bad_database("backupkeyid_backup key is invalid.") + })?) + .map_err(|_| Error::bad_database("backupkeyid_backup room_id is invalid."))?, ) - .map_err(|_| Error::bad_database("backupkeyid_backup room_id is invalid."))?, - ) - .map_err(|_| Error::bad_database("backupkeyid_backup room_id is invalid room id."))?; + .map_err(|_| { + Error::bad_database("backupkeyid_backup room_id is invalid room id.") + })?; - let key_data = serde_json::from_slice(&value).map_err(|_| { - Error::bad_database("KeyBackupData in backupkeyid_backup is invalid.") - })?; + let key_data = serde_json::from_slice(&value).map_err(|_| { + Error::bad_database("KeyBackupData in backupkeyid_backup is invalid.") + })?; - Ok::<_, Error>((room_id, session_id, key_data)) - }) { + Ok::<_, Error>((room_id, session_id, key_data)) + }) + { let (room_id, session_id, key_data) = result?; rooms .entry(room_id) @@ -239,7 +245,7 @@ impl KeyBackups { user_id: &UserId, version: &str, room_id: &RoomId, - ) -> BTreeMap { + ) -> Result> { let mut prefix = user_id.as_bytes().to_vec(); prefix.push(0xff); prefix.extend_from_slice(version.as_bytes()); @@ -247,10 +253,10 @@ impl KeyBackups { prefix.extend_from_slice(room_id.as_bytes()); prefix.push(0xff); - self.backupkeyid_backup - .scan_prefix(&prefix) - .map(|r| { - let (key, value) = r?; + Ok(self + .backupkeyid_backup + .scan_prefix(prefix) + .map(|(key, value)| { let mut parts = key.rsplit(|&b| b == 0xff); let session_id = @@ -268,7 +274,7 @@ impl KeyBackups { Ok::<_, Error>((session_id, key_data)) }) .filter_map(|r| r.ok()) - .collect() + .collect()) } pub fn get_session( @@ -302,13 +308,8 @@ impl KeyBackups { key.extend_from_slice(&version.as_bytes()); key.push(0xff); - for outdated_key in self - .backupkeyid_backup - .scan_prefix(&key) - .keys() - .filter_map(|r| r.ok()) - { - self.backupkeyid_backup.remove(outdated_key)?; + for (outdated_key, _) in self.backupkeyid_backup.scan_prefix(key) { + self.backupkeyid_backup.remove(&outdated_key)?; } Ok(()) @@ -327,13 +328,8 @@ impl KeyBackups { key.extend_from_slice(&room_id.as_bytes()); key.push(0xff); - for outdated_key in self - .backupkeyid_backup - .scan_prefix(&key) - .keys() - .filter_map(|r| r.ok()) - { - self.backupkeyid_backup.remove(outdated_key)?; + for (outdated_key, _) in self.backupkeyid_backup.scan_prefix(key) { + self.backupkeyid_backup.remove(&outdated_key)?; } Ok(()) @@ -354,13 +350,8 @@ impl KeyBackups { key.push(0xff); key.extend_from_slice(&session_id.as_bytes()); - for outdated_key in self - .backupkeyid_backup - .scan_prefix(&key) - .keys() - .filter_map(|r| r.ok()) - { - self.backupkeyid_backup.remove(outdated_key)?; + for (outdated_key, _) in self.backupkeyid_backup.scan_prefix(key) { + self.backupkeyid_backup.remove(&outdated_key)?; } Ok(()) diff --git a/src/database/media.rs b/src/database/media.rs index 28ef88a2..ca45484a 100644 --- a/src/database/media.rs +++ b/src/database/media.rs @@ -1,7 +1,9 @@ use image::{imageops::FilterType, GenericImageView}; use crate::{utils, Error, Result}; -use std::mem; +use std::{mem, sync::Arc}; + +use super::abstraction::Tree; pub struct FileMeta { pub content_disposition: Option, @@ -9,9 +11,8 @@ pub struct FileMeta { pub file: Vec, } -#[derive(Clone)] pub struct Media { - pub(super) mediaid_file: sled::Tree, // MediaId = MXC + WidthHeight + ContentDisposition + ContentType + pub(super) mediaid_file: Arc, // MediaId = MXC + WidthHeight + ContentDisposition + ContentType } impl Media { @@ -42,7 +43,7 @@ impl Media { .unwrap_or_default(), ); - self.mediaid_file.insert(key, file)?; + self.mediaid_file.insert(&key, file)?; Ok(()) } @@ -76,7 +77,7 @@ impl Media { .unwrap_or_default(), ); - self.mediaid_file.insert(key, file)?; + self.mediaid_file.insert(&key, file)?; Ok(()) } @@ -89,8 +90,7 @@ impl Media { prefix.extend_from_slice(&0_u32.to_be_bytes()); // Height = 0 if it's not a thumbnail prefix.push(0xff); - if let Some(r) = self.mediaid_file.scan_prefix(&prefix).next() { - let (key, file) = r?; + if let Some((key, file)) = self.mediaid_file.scan_prefix(prefix).next() { let mut parts = key.rsplit(|&b| b == 0xff); let content_type = parts @@ -169,9 +169,8 @@ impl Media { original_prefix.extend_from_slice(&0_u32.to_be_bytes()); // Height = 0 if it's not a thumbnail original_prefix.push(0xff); - if let Some(r) = self.mediaid_file.scan_prefix(&thumbnail_prefix).next() { + if let Some((key, file)) = self.mediaid_file.scan_prefix(thumbnail_prefix).next() { // Using saved thumbnail - let (key, file) = r?; let mut parts = key.rsplit(|&b| b == 0xff); let content_type = parts @@ -202,10 +201,8 @@ impl Media { content_type, file: file.to_vec(), })) - } else if let Some(r) = self.mediaid_file.scan_prefix(&original_prefix).next() { + } else if let Some((key, file)) = self.mediaid_file.scan_prefix(original_prefix).next() { // Generate a thumbnail - - let (key, file) = r?; let mut parts = key.rsplit(|&b| b == 0xff); let content_type = parts @@ -302,7 +299,7 @@ impl Media { widthheight, ); - self.mediaid_file.insert(thumbnail_key, &*thumbnail_bytes)?; + self.mediaid_file.insert(&thumbnail_key, &thumbnail_bytes)?; Ok(Some(FileMeta { content_disposition, diff --git a/src/database/pusher.rs b/src/database/pusher.rs index 51f55a17..39b631dc 100644 --- a/src/database/pusher.rs +++ b/src/database/pusher.rs @@ -14,23 +14,17 @@ use ruma::{ push::{Action, PushConditionRoomCtx, PushFormat, Ruleset, Tweak}, uint, UInt, UserId, }; -use sled::IVec; -use std::{convert::TryFrom, fmt::Debug, mem}; +use std::{convert::TryFrom, fmt::Debug, mem, sync::Arc}; + +use super::abstraction::Tree; -#[derive(Debug, Clone)] pub struct PushData { /// UserId + pushkey -> Pusher - pub(super) senderkey_pusher: sled::Tree, + pub(super) senderkey_pusher: Arc, } impl PushData { - pub fn new(db: &sled::Db) -> Result { - Ok(Self { - senderkey_pusher: db.open_tree("senderkey_pusher")?, - }) - } - pub fn set_pusher(&self, sender: &UserId, pusher: set_pusher::Pusher) -> Result<()> { let mut key = sender.as_bytes().to_vec(); key.push(0xff); @@ -40,14 +34,14 @@ impl PushData { if pusher.kind.is_none() { return self .senderkey_pusher - .remove(key) + .remove(&key) .map(|_| ()) .map_err(Into::into); } self.senderkey_pusher.insert( - key, - &*serde_json::to_string(&pusher).expect("Pusher is valid JSON string"), + &key, + &serde_json::to_vec(&pusher).expect("Pusher is valid JSON value"), )?; Ok(()) @@ -69,23 +63,21 @@ impl PushData { self.senderkey_pusher .scan_prefix(prefix) - .values() - .map(|push| { - let push = push.map_err(|_| Error::bad_database("Invalid push bytes in db."))?; + .map(|(_, push)| { Ok(serde_json::from_slice(&*push) .map_err(|_| Error::bad_database("Invalid Pusher in db."))?) }) .collect() } - pub fn get_pusher_senderkeys(&self, sender: &UserId) -> impl Iterator> { + pub fn get_pusher_senderkeys<'a>( + &'a self, + sender: &UserId, + ) -> impl Iterator> + 'a { let mut prefix = sender.as_bytes().to_vec(); prefix.push(0xff); - self.senderkey_pusher - .scan_prefix(prefix) - .keys() - .map(|r| Ok(r?)) + self.senderkey_pusher.scan_prefix(prefix).map(|(k, _)| k) } } diff --git a/src/database/rooms.rs b/src/database/rooms.rs index 703314e6..0a8239d4 100644 --- a/src/database/rooms.rs +++ b/src/database/rooms.rs @@ -28,7 +28,7 @@ use std::{ sync::Arc, }; -use super::{admin::AdminCommand, pusher}; +use super::{abstraction::Tree, admin::AdminCommand, pusher}; /// The unique identifier of each state group. /// @@ -36,54 +36,53 @@ use super::{admin::AdminCommand, pusher}; /// hashing the entire state. pub type StateHashId = IVec; -#[derive(Clone)] pub struct Rooms { pub edus: edus::RoomEdus, - pub(super) pduid_pdu: sled::Tree, // PduId = RoomId + Count - pub(super) eventid_pduid: sled::Tree, - pub(super) roomid_pduleaves: sled::Tree, - pub(super) alias_roomid: sled::Tree, - pub(super) aliasid_alias: sled::Tree, // AliasId = RoomId + Count - pub(super) publicroomids: sled::Tree, + pub(super) pduid_pdu: Arc, // PduId = RoomId + Count + pub(super) eventid_pduid: Arc, + pub(super) roomid_pduleaves: Arc, + pub(super) alias_roomid: Arc, + pub(super) aliasid_alias: Arc, // AliasId = RoomId + Count + pub(super) publicroomids: Arc, - pub(super) tokenids: sled::Tree, // TokenId = RoomId + Token + PduId + pub(super) tokenids: Arc, // TokenId = RoomId + Token + PduId /// Participating servers in a room. - pub(super) roomserverids: sled::Tree, // RoomServerId = RoomId + ServerName - pub(super) serverroomids: sled::Tree, // ServerRoomId = ServerName + RoomId + pub(super) roomserverids: Arc, // RoomServerId = RoomId + ServerName + pub(super) serverroomids: Arc, // ServerRoomId = ServerName + RoomId - pub(super) userroomid_joined: sled::Tree, - pub(super) roomuserid_joined: sled::Tree, - pub(super) roomuseroncejoinedids: sled::Tree, - pub(super) userroomid_invitestate: sled::Tree, // InviteState = Vec> - pub(super) roomuserid_invitecount: sled::Tree, // InviteCount = Count - pub(super) userroomid_leftstate: sled::Tree, - pub(super) roomuserid_leftcount: sled::Tree, + pub(super) userroomid_joined: Arc, + pub(super) roomuserid_joined: Arc, + pub(super) roomuseroncejoinedids: Arc, + pub(super) userroomid_invitestate: Arc, // InviteState = Vec> + pub(super) roomuserid_invitecount: Arc, // InviteCount = Count + pub(super) userroomid_leftstate: Arc, + pub(super) roomuserid_leftcount: Arc, - pub(super) userroomid_notificationcount: sled::Tree, // NotifyCount = u64 - pub(super) userroomid_highlightcount: sled::Tree, // HightlightCount = u64 + pub(super) userroomid_notificationcount: Arc, // NotifyCount = u64 + pub(super) userroomid_highlightcount: Arc, // HightlightCount = u64 /// Remember the current state hash of a room. - pub(super) roomid_shortstatehash: sled::Tree, + pub(super) roomid_shortstatehash: Arc, /// Remember the state hash at events in the past. - pub(super) shorteventid_shortstatehash: sled::Tree, + pub(super) shorteventid_shortstatehash: Arc, /// StateKey = EventType + StateKey, ShortStateKey = Count - pub(super) statekey_shortstatekey: sled::Tree, - pub(super) shorteventid_eventid: sled::Tree, + pub(super) statekey_shortstatekey: Arc, + pub(super) shorteventid_eventid: Arc, /// ShortEventId = Count - pub(super) eventid_shorteventid: sled::Tree, + pub(super) eventid_shorteventid: Arc, /// ShortEventId = Count - pub(super) statehash_shortstatehash: sled::Tree, + pub(super) statehash_shortstatehash: Arc, /// ShortStateHash = Count /// StateId = ShortStateHash + ShortStateKey - pub(super) stateid_shorteventid: sled::Tree, + pub(super) stateid_shorteventid: Arc, /// RoomId + EventId -> outlier PDU. /// Any pdu that has passed the steps 1-8 in the incoming event /federation/send/txn. - pub(super) eventid_outlierpdu: sled::Tree, + pub(super) eventid_outlierpdu: Arc, /// RoomId + EventId -> Parent PDU EventId. - pub(super) prevevent_parent: sled::Tree, + pub(super) prevevent_parent: Arc, } impl Rooms { @@ -92,10 +91,8 @@ impl Rooms { pub fn state_full_ids(&self, shortstatehash: u64) -> Result> { Ok(self .stateid_shorteventid - .scan_prefix(&shortstatehash.to_be_bytes()) - .values() - .filter_map(|r| r.ok()) - .map(|bytes| self.shorteventid_eventid.get(&bytes).ok().flatten()) + .scan_prefix(shortstatehash.to_be_bytes().to_vec()) + .map(|(_, bytes)| self.shorteventid_eventid.get(&bytes).ok().flatten()) .flatten() .map(|bytes| { Ok::<_, Error>( @@ -117,10 +114,8 @@ impl Rooms { ) -> Result> { Ok(self .stateid_shorteventid - .scan_prefix(shortstatehash.to_be_bytes()) - .values() - .filter_map(|r| r.ok()) - .map(|bytes| self.shorteventid_eventid.get(&bytes).ok().flatten()) + .scan_prefix(shortstatehash.to_be_bytes().to_vec()) + .map(|(_, bytes)| self.shorteventid_eventid.get(&bytes).ok().flatten()) .flatten() .map(|bytes| { Ok::<_, Error>( @@ -211,16 +206,16 @@ impl Rooms { self.eventid_shorteventid .get(event_id.as_bytes())? .map_or(Ok(None), |shorteventid| { - Ok(self.shorteventid_shortstatehash.get(shorteventid)?.map_or( - Ok::<_, Error>(None), - |bytes| { + Ok(self + .shorteventid_shortstatehash + .get(&shorteventid)? + .map_or(Ok::<_, Error>(None), |bytes| { Ok(Some(utils::u64_from_bytes(&bytes).map_err(|_| { Error::bad_database( "Invalid shortstatehash bytes in shorteventid_shortstatehash", ) })?)) - }, - )?) + })?) }) } @@ -285,7 +280,8 @@ impl Rooms { // Look for PDUs in that room. Ok(self .pduid_pdu - .get_gt(&prefix)? + .iter_from(&prefix, false) + .next() .filter(|(k, _)| k.starts_with(&prefix)) .is_some()) } @@ -471,10 +467,17 @@ impl Rooms { } pub fn latest_pdu_count(&self, room_id: &RoomId) -> Result { + let mut prefix = room_id.as_bytes().to_vec(); + prefix.push(0xff); + + let mut last_possible_key = prefix.clone(); + last_possible_key.extend_from_slice(&u64::MAX.to_be_bytes()); + self.pduid_pdu - .scan_prefix(room_id.as_bytes()) - .last() - .map(|b| self.pdu_count(&b?.0)) + .iter_from(&last_possible_key, true) + .take_while(move |(k, _)| k.starts_with(&prefix)) + .next() + .map(|b| self.pdu_count(&b.0)) .transpose() .map(|op| op.unwrap_or_default()) } @@ -499,7 +502,7 @@ impl Rooms { } /// Returns the pdu's id. - pub fn get_pdu_id(&self, event_id: &EventId) -> Result> { + pub fn get_pdu_id(&self, event_id: &EventId) -> Result>> { self.eventid_pduid .get(event_id.as_bytes())? .map_or(Ok(None), |pdu_id| Ok(Some(pdu_id))) @@ -570,11 +573,11 @@ impl Rooms { } /// Removes a pdu and creates a new one with the same id. - fn replace_pdu(&self, pdu_id: &IVec, pdu: &PduEvent) -> Result<()> { + fn replace_pdu(&self, pdu_id: &[u8], pdu: &PduEvent) -> Result<()> { if self.pduid_pdu.get(&pdu_id)?.is_some() { self.pduid_pdu.insert( &pdu_id, - &*serde_json::to_string(pdu).expect("PduEvent::to_string always works"), + &serde_json::to_vec(pdu).expect("PduEvent::to_vec always works"), )?; Ok(()) } else { @@ -591,11 +594,11 @@ impl Rooms { prefix.push(0xff); self.roomid_pduleaves - .scan_prefix(prefix) - .values() - .map(|bytes| { + .scan_prefix(dbg!(prefix)) + .map(|(key, bytes)| { + dbg!(key); Ok::<_, Error>( - EventId::try_from(utils::string_from_bytes(&bytes?).map_err(|_| { + EventId::try_from(utils::string_from_bytes(&bytes).map_err(|_| { Error::bad_database("EventID in roomid_pduleaves is invalid unicode.") })?) .map_err(|_| Error::bad_database("EventId in roomid_pduleaves is invalid."))?, @@ -612,8 +615,8 @@ impl Rooms { let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); - for key in self.roomid_pduleaves.scan_prefix(&prefix).keys() { - self.roomid_pduleaves.remove(key?)?; + for (key, _) in self.roomid_pduleaves.scan_prefix(prefix.clone()) { + self.roomid_pduleaves.remove(&key)?; } for event_id in event_ids { @@ -628,7 +631,7 @@ impl Rooms { pub fn is_pdu_referenced(&self, pdu: &PduEvent) -> Result { let mut key = pdu.room_id().as_bytes().to_vec(); key.extend_from_slice(pdu.event_id().as_bytes()); - self.prevevent_parent.contains_key(key).map_err(Into::into) + Ok(self.prevevent_parent.get(&key)?.is_some()) } /// Returns the pdu from the outlier tree. @@ -646,7 +649,7 @@ impl Rooms { pub fn add_pdu_outlier(&self, event_id: &EventId, pdu: &CanonicalJsonObject) -> Result<()> { self.eventid_outlierpdu.insert( &event_id.as_bytes(), - &*serde_json::to_string(&pdu).expect("CanonicalJsonObject is valid string"), + &serde_json::to_vec(&pdu).expect("CanonicalJsonObject is valid"), )?; Ok(()) @@ -698,7 +701,7 @@ impl Rooms { let mut key = pdu.room_id().as_bytes().to_vec(); key.extend_from_slice(leaf.as_bytes()); self.prevevent_parent - .insert(key, pdu.event_id().as_bytes())?; + .insert(&key, pdu.event_id().as_bytes())?; } self.replace_pdu_leaves(&pdu.room_id, leaves)?; @@ -711,8 +714,7 @@ impl Rooms { self.pduid_pdu.insert( &pdu_id, - &*serde_json::to_string(&pdu_json) - .expect("CanonicalJsonObject is always a valid String"), + &serde_json::to_vec(&pdu_json).expect("CanonicalJsonObject is always a valid"), )?; // This also replaces the eventid of any outliers with the correct @@ -760,22 +762,14 @@ impl Rooms { userroom_id.extend_from_slice(pdu.room_id.as_bytes()); if notify { - self.userroomid_notificationcount - .update_and_fetch(&userroom_id, utils::increment)? - .expect("utils::increment will always put in a value"); + self.userroomid_notificationcount.increment(&userroom_id)?; } if highlight { - self.userroomid_highlightcount - .update_and_fetch(&userroom_id, utils::increment)? - .expect("utils::increment will always put in a value"); + self.userroomid_highlightcount.increment(&userroom_id)?; } - for senderkey in db - .pusher - .get_pusher_senderkeys(&user) - .filter_map(|r| r.ok()) - { + for senderkey in db.pusher.get_pusher_senderkeys(&user) { db.sending.send_push_pdu(&*pdu_id, senderkey)?; } } @@ -840,7 +834,7 @@ impl Rooms { key.extend_from_slice(word.as_bytes()); key.push(0xff); key.extend_from_slice(&pdu_id); - self.tokenids.insert(key, &[])?; + self.tokenids.insert(&key, &[])?; } if body.starts_with(&format!("@conduit:{}: ", db.globals.server_name())) @@ -991,7 +985,7 @@ impl Rooms { Some(shortstatehash) => { // State already existed in db self.shorteventid_shortstatehash - .insert(shorteventid, &*shortstatehash)?; + .insert(&shorteventid, &*shortstatehash)?; return Ok(()); } None => { @@ -1037,7 +1031,7 @@ impl Rooms { } self.shorteventid_shortstatehash - .insert(shorteventid, &*shortstatehash)?; + .insert(&shorteventid, &*shortstatehash)?; Ok(()) } @@ -1070,7 +1064,7 @@ impl Rooms { }; self.shorteventid_shortstatehash - .insert(shorteventid, &old_shortstatehash)?; + .insert(&shorteventid, &old_shortstatehash)?; if new_pdu.state_key.is_none() { return utils::u64_from_bytes(&old_shortstatehash).map_err(|_| { Error::bad_database("Invalid shortstatehash in roomid_shortstatehash.") @@ -1078,17 +1072,16 @@ impl Rooms { } self.stateid_shorteventid - .scan_prefix(&old_shortstatehash) - .filter_map(|pdu| pdu.map_err(|e| error!("{}", e)).ok()) + .scan_prefix(old_shortstatehash.clone()) // Chop the old_shortstatehash out leaving behind the short state key .map(|(k, v)| (k[old_shortstatehash.len()..].to_vec(), v)) - .collect::, IVec>>() + .collect::, Box<[u8]>>>() } else { HashMap::new() }; if let Some(state_key) = &new_pdu.state_key { - let mut new_state: HashMap, IVec> = old_state; + let mut new_state: HashMap, Box<[u8]>> = old_state; let mut new_state_key = new_pdu.kind.as_ref().as_bytes().to_vec(); new_state_key.push(0xff); @@ -1205,6 +1198,7 @@ impl Rooms { room_id: &RoomId, db: &Database, ) -> Result { + dbg!(&pdu_builder); let PduBuilder { event_type, content, @@ -1385,7 +1379,7 @@ impl Rooms { db.sending.send_pdu(&server, &pdu_id)?; } - for appservice in db.appservice.iter_all().filter_map(|r| r.ok()) { + for appservice in db.appservice.iter_all()?.filter_map(|r| r.ok()) { if let Some(namespaces) = appservice.1.get("namespaces") { let users = namespaces .get("users") @@ -1464,23 +1458,23 @@ impl Rooms { /// Returns an iterator over all PDUs in a room. #[tracing::instrument(skip(self))] - pub fn all_pdus( - &self, + pub fn all_pdus<'a>( + &'a self, user_id: &UserId, room_id: &RoomId, - ) -> Result>> { + ) -> impl Iterator, PduEvent)>> + 'a { self.pdus_since(user_id, room_id, 0) } - /// Returns a double-ended iterator over all events in a room that happened after the event with id `since` + /// Returns an iterator over all events in a room that happened after the event with id `since` /// in chronological order. #[tracing::instrument(skip(self))] - pub fn pdus_since( - &self, + pub fn pdus_since<'a>( + &'a self, user_id: &UserId, room_id: &RoomId, since: u64, - ) -> Result>> { + ) -> impl Iterator, PduEvent)>> + 'a { let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); @@ -1488,19 +1482,10 @@ impl Rooms { let mut first_pdu_id = prefix.clone(); first_pdu_id.extend_from_slice(&(since + 1).to_be_bytes()); - let mut last_pdu_id = prefix; - last_pdu_id.extend_from_slice(&u64::MAX.to_be_bytes()); - let user_id = user_id.clone(); - Ok(self - .pduid_pdu - .range(first_pdu_id..last_pdu_id) - .filter_map(|r| { - if r.is_err() { - error!("Bad pdu in pduid_pdu: {:?}", r); - } - r.ok() - }) + self.pduid_pdu + .iter_from(&first_pdu_id, false) + .take_while(move |(k, _)| k.starts_with(&prefix)) .map(move |(pdu_id, v)| { let mut pdu = serde_json::from_slice::(&v) .map_err(|_| Error::bad_database("PDU in db is invalid."))?; @@ -1508,17 +1493,17 @@ impl Rooms { pdu.unsigned.remove("transaction_id"); } Ok((pdu_id, pdu)) - })) + }) } /// Returns an iterator over all events and their tokens in a room that happened before the /// event with id `until` in reverse-chronological order. - pub fn pdus_until( - &self, + pub fn pdus_until<'a>( + &'a self, user_id: &UserId, room_id: &RoomId, until: u64, - ) -> impl Iterator> { + ) -> impl Iterator, PduEvent)>> + 'a { // Create the first part of the full pdu id let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); @@ -1530,9 +1515,7 @@ impl Rooms { let user_id = user_id.clone(); self.pduid_pdu - .range(..current) - .rev() - .filter_map(|r| r.ok()) + .iter_from(current, true) .take_while(move |(k, _)| k.starts_with(&prefix)) .map(move |(pdu_id, v)| { let mut pdu = serde_json::from_slice::(&v) @@ -1547,12 +1530,12 @@ impl Rooms { /// Returns an iterator over all events and their token in a room that happened after the event /// with id `from` in chronological order. #[tracing::instrument(skip(self))] - pub fn pdus_after( - &self, + pub fn pdus_after<'a>( + &'a self, user_id: &UserId, room_id: &RoomId, from: u64, - ) -> impl Iterator> { + ) -> impl Iterator, PduEvent)>> + 'a { // Create the first part of the full pdu id let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); @@ -1564,8 +1547,7 @@ impl Rooms { let user_id = user_id.clone(); self.pduid_pdu - .range(current..) - .filter_map(|r| r.ok()) + .iter_from(current, false) .take_while(move |(k, _)| k.starts_with(&prefix)) .map(move |(pdu_id, v)| { let mut pdu = serde_json::from_slice::(&v) @@ -1744,7 +1726,7 @@ impl Rooms { self.serverroomids.insert(&serverroom_id, &[])?; self.userroomid_invitestate.insert( &userroom_id, - serde_json::to_vec(&last_state.unwrap_or_default()) + &serde_json::to_vec(&last_state.unwrap_or_default()) .expect("state to bytes always works"), )?; self.roomuserid_invitecount @@ -1766,7 +1748,7 @@ impl Rooms { } self.userroomid_leftstate.insert( &userroom_id, - serde_json::to_vec(&Vec::>::new()).unwrap(), + &serde_json::to_vec(&Vec::>::new()).unwrap(), )?; // TODO self.roomuserid_leftcount .insert(&roomuser_id, &db.globals.next_count()?.to_be_bytes())?; @@ -1966,8 +1948,8 @@ impl Rooms { roomuser_id.push(0xff); roomuser_id.extend_from_slice(user_id.as_bytes()); - self.userroomid_leftstate.remove(userroom_id)?; - self.roomuserid_leftcount.remove(roomuser_id)?; + self.userroomid_leftstate.remove(&userroom_id)?; + self.roomuserid_leftcount.remove(&roomuser_id)?; Ok(()) } @@ -1981,26 +1963,26 @@ impl Rooms { if let Some(room_id) = room_id { // New alias self.alias_roomid - .insert(alias.alias(), room_id.as_bytes())?; + .insert(&alias.alias().as_bytes(), room_id.as_bytes())?; let mut aliasid = room_id.as_bytes().to_vec(); aliasid.push(0xff); aliasid.extend_from_slice(&globals.next_count()?.to_be_bytes()); - self.aliasid_alias.insert(aliasid, &*alias.as_bytes())?; + self.aliasid_alias.insert(&aliasid, &*alias.as_bytes())?; } else { // room_id=None means remove alias - let room_id = self - .alias_roomid - .remove(alias.alias())? - .ok_or(Error::BadRequest( + if let Some(room_id) = self.alias_roomid.get(&alias.alias().as_bytes())? { + let mut prefix = room_id.to_vec(); + prefix.push(0xff); + + for (key, _) in self.aliasid_alias.scan_prefix(prefix) { + self.aliasid_alias.remove(&key)?; + } + self.alias_roomid.remove(&alias.alias().as_bytes())?; + } else { + return Err(Error::BadRequest( ErrorKind::NotFound, "Alias does not exist.", - ))?; - - let mut prefix = room_id.to_vec(); - prefix.push(0xff); - - for key in self.aliasid_alias.scan_prefix(prefix).keys() { - self.aliasid_alias.remove(key?)?; + )); } } @@ -2009,7 +1991,7 @@ impl Rooms { pub fn id_from_alias(&self, alias: &RoomAliasId) -> Result> { self.alias_roomid - .get(alias.alias())? + .get(alias.alias().as_bytes())? .map_or(Ok(None), |bytes| { Ok(Some( RoomId::try_from(utils::string_from_bytes(&bytes).map_err(|_| { @@ -2020,19 +2002,19 @@ impl Rooms { }) } - pub fn room_aliases(&self, room_id: &RoomId) -> impl Iterator> { + pub fn room_aliases<'a>( + &'a self, + room_id: &RoomId, + ) -> impl Iterator> + 'a { let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); - self.aliasid_alias - .scan_prefix(prefix) - .values() - .map(|bytes| { - Ok(utils::string_from_bytes(&bytes?) - .map_err(|_| Error::bad_database("Invalid alias bytes in aliasid_alias."))? - .try_into() - .map_err(|_| Error::bad_database("Invalid alias in aliasid_alias."))?) - }) + self.aliasid_alias.scan_prefix(prefix).map(|(_, bytes)| { + Ok(utils::string_from_bytes(&bytes) + .map_err(|_| Error::bad_database("Invalid alias bytes in aliasid_alias."))? + .try_into() + .map_err(|_| Error::bad_database("Invalid alias in aliasid_alias."))?) + }) } pub fn set_public(&self, room_id: &RoomId, public: bool) -> Result<()> { @@ -2046,13 +2028,13 @@ impl Rooms { } pub fn is_public_room(&self, room_id: &RoomId) -> Result { - Ok(self.publicroomids.contains_key(room_id.as_bytes())?) + Ok(self.publicroomids.get(room_id.as_bytes())?.is_some()) } - pub fn public_rooms(&self) -> impl Iterator> { - self.publicroomids.iter().keys().map(|bytes| { + pub fn public_rooms<'a>(&'a self) -> impl Iterator> + 'a { + self.publicroomids.iter().map(|(bytes, _)| { Ok( - RoomId::try_from(utils::string_from_bytes(&bytes?).map_err(|_| { + RoomId::try_from(utils::string_from_bytes(&bytes).map_err(|_| { Error::bad_database("Room ID in publicroomids is invalid unicode.") })?) .map_err(|_| Error::bad_database("Room ID in publicroomids is invalid."))?, @@ -2073,31 +2055,39 @@ impl Rooms { .map(str::to_lowercase) .collect::>(); - let iterators = words.clone().into_iter().map(move |word| { - let mut prefix2 = prefix.clone(); - prefix2.extend_from_slice(word.as_bytes()); - prefix2.push(0xff); - self.tokenids - .scan_prefix(&prefix2) - .keys() - .rev() // Newest pdus first - .filter_map(|r| r.ok()) - .map(|key| { - let pduid_index = key - .iter() - .enumerate() - .filter(|(_, &b)| b == 0xff) - .nth(1) - .ok_or_else(|| Error::bad_database("Invalid tokenid in db."))? - .0 - + 1; // +1 because the pdu id starts AFTER the separator - - let pdu_id = key[pduid_index..].to_vec(); - - Ok::<_, Error>(pdu_id) - }) - .filter_map(|r| r.ok()) - }); + let iterators = words + .clone() + .into_iter() + .map(move |word| { + let mut prefix2 = prefix.clone(); + prefix2.extend_from_slice(word.as_bytes()); + prefix2.push(0xff); + + let mut last_possible_id = prefix2.clone(); + last_possible_id.extend_from_slice(&u64::MAX.to_be_bytes()); + + Ok::<_, Error>( + self.tokenids + .iter_from(&last_possible_id, true) // Newest pdus first + .take_while(move |(k, _)| k.starts_with(&prefix2)) + .map(|(key, _)| { + let pduid_index = key + .iter() + .enumerate() + .filter(|(_, &b)| b == 0xff) + .nth(1) + .ok_or_else(|| Error::bad_database("Invalid tokenid in db."))? + .0 + + 1; // +1 because the pdu id starts AFTER the separator + + let pdu_id = key[pduid_index..].to_vec(); + + Ok::<_, Error>(pdu_id) + }) + .filter_map(|r| r.ok()), + ) + }) + .filter_map(|r| r.ok()); Ok(( utils::common_elements(iterators, |a, b| { @@ -2113,52 +2103,59 @@ impl Rooms { pub fn get_shared_rooms<'a>( &'a self, users: Vec, - ) -> impl Iterator> + 'a { - let iterators = users.into_iter().map(move |user_id| { - let mut prefix = user_id.as_bytes().to_vec(); - prefix.push(0xff); - - self.userroomid_joined - .scan_prefix(&prefix) - .keys() - .filter_map(|r| r.ok()) - .map(|key| { - let roomid_index = key - .iter() - .enumerate() - .find(|(_, &b)| b == 0xff) - .ok_or_else(|| Error::bad_database("Invalid userroomid_joined in db."))? - .0 - + 1; // +1 because the room id starts AFTER the separator - - let room_id = key[roomid_index..].to_vec(); - - Ok::<_, Error>(room_id) - }) - .filter_map(|r| r.ok()) - }); + ) -> Result> + 'a> { + let iterators = users + .into_iter() + .map(move |user_id| { + let mut prefix = user_id.as_bytes().to_vec(); + prefix.push(0xff); + + Ok::<_, Error>( + self.userroomid_joined + .scan_prefix(prefix) + .map(|(key, _)| { + let roomid_index = key + .iter() + .enumerate() + .find(|(_, &b)| b == 0xff) + .ok_or_else(|| { + Error::bad_database("Invalid userroomid_joined in db.") + })? + .0 + + 1; // +1 because the room id starts AFTER the separator + + let room_id = key[roomid_index..].to_vec(); + + Ok::<_, Error>(room_id) + }) + .filter_map(|r| r.ok()), + ) + }) + .filter_map(|r| r.ok()); // We use the default compare function because keys are sorted correctly (not reversed) - utils::common_elements(iterators, Ord::cmp) + Ok(utils::common_elements(iterators, Ord::cmp) .expect("users is not empty") .map(|bytes| { RoomId::try_from(utils::string_from_bytes(&*bytes).map_err(|_| { Error::bad_database("Invalid RoomId bytes in userroomid_joined") })?) .map_err(|_| Error::bad_database("Invalid RoomId in userroomid_joined.")) - }) + })) } /// Returns an iterator of all servers participating in this room. - pub fn room_servers(&self, room_id: &RoomId) -> impl Iterator>> { + pub fn room_servers<'a>( + &'a self, + room_id: &RoomId, + ) -> impl Iterator>> + 'a { let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); - self.roomserverids.scan_prefix(prefix).keys().map(|key| { + self.roomserverids.scan_prefix(prefix).map(|(key, _)| { Ok(Box::::try_from( utils::string_from_bytes( - &key? - .rsplit(|&b| b == 0xff) + &key.rsplit(|&b| b == 0xff) .next() .expect("rsplit always returns an element"), ) @@ -2171,15 +2168,17 @@ impl Rooms { } /// Returns an iterator of all rooms a server participates in (as far as we know). - pub fn server_rooms(&self, server: &ServerName) -> impl Iterator> { + pub fn server_rooms<'a>( + &'a self, + server: &ServerName, + ) -> impl Iterator> + 'a { let mut prefix = server.as_bytes().to_vec(); prefix.push(0xff); - self.serverroomids.scan_prefix(prefix).keys().map(|key| { + self.serverroomids.scan_prefix(prefix).map(|(key, _)| { Ok(RoomId::try_from( utils::string_from_bytes( - &key? - .rsplit(|&b| b == 0xff) + &key.rsplit(|&b| b == 0xff) .next() .expect("rsplit always returns an element"), ) @@ -2191,42 +2190,42 @@ impl Rooms { /// Returns an iterator over all joined members of a room. #[tracing::instrument(skip(self))] - pub fn room_members(&self, room_id: &RoomId) -> impl Iterator> { + pub fn room_members<'a>( + &'a self, + room_id: &RoomId, + ) -> impl Iterator> + 'a { let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); - self.roomuserid_joined - .scan_prefix(prefix) - .keys() - .map(|key| { - Ok(UserId::try_from( - utils::string_from_bytes( - &key? - .rsplit(|&b| b == 0xff) - .next() - .expect("rsplit always returns an element"), - ) - .map_err(|_| { - Error::bad_database("User ID in roomuserid_joined is invalid unicode.") - })?, + self.roomuserid_joined.scan_prefix(prefix).map(|(key, _)| { + Ok(UserId::try_from( + utils::string_from_bytes( + &key.rsplit(|&b| b == 0xff) + .next() + .expect("rsplit always returns an element"), ) - .map_err(|_| Error::bad_database("User ID in roomuserid_joined is invalid."))?) - }) + .map_err(|_| { + Error::bad_database("User ID in roomuserid_joined is invalid unicode.") + })?, + ) + .map_err(|_| Error::bad_database("User ID in roomuserid_joined is invalid."))?) + }) } /// Returns an iterator over all User IDs who ever joined a room. - pub fn room_useroncejoined(&self, room_id: &RoomId) -> impl Iterator> { + pub fn room_useroncejoined<'a>( + &'a self, + room_id: &RoomId, + ) -> impl Iterator> + 'a { let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); self.roomuseroncejoinedids .scan_prefix(prefix) - .keys() - .map(|key| { + .map(|(key, _)| { Ok(UserId::try_from( utils::string_from_bytes( - &key? - .rsplit(|&b| b == 0xff) + &key.rsplit(|&b| b == 0xff) .next() .expect("rsplit always returns an element"), ) @@ -2240,18 +2239,19 @@ impl Rooms { /// Returns an iterator over all invited members of a room. #[tracing::instrument(skip(self))] - pub fn room_members_invited(&self, room_id: &RoomId) -> impl Iterator> { + pub fn room_members_invited<'a>( + &'a self, + room_id: &RoomId, + ) -> impl Iterator> + 'a { let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); self.roomuserid_invitecount .scan_prefix(prefix) - .keys() - .map(|key| { + .map(|(key, _)| { Ok(UserId::try_from( utils::string_from_bytes( - &key? - .rsplit(|&b| b == 0xff) + &key.rsplit(|&b| b == 0xff) .next() .expect("rsplit always returns an element"), ) @@ -2270,7 +2270,7 @@ impl Rooms { key.extend_from_slice(user_id.as_bytes()); self.roomuserid_invitecount - .get(key)? + .get(&key)? .map_or(Ok(None), |bytes| { Ok(Some(utils::u64_from_bytes(&bytes).map_err(|_| { Error::bad_database("Invalid invitecount in db.") @@ -2285,7 +2285,7 @@ impl Rooms { key.extend_from_slice(user_id.as_bytes()); self.roomuserid_leftcount - .get(key)? + .get(&key)? .map_or(Ok(None), |bytes| { Ok(Some(utils::u64_from_bytes(&bytes).map_err(|_| { Error::bad_database("Invalid leftcount in db.") @@ -2295,15 +2295,16 @@ impl Rooms { /// Returns an iterator over all rooms this user joined. #[tracing::instrument(skip(self))] - pub fn rooms_joined(&self, user_id: &UserId) -> impl Iterator> { + pub fn rooms_joined<'a>( + &'a self, + user_id: &UserId, + ) -> impl Iterator> + 'a { self.userroomid_joined - .scan_prefix(user_id.as_bytes()) - .keys() - .map(|key| { + .scan_prefix(user_id.as_bytes().to_vec()) + .map(|(key, _)| { Ok(RoomId::try_from( utils::string_from_bytes( - &key? - .rsplit(|&b| b == 0xff) + &key.rsplit(|&b| b == 0xff) .next() .expect("rsplit always returns an element"), ) @@ -2317,32 +2318,33 @@ impl Rooms { /// Returns an iterator over all rooms a user was invited to. #[tracing::instrument(skip(self))] - pub fn rooms_invited( - &self, + pub fn rooms_invited<'a>( + &'a self, user_id: &UserId, - ) -> impl Iterator>)>> { + ) -> impl Iterator>)>> + 'a { let mut prefix = user_id.as_bytes().to_vec(); prefix.push(0xff); - self.userroomid_invitestate.scan_prefix(prefix).map(|r| { - let (key, state) = r?; - let room_id = RoomId::try_from( - utils::string_from_bytes( - &key.rsplit(|&b| b == 0xff) - .next() - .expect("rsplit always returns an element"), + self.userroomid_invitestate + .scan_prefix(prefix) + .map(|(key, state)| { + let room_id = RoomId::try_from( + utils::string_from_bytes( + &key.rsplit(|&b| b == 0xff) + .next() + .expect("rsplit always returns an element"), + ) + .map_err(|_| { + Error::bad_database("Room ID in userroomid_invited is invalid unicode.") + })?, ) - .map_err(|_| { - Error::bad_database("Room ID in userroomid_invited is invalid unicode.") - })?, - ) - .map_err(|_| Error::bad_database("Room ID in userroomid_invited is invalid."))?; + .map_err(|_| Error::bad_database("Room ID in userroomid_invited is invalid."))?; - let state = serde_json::from_slice(&state) - .map_err(|_| Error::bad_database("Invalid state in userroomid_invitestate."))?; + let state = serde_json::from_slice(&state) + .map_err(|_| Error::bad_database("Invalid state in userroomid_invitestate."))?; - Ok((room_id, state)) - }) + Ok((room_id, state)) + }) } #[tracing::instrument(skip(self))] @@ -2356,7 +2358,7 @@ impl Rooms { key.extend_from_slice(&room_id.as_bytes()); self.userroomid_invitestate - .get(key)? + .get(&key)? .map(|state| { let state = serde_json::from_slice(&state) .map_err(|_| Error::bad_database("Invalid state in userroomid_invitestate."))?; @@ -2377,7 +2379,7 @@ impl Rooms { key.extend_from_slice(&room_id.as_bytes()); self.userroomid_leftstate - .get(key)? + .get(&key)? .map(|state| { let state = serde_json::from_slice(&state) .map_err(|_| Error::bad_database("Invalid state in userroomid_leftstate."))?; @@ -2389,32 +2391,33 @@ impl Rooms { /// Returns an iterator over all rooms a user left. #[tracing::instrument(skip(self))] - pub fn rooms_left( - &self, + pub fn rooms_left<'a>( + &'a self, user_id: &UserId, - ) -> impl Iterator>)>> { + ) -> impl Iterator>)>> + 'a { let mut prefix = user_id.as_bytes().to_vec(); prefix.push(0xff); - self.userroomid_leftstate.scan_prefix(prefix).map(|r| { - let (key, state) = r?; - let room_id = RoomId::try_from( - utils::string_from_bytes( - &key.rsplit(|&b| b == 0xff) - .next() - .expect("rsplit always returns an element"), + self.userroomid_leftstate + .scan_prefix(prefix) + .map(|(key, state)| { + let room_id = RoomId::try_from( + utils::string_from_bytes( + &key.rsplit(|&b| b == 0xff) + .next() + .expect("rsplit always returns an element"), + ) + .map_err(|_| { + Error::bad_database("Room ID in userroomid_invited is invalid unicode.") + })?, ) - .map_err(|_| { - Error::bad_database("Room ID in userroomid_invited is invalid unicode.") - })?, - ) - .map_err(|_| Error::bad_database("Room ID in userroomid_invited is invalid."))?; + .map_err(|_| Error::bad_database("Room ID in userroomid_invited is invalid."))?; - let state = serde_json::from_slice(&state) - .map_err(|_| Error::bad_database("Invalid state in userroomid_leftstate."))?; + let state = serde_json::from_slice(&state) + .map_err(|_| Error::bad_database("Invalid state in userroomid_leftstate."))?; - Ok((room_id, state)) - }) + Ok((room_id, state)) + }) } pub fn once_joined(&self, user_id: &UserId, room_id: &RoomId) -> Result { @@ -2422,7 +2425,7 @@ impl Rooms { userroom_id.push(0xff); userroom_id.extend_from_slice(room_id.as_bytes()); - Ok(self.roomuseroncejoinedids.get(userroom_id)?.is_some()) + Ok(self.roomuseroncejoinedids.get(&userroom_id)?.is_some()) } pub fn is_joined(&self, user_id: &UserId, room_id: &RoomId) -> Result { @@ -2430,7 +2433,7 @@ impl Rooms { userroom_id.push(0xff); userroom_id.extend_from_slice(room_id.as_bytes()); - Ok(self.userroomid_joined.get(userroom_id)?.is_some()) + Ok(self.userroomid_joined.get(&userroom_id)?.is_some()) } pub fn is_invited(&self, user_id: &UserId, room_id: &RoomId) -> Result { @@ -2438,7 +2441,7 @@ impl Rooms { userroom_id.push(0xff); userroom_id.extend_from_slice(room_id.as_bytes()); - Ok(self.userroomid_invitestate.get(userroom_id)?.is_some()) + Ok(self.userroomid_invitestate.get(&userroom_id)?.is_some()) } pub fn is_left(&self, user_id: &UserId, room_id: &RoomId) -> Result { @@ -2446,6 +2449,6 @@ impl Rooms { userroom_id.push(0xff); userroom_id.extend_from_slice(room_id.as_bytes()); - Ok(self.userroomid_leftstate.get(userroom_id)?.is_some()) + Ok(self.userroomid_leftstate.get(&userroom_id)?.is_some()) } } diff --git a/src/database/rooms/edus.rs b/src/database/rooms/edus.rs index f4c7075e..677d26eb 100644 --- a/src/database/rooms/edus.rs +++ b/src/database/rooms/edus.rs @@ -1,4 +1,4 @@ -use crate::{utils, Error, Result}; +use crate::{database::abstraction::Tree, utils, Error, Result}; use ruma::{ events::{ presence::{PresenceEvent, PresenceEventContent}, @@ -13,17 +13,17 @@ use std::{ collections::{HashMap, HashSet}, convert::{TryFrom, TryInto}, mem, + sync::Arc, }; -#[derive(Clone)] pub struct RoomEdus { - pub(in super::super) readreceiptid_readreceipt: sled::Tree, // ReadReceiptId = RoomId + Count + UserId - pub(in super::super) roomuserid_privateread: sled::Tree, // RoomUserId = Room + User, PrivateRead = Count - pub(in super::super) roomuserid_lastprivatereadupdate: sled::Tree, // LastPrivateReadUpdate = Count - pub(in super::super) typingid_userid: sled::Tree, // TypingId = RoomId + TimeoutTime + Count - pub(in super::super) roomid_lasttypingupdate: sled::Tree, // LastRoomTypingUpdate = Count - pub(in super::super) presenceid_presence: sled::Tree, // PresenceId = RoomId + Count + UserId - pub(in super::super) userid_lastpresenceupdate: sled::Tree, // LastPresenceUpdate = Count + pub(in super::super) readreceiptid_readreceipt: Arc, // ReadReceiptId = RoomId + Count + UserId + pub(in super::super) roomuserid_privateread: Arc, // RoomUserId = Room + User, PrivateRead = Count + pub(in super::super) roomuserid_lastprivatereadupdate: Arc, // LastPrivateReadUpdate = Count + pub(in super::super) typingid_userid: Arc, // TypingId = RoomId + TimeoutTime + Count + pub(in super::super) roomid_lasttypingupdate: Arc, // LastRoomTypingUpdate = Count + pub(in super::super) presenceid_presence: Arc, // PresenceId = RoomId + Count + UserId + pub(in super::super) userid_lastpresenceupdate: Arc, // LastPresenceUpdate = Count } impl RoomEdus { @@ -38,15 +38,15 @@ impl RoomEdus { let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); + let mut last_possible_key = prefix.clone(); + last_possible_key.extend_from_slice(&u64::MAX.to_be_bytes()); + // Remove old entry - if let Some(old) = self + if let Some((old, _)) = self .readreceiptid_readreceipt - .scan_prefix(&prefix) - .keys() - .rev() - .filter_map(|r| r.ok()) - .take_while(|key| key.starts_with(&prefix)) - .find(|key| { + .iter_from(&last_possible_key, true) + .take_while(|(key, _)| key.starts_with(&prefix)) + .find(|(key, _)| { key.rsplit(|&b| b == 0xff) .next() .expect("rsplit always returns an element") @@ -54,7 +54,7 @@ impl RoomEdus { }) { // This is the old room_latest - self.readreceiptid_readreceipt.remove(old)?; + self.readreceiptid_readreceipt.remove(&old)?; } let mut room_latest_id = prefix; @@ -63,8 +63,8 @@ impl RoomEdus { room_latest_id.extend_from_slice(&user_id.as_bytes()); self.readreceiptid_readreceipt.insert( - room_latest_id, - &*serde_json::to_string(&event).expect("EduEvent::to_string always works"), + &room_latest_id, + &serde_json::to_vec(&event).expect("EduEvent::to_string always works"), )?; Ok(()) @@ -72,13 +72,12 @@ impl RoomEdus { /// Returns an iterator over the most recent read_receipts in a room that happened after the event with id `since`. #[tracing::instrument(skip(self))] - pub fn readreceipts_since( - &self, + pub fn readreceipts_since<'a>( + &'a self, room_id: &RoomId, since: u64, - ) -> Result< - impl Iterator)>>, - > { + ) -> impl Iterator)>> + 'a + { let mut prefix = room_id.as_bytes().to_vec(); prefix.push(0xff); let prefix2 = prefix.clone(); @@ -86,10 +85,8 @@ impl RoomEdus { let mut first_possible_edu = prefix.clone(); first_possible_edu.extend_from_slice(&(since + 1).to_be_bytes()); // +1 so we don't send the event at since - Ok(self - .readreceiptid_readreceipt - .range(&*first_possible_edu..) - .filter_map(|r| r.ok()) + self.readreceiptid_readreceipt + .iter_from(&first_possible_edu, false) .take_while(move |(k, _)| k.starts_with(&prefix2)) .map(move |(k, v)| { let count = @@ -115,7 +112,7 @@ impl RoomEdus { serde_json::value::to_raw_value(&json).expect("json is valid raw value"), ), )) - })) + }) } /// Sets a private read marker at `count`. @@ -146,11 +143,13 @@ impl RoomEdus { key.push(0xff); key.extend_from_slice(&user_id.as_bytes()); - self.roomuserid_privateread.get(key)?.map_or(Ok(None), |v| { - Ok(Some(utils::u64_from_bytes(&v).map_err(|_| { - Error::bad_database("Invalid private read marker bytes") - })?)) - }) + self.roomuserid_privateread + .get(&key)? + .map_or(Ok(None), |v| { + Ok(Some(utils::u64_from_bytes(&v).map_err(|_| { + Error::bad_database("Invalid private read marker bytes") + })?)) + }) } /// Returns the count of the last typing update in this room. @@ -215,11 +214,10 @@ impl RoomEdus { // Maybe there are multiple ones from calling roomtyping_add multiple times for outdated_edu in self .typingid_userid - .scan_prefix(&prefix) - .filter_map(|r| r.ok()) - .filter(|(_, v)| v == user_id.as_bytes()) + .scan_prefix(prefix) + .filter(|(_, v)| &**v == user_id.as_bytes()) { - self.typingid_userid.remove(outdated_edu.0)?; + self.typingid_userid.remove(&outdated_edu.0)?; found_outdated = true; } @@ -247,10 +245,8 @@ impl RoomEdus { // Find all outdated edus before inserting a new one for outdated_edu in self .typingid_userid - .scan_prefix(&prefix) - .keys() - .map(|key| { - let key = key?; + .scan_prefix(prefix) + .map(|(key, _)| { Ok::<_, Error>(( key.clone(), utils::u64_from_bytes( @@ -265,7 +261,7 @@ impl RoomEdus { .take_while(|&(_, timestamp)| timestamp < current_timestamp) { // This is an outdated edu (time > timestamp) - self.typingid_userid.remove(outdated_edu.0)?; + self.typingid_userid.remove(&outdated_edu.0)?; found_outdated = true; } @@ -309,10 +305,9 @@ impl RoomEdus { for user_id in self .typingid_userid .scan_prefix(prefix) - .values() - .map(|user_id| { + .map(|(_, user_id)| { Ok::<_, Error>( - UserId::try_from(utils::string_from_bytes(&user_id?).map_err(|_| { + UserId::try_from(utils::string_from_bytes(&user_id).map_err(|_| { Error::bad_database("User ID in typingid_userid is invalid unicode.") })?) .map_err(|_| Error::bad_database("User ID in typingid_userid is invalid."))?, @@ -351,12 +346,12 @@ impl RoomEdus { presence_id.extend_from_slice(&presence.sender.as_bytes()); self.presenceid_presence.insert( - presence_id, - &*serde_json::to_string(&presence).expect("PresenceEvent can be serialized"), + &presence_id, + &serde_json::to_vec(&presence).expect("PresenceEvent can be serialized"), )?; self.userid_lastpresenceupdate.insert( - &user_id.as_bytes(), + user_id.as_bytes(), &utils::millis_since_unix_epoch().to_be_bytes(), )?; @@ -403,7 +398,7 @@ impl RoomEdus { presence_id.extend_from_slice(&user_id.as_bytes()); self.presenceid_presence - .get(presence_id)? + .get(&presence_id)? .map(|value| { let mut presence = serde_json::from_slice::(&value) .map_err(|_| Error::bad_database("Invalid presence event in db."))?; @@ -438,7 +433,6 @@ impl RoomEdus { for (user_id_bytes, last_timestamp) in self .userid_lastpresenceupdate .iter() - .filter_map(|r| r.ok()) .filter_map(|(k, bytes)| { Some(( k, @@ -468,8 +462,8 @@ impl RoomEdus { presence_id.extend_from_slice(&user_id_bytes); self.presenceid_presence.insert( - presence_id, - &*serde_json::to_string(&PresenceEvent { + &presence_id, + &serde_json::to_vec(&PresenceEvent { content: PresenceEventContent { avatar_url: None, currently_active: None, @@ -515,8 +509,7 @@ impl RoomEdus { for (key, value) in self .presenceid_presence - .range(&*first_possible_edu..) - .filter_map(|r| r.ok()) + .iter_from(&*first_possible_edu, false) .take_while(|(key, _)| key.starts_with(&prefix)) { let user_id = UserId::try_from( diff --git a/src/database/sending.rs b/src/database/sending.rs index ed5b5ef8..77f6ed78 100644 --- a/src/database/sending.rs +++ b/src/database/sending.rs @@ -12,7 +12,10 @@ use crate::{ use federation::transactions::send_transaction_message; use log::{error, warn}; use ring::digest; -use rocket::futures::stream::{FuturesUnordered, StreamExt}; +use rocket::futures::{ + channel::mpsc, + stream::{FuturesUnordered, StreamExt}, +}; use ruma::{ api::{ appservice, @@ -27,9 +30,10 @@ use ruma::{ receipt::ReceiptType, MilliSecondsSinceUnixEpoch, ServerName, UInt, UserId, }; -use sled::IVec; use tokio::{select, sync::Semaphore}; +use super::abstraction::Tree; + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum OutgoingKind { Appservice(Box), @@ -70,13 +74,13 @@ pub enum SendingEventType { Edu(Vec), } -#[derive(Clone)] pub struct Sending { /// The state for a given state hash. - pub(super) servername_educount: sled::Tree, // EduCount: Count of last EDU sync - pub(super) servernamepduids: sled::Tree, // ServernamePduId = (+ / $)SenderKey / ServerName / UserId + PduId - pub(super) servercurrentevents: sled::Tree, // ServerCurrentEvents = (+ / $)ServerName / UserId + PduId / (*)EduEvent + pub(super) servername_educount: Arc, // EduCount: Count of last EDU sync + pub(super) servernamepduids: Arc, // ServernamePduId = (+ / $)SenderKey / ServerName / UserId + PduId + pub(super) servercurrentevents: Arc, // ServerCurrentEvents = (+ / $)ServerName / UserId + PduId / (*)EduEvent pub(super) maximum_requests: Arc, + pub sender: mpsc::UnboundedSender>, } enum TransactionStatus { @@ -86,28 +90,25 @@ enum TransactionStatus { } impl Sending { - pub fn start_handler(&self, db: &Database) { - let servernamepduids = self.servernamepduids.clone(); - let servercurrentevents = self.servercurrentevents.clone(); - + pub fn start_handler(&self, db: Arc, mut receiver: mpsc::UnboundedReceiver>) { let db = db.clone(); tokio::spawn(async move { let mut futures = FuturesUnordered::new(); - // Retry requests we could not finish yet - let mut subscriber = servernamepduids.watch_prefix(b""); let mut current_transaction_status = HashMap::, TransactionStatus>::new(); + // Retry requests we could not finish yet let mut initial_transactions = HashMap::>::new(); - for (key, outgoing_kind, event) in servercurrentevents - .iter() - .filter_map(|r| r.ok()) - .filter_map(|(key, _)| { - Self::parse_servercurrentevent(&key) - .ok() - .map(|(k, e)| (key, k, e)) - }) + for (key, outgoing_kind, event) in + db.sending + .servercurrentevents + .iter() + .filter_map(|(key, _)| { + Self::parse_servercurrentevent(&key) + .ok() + .map(|(k, e)| (key, k, e)) + }) { let entry = initial_transactions .entry(outgoing_kind.clone()) @@ -118,7 +119,7 @@ impl Sending { "Dropping some current events: {:?} {:?} {:?}", key, outgoing_kind, event ); - servercurrentevents.remove(key).unwrap(); + db.sending.servercurrentevents.remove(&key).unwrap(); continue; } @@ -137,20 +138,16 @@ impl Sending { match response { Ok(outgoing_kind) => { let prefix = outgoing_kind.get_prefix(); - for key in servercurrentevents - .scan_prefix(&prefix) - .keys() - .filter_map(|r| r.ok()) + for (key, _) in db.sending.servercurrentevents + .scan_prefix(prefix.clone()) { - servercurrentevents.remove(key).unwrap(); + db.sending.servercurrentevents.remove(&key).unwrap(); } // Find events that have been added since starting the last request - let new_events = servernamepduids - .scan_prefix(&prefix) - .keys() - .filter_map(|r| r.ok()) - .map(|k| { + let new_events = db.sending.servernamepduids + .scan_prefix(prefix.clone()) + .map(|(k, _)| { SendingEventType::Pdu(k[prefix.len()..].to_vec()) }) .take(30) @@ -166,8 +163,8 @@ impl Sending { SendingEventType::Pdu(b) | SendingEventType::Edu(b) => { current_key.extend_from_slice(&b); - servercurrentevents.insert(¤t_key, &[]).unwrap(); - servernamepduids.remove(¤t_key).unwrap(); + db.sending.servercurrentevents.insert(¤t_key, &[]).unwrap(); + db.sending.servernamepduids.remove(¤t_key).unwrap(); } } } @@ -195,18 +192,15 @@ impl Sending { } }; }, - Some(event) = &mut subscriber => { - // New sled version: - //for (_tree, key, value_opt) in &event { - // if value_opt.is_none() { - // continue; - // } - - if let sled::Event::Insert { key, .. } = event { - if let Ok((outgoing_kind, event)) = Self::parse_servercurrentevent(&key) { - if let Some(events) = Self::select_events(&outgoing_kind, vec![(event, key)], &mut current_transaction_status, &servercurrentevents, &servernamepduids, &db) { - futures.push(Self::handle_events(outgoing_kind, events, &db)); - } + Some(key) = receiver.next() => { + if let Ok((outgoing_kind, event)) = Self::parse_servercurrentevent(&key) { + if let Ok(Some(events)) = Self::select_events( + &outgoing_kind, + vec![(event, key)], + &mut current_transaction_status, + &db + ) { + futures.push(Self::handle_events(outgoing_kind, events, &db)); } } } @@ -217,12 +211,10 @@ impl Sending { fn select_events( outgoing_kind: &OutgoingKind, - new_events: Vec<(SendingEventType, IVec)>, // Events we want to send: event and full key + new_events: Vec<(SendingEventType, Vec)>, // Events we want to send: event and full key current_transaction_status: &mut HashMap, TransactionStatus>, - servercurrentevents: &sled::Tree, - servernamepduids: &sled::Tree, db: &Database, - ) -> Option> { + ) -> Result>> { let mut retry = false; let mut allow = true; @@ -252,29 +244,25 @@ impl Sending { .or_insert(TransactionStatus::Running); if !allow { - return None; + return Ok(None); } let mut events = Vec::new(); if retry { // We retry the previous transaction - for key in servercurrentevents - .scan_prefix(&prefix) - .keys() - .filter_map(|r| r.ok()) - { + for (key, _) in db.sending.servercurrentevents.scan_prefix(prefix) { if let Ok((_, e)) = Self::parse_servercurrentevent(&key) { events.push(e); } } } else { for (e, full_key) in new_events { - servercurrentevents.insert(&full_key, &[]).unwrap(); + db.sending.servercurrentevents.insert(&full_key, &[])?; // If it was a PDU we have to unqueue it // TODO: don't try to unqueue EDUs - servernamepduids.remove(&full_key).unwrap(); + db.sending.servernamepduids.remove(&full_key)?; events.push(e); } @@ -284,13 +272,12 @@ impl Sending { events.extend_from_slice(&select_edus); db.sending .servername_educount - .insert(server_name.as_bytes(), &last_count.to_be_bytes()) - .unwrap(); + .insert(server_name.as_bytes(), &last_count.to_be_bytes())?; } } } - Some(events) + Ok(Some(events)) } pub fn select_edus(db: &Database, server: &ServerName) -> Result<(Vec, u64)> { @@ -307,7 +294,7 @@ impl Sending { let mut max_edu_count = since; 'outer: for room_id in db.rooms.server_rooms(server) { let room_id = room_id?; - for r in db.rooms.edus.readreceipts_since(&room_id, since)? { + for r in db.rooms.edus.readreceipts_since(&room_id, since) { let (user_id, count, read_receipt) = r?; if count > max_edu_count { @@ -372,12 +359,13 @@ impl Sending { } #[tracing::instrument(skip(self))] - pub fn send_push_pdu(&self, pdu_id: &[u8], senderkey: IVec) -> Result<()> { + pub fn send_push_pdu(&self, pdu_id: &[u8], senderkey: Box<[u8]>) -> Result<()> { let mut key = b"$".to_vec(); key.extend_from_slice(&senderkey); key.push(0xff); key.extend_from_slice(pdu_id); - self.servernamepduids.insert(key, b"")?; + self.servernamepduids.insert(&key, b"")?; + self.sender.unbounded_send(key).unwrap(); Ok(()) } @@ -387,7 +375,8 @@ impl Sending { let mut key = server.as_bytes().to_vec(); key.push(0xff); key.extend_from_slice(pdu_id); - self.servernamepduids.insert(key, b"")?; + self.servernamepduids.insert(&key, b"")?; + self.sender.unbounded_send(key).unwrap(); Ok(()) } @@ -398,7 +387,8 @@ impl Sending { key.extend_from_slice(appservice_id.as_bytes()); key.push(0xff); key.extend_from_slice(pdu_id); - self.servernamepduids.insert(key, b"")?; + self.servernamepduids.insert(&key, b"")?; + self.sender.unbounded_send(key).unwrap(); Ok(()) } @@ -641,7 +631,7 @@ impl Sending { } } - fn parse_servercurrentevent(key: &IVec) -> Result<(OutgoingKind, SendingEventType)> { + fn parse_servercurrentevent(key: &[u8]) -> Result<(OutgoingKind, SendingEventType)> { // Appservices start with a plus Ok::<_, Error>(if key.starts_with(b"+") { let mut parts = key[1..].splitn(2, |&b| b == 0xff); diff --git a/src/database/transaction_ids.rs b/src/database/transaction_ids.rs index 1f8ba7de..3e377798 100644 --- a/src/database/transaction_ids.rs +++ b/src/database/transaction_ids.rs @@ -1,10 +1,12 @@ +use std::sync::Arc; + use crate::Result; use ruma::{DeviceId, UserId}; -use sled::IVec; -#[derive(Clone)] +use super::abstraction::Tree; + pub struct TransactionIds { - pub(super) userdevicetxnid_response: sled::Tree, // Response can be empty (/sendToDevice) or the event id (/send) + pub(super) userdevicetxnid_response: Arc, // Response can be empty (/sendToDevice) or the event id (/send) } impl TransactionIds { @@ -21,7 +23,7 @@ impl TransactionIds { key.push(0xff); key.extend_from_slice(txn_id.as_bytes()); - self.userdevicetxnid_response.insert(key, data)?; + self.userdevicetxnid_response.insert(&key, data)?; Ok(()) } @@ -31,7 +33,7 @@ impl TransactionIds { user_id: &UserId, device_id: Option<&DeviceId>, txn_id: &str, - ) -> Result> { + ) -> Result>> { let mut key = user_id.as_bytes().to_vec(); key.push(0xff); key.extend_from_slice(device_id.map(|d| d.as_bytes()).unwrap_or_default()); @@ -39,6 +41,6 @@ impl TransactionIds { key.extend_from_slice(txn_id.as_bytes()); // If there's no entry, this is a new transaction - Ok(self.userdevicetxnid_response.get(key)?) + Ok(self.userdevicetxnid_response.get(&key)?) } } diff --git a/src/database/uiaa.rs b/src/database/uiaa.rs index 3b778402..f7f3d1f8 100644 --- a/src/database/uiaa.rs +++ b/src/database/uiaa.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use crate::{client_server::SESSION_ID_LENGTH, utils, Error, Result}; use ruma::{ api::client::{ @@ -8,10 +10,11 @@ use ruma::{ DeviceId, UserId, }; -#[derive(Clone)] +use super::abstraction::Tree; + pub struct Uiaa { - pub(super) userdevicesessionid_uiaainfo: sled::Tree, // User-interactive authentication - pub(super) userdevicesessionid_uiaarequest: sled::Tree, // UiaaRequest = canonical json value + pub(super) userdevicesessionid_uiaainfo: Arc, // User-interactive authentication + pub(super) userdevicesessionid_uiaarequest: Arc, // UiaaRequest = canonical json value } impl Uiaa { @@ -185,7 +188,7 @@ impl Uiaa { self.userdevicesessionid_uiaarequest.insert( &userdevicesessionid, - &*serde_json::to_string(request).expect("json value to string always works"), + &serde_json::to_vec(request).expect("json value to vec always works"), )?; Ok(()) @@ -233,7 +236,7 @@ impl Uiaa { if let Some(uiaainfo) = uiaainfo { self.userdevicesessionid_uiaainfo.insert( &userdevicesessionid, - &*serde_json::to_string(&uiaainfo).expect("UiaaInfo::to_string always works"), + &serde_json::to_vec(&uiaainfo).expect("UiaaInfo::to_vec always works"), )?; } else { self.userdevicesessionid_uiaainfo diff --git a/src/database/users.rs b/src/database/users.rs index 52e6e33b..b6d3b3c4 100644 --- a/src/database/users.rs +++ b/src/database/users.rs @@ -7,40 +7,41 @@ use ruma::{ serde::Raw, DeviceId, DeviceKeyAlgorithm, DeviceKeyId, MilliSecondsSinceUnixEpoch, UInt, UserId, }; -use std::{collections::BTreeMap, convert::TryFrom, mem}; +use std::{collections::BTreeMap, convert::TryFrom, mem, sync::Arc}; + +use super::abstraction::Tree; -#[derive(Clone)] pub struct Users { - pub(super) userid_password: sled::Tree, - pub(super) userid_displayname: sled::Tree, - pub(super) userid_avatarurl: sled::Tree, - pub(super) userdeviceid_token: sled::Tree, - pub(super) userdeviceid_metadata: sled::Tree, // This is also used to check if a device exists - pub(super) userid_devicelistversion: sled::Tree, // DevicelistVersion = u64 - pub(super) token_userdeviceid: sled::Tree, - - pub(super) onetimekeyid_onetimekeys: sled::Tree, // OneTimeKeyId = UserId + DeviceKeyId - pub(super) userid_lastonetimekeyupdate: sled::Tree, // LastOneTimeKeyUpdate = Count - pub(super) keychangeid_userid: sled::Tree, // KeyChangeId = UserId/RoomId + Count - pub(super) keyid_key: sled::Tree, // KeyId = UserId + KeyId (depends on key type) - pub(super) userid_masterkeyid: sled::Tree, - pub(super) userid_selfsigningkeyid: sled::Tree, - pub(super) userid_usersigningkeyid: sled::Tree, - - pub(super) todeviceid_events: sled::Tree, // ToDeviceId = UserId + DeviceId + Count + pub(super) userid_password: Arc, + pub(super) userid_displayname: Arc, + pub(super) userid_avatarurl: Arc, + pub(super) userdeviceid_token: Arc, + pub(super) userdeviceid_metadata: Arc, // This is also used to check if a device exists + pub(super) userid_devicelistversion: Arc, // DevicelistVersion = u64 + pub(super) token_userdeviceid: Arc, + + pub(super) onetimekeyid_onetimekeys: Arc, // OneTimeKeyId = UserId + DeviceKeyId + pub(super) userid_lastonetimekeyupdate: Arc, // LastOneTimeKeyUpdate = Count + pub(super) keychangeid_userid: Arc, // KeyChangeId = UserId/RoomId + Count + pub(super) keyid_key: Arc, // KeyId = UserId + KeyId (depends on key type) + pub(super) userid_masterkeyid: Arc, + pub(super) userid_selfsigningkeyid: Arc, + pub(super) userid_usersigningkeyid: Arc, + + pub(super) todeviceid_events: Arc, // ToDeviceId = UserId + DeviceId + Count } impl Users { /// Check if a user has an account on this homeserver. pub fn exists(&self, user_id: &UserId) -> Result { - Ok(self.userid_password.contains_key(user_id.to_string())?) + Ok(self.userid_password.get(user_id.as_bytes())?.is_some()) } /// Check if account is deactivated pub fn is_deactivated(&self, user_id: &UserId) -> Result { Ok(self .userid_password - .get(user_id.to_string())? + .get(user_id.as_bytes())? .ok_or(Error::BadRequest( ErrorKind::InvalidParam, "User does not exist.", @@ -55,14 +56,14 @@ impl Users { } /// Returns the number of users registered on this server. - pub fn count(&self) -> usize { - self.userid_password.iter().count() + pub fn count(&self) -> Result { + Ok(self.userid_password.iter().count()) } /// Find out which user an access token belongs to. pub fn find_from_token(&self, token: &str) -> Result> { self.token_userdeviceid - .get(token)? + .get(token.as_bytes())? .map_or(Ok(None), |bytes| { let mut parts = bytes.split(|&b| b == 0xff); let user_bytes = parts.next().ok_or_else(|| { @@ -87,10 +88,10 @@ impl Users { } /// Returns an iterator over all users on this homeserver. - pub fn iter(&self) -> impl Iterator> { - self.userid_password.iter().keys().map(|bytes| { + pub fn iter<'a>(&'a self) -> impl Iterator> + 'a { + self.userid_password.iter().map(|(bytes, _)| { Ok( - UserId::try_from(utils::string_from_bytes(&bytes?).map_err(|_| { + UserId::try_from(utils::string_from_bytes(&bytes).map_err(|_| { Error::bad_database("User ID in userid_password is invalid unicode.") })?) .map_err(|_| Error::bad_database("User ID in userid_password is invalid."))?, @@ -101,7 +102,7 @@ impl Users { /// Returns the password hash for the given user. pub fn password_hash(&self, user_id: &UserId) -> Result> { self.userid_password - .get(user_id.to_string())? + .get(user_id.as_bytes())? .map_or(Ok(None), |bytes| { Ok(Some(utils::string_from_bytes(&bytes).map_err(|_| { Error::bad_database("Password hash in db is not valid string.") @@ -113,7 +114,8 @@ impl Users { pub fn set_password(&self, user_id: &UserId, password: Option<&str>) -> Result<()> { if let Some(password) = password { if let Ok(hash) = utils::calculate_hash(&password) { - self.userid_password.insert(user_id.to_string(), &*hash)?; + self.userid_password + .insert(user_id.as_bytes(), hash.as_bytes())?; Ok(()) } else { Err(Error::BadRequest( @@ -122,7 +124,7 @@ impl Users { )) } } else { - self.userid_password.insert(user_id.to_string(), "")?; + self.userid_password.insert(user_id.as_bytes(), b"")?; Ok(()) } } @@ -130,7 +132,7 @@ impl Users { /// Returns the displayname of a user on this homeserver. pub fn displayname(&self, user_id: &UserId) -> Result> { self.userid_displayname - .get(user_id.to_string())? + .get(user_id.as_bytes())? .map_or(Ok(None), |bytes| { Ok(Some(utils::string_from_bytes(&bytes).map_err(|_| { Error::bad_database("Displayname in db is invalid.") @@ -142,9 +144,9 @@ impl Users { pub fn set_displayname(&self, user_id: &UserId, displayname: Option) -> Result<()> { if let Some(displayname) = displayname { self.userid_displayname - .insert(user_id.to_string(), &*displayname)?; + .insert(user_id.as_bytes(), displayname.as_bytes())?; } else { - self.userid_displayname.remove(user_id.to_string())?; + self.userid_displayname.remove(user_id.as_bytes())?; } Ok(()) @@ -153,7 +155,7 @@ impl Users { /// Get a the avatar_url of a user. pub fn avatar_url(&self, user_id: &UserId) -> Result> { self.userid_avatarurl - .get(user_id.to_string())? + .get(user_id.as_bytes())? .map(|bytes| { let s = utils::string_from_bytes(&bytes) .map_err(|_| Error::bad_database("Avatar URL in db is invalid."))?; @@ -166,9 +168,9 @@ impl Users { pub fn set_avatar_url(&self, user_id: &UserId, avatar_url: Option) -> Result<()> { if let Some(avatar_url) = avatar_url { self.userid_avatarurl - .insert(user_id.to_string(), avatar_url.to_string().as_str())?; + .insert(user_id.as_bytes(), avatar_url.to_string().as_bytes())?; } else { - self.userid_avatarurl.remove(user_id.to_string())?; + self.userid_avatarurl.remove(user_id.as_bytes())?; } Ok(()) @@ -190,19 +192,17 @@ impl Users { userdeviceid.extend_from_slice(device_id.as_bytes()); self.userid_devicelistversion - .update_and_fetch(&user_id.as_bytes(), utils::increment)? - .expect("utils::increment will always put in a value"); + .increment(user_id.as_bytes())?; self.userdeviceid_metadata.insert( - userdeviceid, - serde_json::to_string(&Device { + &userdeviceid, + &serde_json::to_vec(&Device { device_id: device_id.into(), display_name: initial_device_display_name, last_seen_ip: None, // TODO last_seen_ts: Some(MilliSecondsSinceUnixEpoch::now()), }) - .expect("Device::to_string never fails.") - .as_bytes(), + .expect("Device::to_string never fails."), )?; self.set_token(user_id, &device_id, token)?; @@ -217,7 +217,8 @@ impl Users { userdeviceid.extend_from_slice(device_id.as_bytes()); // Remove tokens - if let Some(old_token) = self.userdeviceid_token.remove(&userdeviceid)? { + if let Some(old_token) = self.userdeviceid_token.get(&userdeviceid)? { + self.userdeviceid_token.remove(&userdeviceid)?; self.token_userdeviceid.remove(&old_token)?; } @@ -225,15 +226,14 @@ impl Users { let mut prefix = userdeviceid.clone(); prefix.push(0xff); - for key in self.todeviceid_events.scan_prefix(&prefix).keys() { - self.todeviceid_events.remove(key?)?; + for (key, _) in self.todeviceid_events.scan_prefix(prefix) { + self.todeviceid_events.remove(&key)?; } // TODO: Remove onetimekeys self.userid_devicelistversion - .update_and_fetch(&user_id.as_bytes(), utils::increment)? - .expect("utils::increment will always put in a value"); + .increment(user_id.as_bytes())?; self.userdeviceid_metadata.remove(&userdeviceid)?; @@ -241,16 +241,18 @@ impl Users { } /// Returns an iterator over all device ids of this user. - pub fn all_device_ids(&self, user_id: &UserId) -> impl Iterator>> { + pub fn all_device_ids<'a>( + &'a self, + user_id: &UserId, + ) -> impl Iterator>> + 'a { let mut prefix = user_id.as_bytes().to_vec(); prefix.push(0xff); // All devices have metadata self.userdeviceid_metadata .scan_prefix(prefix) - .keys() - .map(|bytes| { + .map(|(bytes, _)| { Ok(utils::string_from_bytes( - &*bytes? + &bytes .rsplit(|&b| b == 0xff) .next() .ok_or_else(|| Error::bad_database("UserDevice ID in db is invalid."))?, @@ -271,13 +273,15 @@ impl Users { // Remove old token if let Some(old_token) = self.userdeviceid_token.get(&userdeviceid)? { - self.token_userdeviceid.remove(old_token)?; + self.token_userdeviceid.remove(&old_token)?; // It will be removed from userdeviceid_token by the insert later } // Assign token to user device combination - self.userdeviceid_token.insert(&userdeviceid, &*token)?; - self.token_userdeviceid.insert(token, userdeviceid)?; + self.userdeviceid_token + .insert(&userdeviceid, token.as_bytes())?; + self.token_userdeviceid + .insert(token.as_bytes(), &userdeviceid)?; Ok(()) } @@ -309,8 +313,7 @@ impl Users { self.onetimekeyid_onetimekeys.insert( &key, - &*serde_json::to_string(&one_time_key_value) - .expect("OneTimeKey::to_string always works"), + &serde_json::to_vec(&one_time_key_value).expect("OneTimeKey::to_vec always works"), )?; self.userid_lastonetimekeyupdate @@ -350,10 +353,9 @@ impl Users { .insert(&user_id.as_bytes(), &globals.next_count()?.to_be_bytes())?; self.onetimekeyid_onetimekeys - .scan_prefix(&prefix) + .scan_prefix(prefix) .next() - .map(|r| { - let (key, value) = r?; + .map(|(key, value)| { self.onetimekeyid_onetimekeys.remove(&key)?; Ok(( @@ -383,21 +385,20 @@ impl Users { let mut counts = BTreeMap::new(); - for algorithm in self - .onetimekeyid_onetimekeys - .scan_prefix(&userdeviceid) - .keys() - .map(|bytes| { - Ok::<_, Error>( - serde_json::from_slice::( - &*bytes?.rsplit(|&b| b == 0xff).next().ok_or_else(|| { - Error::bad_database("OneTimeKey ID in db is invalid.") - })?, + for algorithm in + self.onetimekeyid_onetimekeys + .scan_prefix(userdeviceid) + .map(|(bytes, _)| { + Ok::<_, Error>( + serde_json::from_slice::( + &*bytes.rsplit(|&b| b == 0xff).next().ok_or_else(|| { + Error::bad_database("OneTimeKey ID in db is invalid.") + })?, + ) + .map_err(|_| Error::bad_database("DeviceKeyId in db is invalid."))? + .algorithm(), ) - .map_err(|_| Error::bad_database("DeviceKeyId in db is invalid."))? - .algorithm(), - ) - }) + }) { *counts.entry(algorithm?).or_default() += UInt::from(1_u32); } @@ -419,7 +420,7 @@ impl Users { self.keyid_key.insert( &userdeviceid, - &*serde_json::to_string(&device_keys).expect("DeviceKeys::to_string always works"), + &serde_json::to_vec(&device_keys).expect("DeviceKeys::to_vec always works"), )?; self.mark_device_key_update(user_id, rooms, globals)?; @@ -460,11 +461,11 @@ impl Users { self.keyid_key.insert( &master_key_key, - &*serde_json::to_string(&master_key).expect("CrossSigningKey::to_string always works"), + &serde_json::to_vec(&master_key).expect("CrossSigningKey::to_vec always works"), )?; self.userid_masterkeyid - .insert(&*user_id.to_string(), master_key_key)?; + .insert(user_id.as_bytes(), &master_key_key)?; // Self-signing key if let Some(self_signing_key) = self_signing_key { @@ -486,12 +487,12 @@ impl Users { self.keyid_key.insert( &self_signing_key_key, - &*serde_json::to_string(&self_signing_key) - .expect("CrossSigningKey::to_string always works"), + &serde_json::to_vec(&self_signing_key) + .expect("CrossSigningKey::to_vec always works"), )?; self.userid_selfsigningkeyid - .insert(&*user_id.to_string(), self_signing_key_key)?; + .insert(user_id.as_bytes(), &self_signing_key_key)?; } // User-signing key @@ -514,12 +515,12 @@ impl Users { self.keyid_key.insert( &user_signing_key_key, - &*serde_json::to_string(&user_signing_key) - .expect("CrossSigningKey::to_string always works"), + &serde_json::to_vec(&user_signing_key) + .expect("CrossSigningKey::to_vec always works"), )?; self.userid_usersigningkeyid - .insert(&*user_id.to_string(), user_signing_key_key)?; + .insert(user_id.as_bytes(), &user_signing_key_key)?; } self.mark_device_key_update(user_id, rooms, globals)?; @@ -561,8 +562,7 @@ impl Users { self.keyid_key.insert( &key, - &*serde_json::to_string(&cross_signing_key) - .expect("CrossSigningKey::to_string always works"), + &serde_json::to_vec(&cross_signing_key).expect("CrossSigningKey::to_vec always works"), )?; // TODO: Should we notify about this change? @@ -572,24 +572,20 @@ impl Users { } #[tracing::instrument(skip(self))] - pub fn keys_changed( - &self, + pub fn keys_changed<'a>( + &'a self, user_or_room_id: &str, from: u64, to: Option, - ) -> impl Iterator> { + ) -> impl Iterator> + 'a { let mut prefix = user_or_room_id.as_bytes().to_vec(); prefix.push(0xff); let mut start = prefix.clone(); start.extend_from_slice(&(from + 1).to_be_bytes()); - let mut end = prefix.clone(); - end.extend_from_slice(&to.unwrap_or(u64::MAX).to_be_bytes()); - self.keychangeid_userid - .range(start..end) - .filter_map(|r| r.ok()) + .iter_from(&start, false) .take_while(move |(k, _)| k.starts_with(&prefix)) .map(|(_, bytes)| { Ok( @@ -625,13 +621,13 @@ impl Users { key.push(0xff); key.extend_from_slice(&count); - self.keychangeid_userid.insert(key, &*user_id.to_string())?; + self.keychangeid_userid.insert(&key, user_id.as_bytes())?; } let mut key = user_id.as_bytes().to_vec(); key.push(0xff); key.extend_from_slice(&count); - self.keychangeid_userid.insert(key, &*user_id.to_string())?; + self.keychangeid_userid.insert(&key, user_id.as_bytes())?; Ok(()) } @@ -645,7 +641,7 @@ impl Users { key.push(0xff); key.extend_from_slice(device_id.as_bytes()); - self.keyid_key.get(key)?.map_or(Ok(None), |bytes| { + self.keyid_key.get(&key)?.map_or(Ok(None), |bytes| { Ok(Some(serde_json::from_slice(&bytes).map_err(|_| { Error::bad_database("DeviceKeys in db are invalid.") })?)) @@ -658,9 +654,9 @@ impl Users { allowed_signatures: F, ) -> Result> { self.userid_masterkeyid - .get(user_id.to_string())? + .get(user_id.as_bytes())? .map_or(Ok(None), |key| { - self.keyid_key.get(key)?.map_or(Ok(None), |bytes| { + self.keyid_key.get(&key)?.map_or(Ok(None), |bytes| { let mut cross_signing_key = serde_json::from_slice::(&bytes) .map_err(|_| { Error::bad_database("CrossSigningKey in db is invalid.") @@ -685,9 +681,9 @@ impl Users { allowed_signatures: F, ) -> Result> { self.userid_selfsigningkeyid - .get(user_id.to_string())? + .get(user_id.as_bytes())? .map_or(Ok(None), |key| { - self.keyid_key.get(key)?.map_or(Ok(None), |bytes| { + self.keyid_key.get(&key)?.map_or(Ok(None), |bytes| { let mut cross_signing_key = serde_json::from_slice::(&bytes) .map_err(|_| { Error::bad_database("CrossSigningKey in db is invalid.") @@ -708,9 +704,9 @@ impl Users { pub fn get_user_signing_key(&self, user_id: &UserId) -> Result> { self.userid_usersigningkeyid - .get(user_id.to_string())? + .get(user_id.as_bytes())? .map_or(Ok(None), |key| { - self.keyid_key.get(key)?.map_or(Ok(None), |bytes| { + self.keyid_key.get(&key)?.map_or(Ok(None), |bytes| { Ok(Some(serde_json::from_slice(&bytes).map_err(|_| { Error::bad_database("CrossSigningKey in db is invalid.") })?)) @@ -740,7 +736,7 @@ impl Users { self.todeviceid_events.insert( &key, - &*serde_json::to_string(&json).expect("Map::to_string always works"), + &serde_json::to_vec(&json).expect("Map::to_vec always works"), )?; Ok(()) @@ -759,9 +755,9 @@ impl Users { prefix.extend_from_slice(device_id.as_bytes()); prefix.push(0xff); - for value in self.todeviceid_events.scan_prefix(&prefix).values() { + for (_, value) in self.todeviceid_events.scan_prefix(prefix) { events.push( - serde_json::from_slice(&*value?) + serde_json::from_slice(&value) .map_err(|_| Error::bad_database("Event in todeviceid_events is invalid."))?, ); } @@ -786,10 +782,9 @@ impl Users { for (key, _) in self .todeviceid_events - .range(&*prefix..=&*last) - .keys() - .map(|key| { - let key = key?; + .iter_from(&last, true) + .take_while(move |(k, _)| k.starts_with(&prefix)) + .map(|(key, _)| { Ok::<_, Error>(( key.clone(), utils::u64_from_bytes(&key[key.len() - mem::size_of::()..key.len()]) @@ -799,7 +794,7 @@ impl Users { .filter_map(|r| r.ok()) .take_while(|&(_, count)| count <= until) { - self.todeviceid_events.remove(key)?; + self.todeviceid_events.remove(&key)?; } Ok(()) @@ -819,14 +814,11 @@ impl Users { assert!(self.userdeviceid_metadata.get(&userdeviceid)?.is_some()); self.userid_devicelistversion - .update_and_fetch(&user_id.as_bytes(), utils::increment)? - .expect("utils::increment will always put in a value"); + .increment(user_id.as_bytes())?; self.userdeviceid_metadata.insert( - userdeviceid, - serde_json::to_string(device) - .expect("Device::to_string always works") - .as_bytes(), + &userdeviceid, + &serde_json::to_vec(device).expect("Device::to_string always works"), )?; Ok(()) @@ -861,15 +853,17 @@ impl Users { }) } - pub fn all_devices_metadata(&self, user_id: &UserId) -> impl Iterator> { + pub fn all_devices_metadata<'a>( + &'a self, + user_id: &UserId, + ) -> impl Iterator> + 'a { let mut key = user_id.as_bytes().to_vec(); key.push(0xff); self.userdeviceid_metadata .scan_prefix(key) - .values() - .map(|bytes| { - Ok(serde_json::from_slice::(&bytes?).map_err(|_| { + .map(|(_, bytes)| { + Ok(serde_json::from_slice::(&bytes).map_err(|_| { Error::bad_database("Device in userdeviceid_metadata is invalid.") })?) }) @@ -885,7 +879,7 @@ impl Users { // Set the password to "" to indicate a deactivated account. Hashes will never result in an // empty string, so the user will not be able to log in again. Systems like changing the // password without logging in should check if the account is deactivated. - self.userid_password.insert(user_id.to_string(), "")?; + self.userid_password.insert(user_id.as_bytes(), &[])?; // TODO: Unhook 3PID Ok(()) diff --git a/src/error.rs b/src/error.rs index e2664e21..93c67c12 100644 --- a/src/error.rs +++ b/src/error.rs @@ -23,11 +23,16 @@ pub type Result = std::result::Result; #[derive(Error, Debug)] pub enum Error { - #[error("There was a problem with the connection to the database.")] + #[error("There was a problem with the connection to the sled database.")] SledError { #[from] source: sled::Error, }, + #[error("There was a problem with the connection to the rocksdb database: {source}")] + RocksDbError { + #[from] + source: rocksdb::Error, + }, #[error("Could not generate an image.")] ImageError { #[from] diff --git a/src/main.rs b/src/main.rs index e76cea4e..8b63d1d6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,6 +12,8 @@ mod pdu; mod ruma_wrapper; mod utils; +use std::sync::Arc; + use database::Config; pub use database::Database; pub use error::{Error, Result}; @@ -31,7 +33,7 @@ use rocket::{ use tracing::span; use tracing_subscriber::{prelude::*, Registry}; -fn setup_rocket(config: Figment, data: Database) -> rocket::Rocket { +fn setup_rocket(config: Figment, data: Arc) -> rocket::Rocket { rocket::custom(config) .manage(data) .mount( @@ -197,8 +199,6 @@ async fn main() { .await .expect("config is valid"); - db.sending.start_handler(&db); - if config.allow_jaeger { let (tracer, _uninstall) = opentelemetry_jaeger::new_pipeline() .with_service_name("conduit") diff --git a/src/ruma_wrapper.rs b/src/ruma_wrapper.rs index 147df3c3..ba2c37ed 100644 --- a/src/ruma_wrapper.rs +++ b/src/ruma_wrapper.rs @@ -1,10 +1,10 @@ -use crate::Error; +use crate::{Database, Error}; use ruma::{ api::OutgoingResponse, identifiers::{DeviceId, UserId}, Outgoing, }; -use std::ops::Deref; +use std::{ops::Deref, sync::Arc}; #[cfg(feature = "conduit_bin")] use { @@ -51,7 +51,7 @@ where async fn from_data(request: &'a Request<'_>, data: Data) -> data::Outcome { let metadata = T::Incoming::METADATA; let db = request - .guard::>() + .guard::>>() .await .expect("database was loaded"); @@ -75,6 +75,7 @@ where )) = db .appservice .iter_all() + .unwrap() .filter_map(|r| r.ok()) .find(|(_id, registration)| { registration diff --git a/src/server_server.rs b/src/server_server.rs index b405c1ab..7a338dc5 100644 --- a/src/server_server.rs +++ b/src/server_server.rs @@ -433,7 +433,7 @@ pub async fn request_well_known( #[cfg_attr(feature = "conduit_bin", get("/_matrix/federation/v1/version"))] #[tracing::instrument(skip(db))] pub fn get_server_version_route( - db: State<'_, Database>, + db: State<'_, Arc>, ) -> ConduitResult { if !db.globals.allow_federation() { return Err(Error::bad_config("Federation is disabled.")); @@ -451,7 +451,7 @@ pub fn get_server_version_route( // Response type for this endpoint is Json because we need to calculate a signature for the response #[cfg_attr(feature = "conduit_bin", get("/_matrix/key/v2/server"))] #[tracing::instrument(skip(db))] -pub fn get_server_keys_route(db: State<'_, Database>) -> Json { +pub fn get_server_keys_route(db: State<'_, Arc>) -> Json { if !db.globals.allow_federation() { // TODO: Use proper types return Json("Federation is disabled.".to_owned()); @@ -498,7 +498,7 @@ pub fn get_server_keys_route(db: State<'_, Database>) -> Json { #[cfg_attr(feature = "conduit_bin", get("/_matrix/key/v2/server/<_>"))] #[tracing::instrument(skip(db))] -pub fn get_server_keys_deprecated_route(db: State<'_, Database>) -> Json { +pub fn get_server_keys_deprecated_route(db: State<'_, Arc>) -> Json { get_server_keys_route(db) } @@ -508,7 +508,7 @@ pub fn get_server_keys_deprecated_route(db: State<'_, Database>) -> Json )] #[tracing::instrument(skip(db, body))] pub async fn get_public_rooms_filtered_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -556,7 +556,7 @@ pub async fn get_public_rooms_filtered_route( )] #[tracing::instrument(skip(db, body))] pub async fn get_public_rooms_route( - db: State<'_, Database>, + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -603,8 +603,8 @@ pub async fn get_public_rooms_route( put("/_matrix/federation/v1/send/<_>", data = "") )] #[tracing::instrument(skip(db, body))] -pub async fn send_transaction_message_route<'a>( - db: State<'a, Database>, +pub async fn send_transaction_message_route( + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -1585,7 +1585,7 @@ pub(crate) async fn fetch_signing_keys( .await { db.globals - .add_signing_key(origin, &get_keys_response.server_key)?; + .add_signing_key(origin, get_keys_response.server_key.clone())?; result.extend( get_keys_response @@ -1628,7 +1628,7 @@ pub(crate) async fn fetch_signing_keys( { trace!("Got signing keys: {:?}", keys); for k in keys.server_keys { - db.globals.add_signing_key(origin, &k)?; + db.globals.add_signing_key(origin, k.clone())?; result.extend( k.verify_keys .into_iter() @@ -1686,7 +1686,7 @@ pub(crate) fn append_incoming_pdu( &db, )?; - for appservice in db.appservice.iter_all().filter_map(|r| r.ok()) { + for appservice in db.appservice.iter_all()?.filter_map(|r| r.ok()) { if let Some(namespaces) = appservice.1.get("namespaces") { let users = namespaces .get("users") @@ -1758,8 +1758,8 @@ pub(crate) fn append_incoming_pdu( get("/_matrix/federation/v1/event/<_>", data = "") )] #[tracing::instrument(skip(db, body))] -pub fn get_event_route<'a>( - db: State<'a, Database>, +pub fn get_event_route( + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -1783,8 +1783,8 @@ pub fn get_event_route<'a>( post("/_matrix/federation/v1/get_missing_events/<_>", data = "") )] #[tracing::instrument(skip(db, body))] -pub fn get_missing_events_route<'a>( - db: State<'a, Database>, +pub fn get_missing_events_route( + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -1832,8 +1832,8 @@ pub fn get_missing_events_route<'a>( get("/_matrix/federation/v1/state_ids/<_>", data = "") )] #[tracing::instrument(skip(db, body))] -pub fn get_room_state_ids_route<'a>( - db: State<'a, Database>, +pub fn get_room_state_ids_route( + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -1884,8 +1884,8 @@ pub fn get_room_state_ids_route<'a>( get("/_matrix/federation/v1/make_join/<_>/<_>", data = "") )] #[tracing::instrument(skip(db, body))] -pub fn create_join_event_template_route<'a>( - db: State<'a, Database>, +pub fn create_join_event_template_route( + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -2055,8 +2055,8 @@ pub fn create_join_event_template_route<'a>( put("/_matrix/federation/v2/send_join/<_>/<_>", data = "") )] #[tracing::instrument(skip(db, body))] -pub async fn create_join_event_route<'a>( - db: State<'a, Database>, +pub async fn create_join_event_route( + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -2171,8 +2171,8 @@ pub async fn create_join_event_route<'a>( put("/_matrix/federation/v2/invite/<_>/<_>", data = "") )] #[tracing::instrument(skip(db, body))] -pub async fn create_invite_route<'a>( - db: State<'a, Database>, +pub async fn create_invite_route( + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -2276,8 +2276,8 @@ pub async fn create_invite_route<'a>( get("/_matrix/federation/v1/user/devices/<_>", data = "") )] #[tracing::instrument(skip(db, body))] -pub fn get_devices_route<'a>( - db: State<'a, Database>, +pub fn get_devices_route( + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -2316,8 +2316,8 @@ pub fn get_devices_route<'a>( get("/_matrix/federation/v1/query/directory", data = "") )] #[tracing::instrument(skip(db, body))] -pub fn get_room_information_route<'a>( - db: State<'a, Database>, +pub fn get_room_information_route( + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -2344,8 +2344,8 @@ pub fn get_room_information_route<'a>( get("/_matrix/federation/v1/query/profile", data = "") )] #[tracing::instrument(skip(db, body))] -pub fn get_profile_information_route<'a>( - db: State<'a, Database>, +pub fn get_profile_information_route( + db: State<'_, Arc>, body: Ruma>, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -2378,8 +2378,8 @@ pub fn get_profile_information_route<'a>( post("/_matrix/federation/v1/user/keys/query", data = "") )] #[tracing::instrument(skip(db, body))] -pub fn get_keys_route<'a>( - db: State<'a, Database>, +pub fn get_keys_route( + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { if !db.globals.allow_federation() { @@ -2406,8 +2406,8 @@ pub fn get_keys_route<'a>( post("/_matrix/federation/v1/user/keys/claim", data = "") )] #[tracing::instrument(skip(db, body))] -pub async fn claim_keys_route<'a>( - db: State<'a, Database>, +pub async fn claim_keys_route( + db: State<'_, Arc>, body: Ruma, ) -> ConduitResult { if !db.globals.allow_federation() { diff --git a/src/utils.rs b/src/utils.rs index 106baffd..f59afb3a 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,6 +1,7 @@ use argon2::{Config, Variant}; use cmp::Ordering; use rand::prelude::*; +use rocksdb::MergeOperands; use ruma::serde::{try_from_json_map, CanonicalJsonError, CanonicalJsonObject}; use std::{ cmp, @@ -15,6 +16,14 @@ pub fn millis_since_unix_epoch() -> u64 { .as_millis() as u64 } +pub fn increment_rocksdb( + _new_key: &[u8], + old: Option<&[u8]>, + _operands: &mut MergeOperands, +) -> Option> { + increment(old) +} + pub fn increment(old: Option<&[u8]>) -> Option> { let number = match old.map(|bytes| bytes.try_into()) { Some(Ok(bytes)) => { @@ -27,16 +36,14 @@ pub fn increment(old: Option<&[u8]>) -> Option> { Some(number.to_be_bytes().to_vec()) } -pub fn generate_keypair(old: Option<&[u8]>) -> Option> { - Some(old.map(|s| s.to_vec()).unwrap_or_else(|| { - let mut value = random_string(8).as_bytes().to_vec(); - value.push(0xff); - value.extend_from_slice( - &ruma::signatures::Ed25519KeyPair::generate() - .expect("Ed25519KeyPair generation always works (?)"), - ); - value - })) +pub fn generate_keypair() -> Vec { + let mut value = random_string(8).as_bytes().to_vec(); + value.push(0xff); + value.extend_from_slice( + &ruma::signatures::Ed25519KeyPair::generate() + .expect("Ed25519KeyPair generation always works (?)"), + ); + value } /// Parses the bytes into an u64. From 972caacdc2de95183ddddd4084282a069a75b89a Mon Sep 17 00:00:00 2001 From: hamidreza kalbasi Date: Fri, 4 Jun 2021 08:06:12 +0430 Subject: [PATCH 2/8] put media in filesystem --- src/client_server/media.rs | 14 +++++---- src/database.rs | 4 +-- src/database/abstraction.rs | 6 ++-- src/database/globals.rs | 21 ++++++++++---- src/database/media.rs | 58 +++++++++++++++++++++++++++---------- src/error.rs | 5 ++++ 6 files changed, 78 insertions(+), 30 deletions(-) diff --git a/src/client_server/media.rs b/src/client_server/media.rs index 0673787c..14ab6db0 100644 --- a/src/client_server/media.rs +++ b/src/client_server/media.rs @@ -38,6 +38,7 @@ pub async fn create_content_route( ); db.media.create( mxc.clone(), + &db.globals, &body .filename .as_ref() @@ -45,7 +46,7 @@ pub async fn create_content_route( .as_deref(), &body.content_type.as_deref(), &body.file, - )?; + ).await?; db.flush().await?; @@ -71,7 +72,7 @@ pub async fn get_content_route( content_disposition, content_type, file, - }) = db.media.get(&mxc)? + }) = db.media.get(&db.globals, &mxc).await? { Ok(get_content::Response { file, @@ -95,10 +96,11 @@ pub async fn get_content_route( db.media.create( mxc, + &db.globals, &get_content_response.content_disposition.as_deref(), &get_content_response.content_type.as_deref(), &get_content_response.file, - )?; + ).await?; Ok(get_content_response.into()) } else { @@ -121,13 +123,14 @@ pub async fn get_content_thumbnail_route( content_type, file, .. }) = db.media.get_thumbnail( mxc.clone(), + &db.globals, body.width .try_into() .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Width is invalid."))?, body.height .try_into() .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Width is invalid."))?, - )? { + ).await? { Ok(get_content_thumbnail::Response { file, content_type }.into()) } else if &*body.server_name != db.globals.server_name() && body.allow_remote { let get_thumbnail_response = db @@ -148,12 +151,13 @@ pub async fn get_content_thumbnail_route( db.media.upload_thumbnail( mxc, + &db.globals, &None, &get_thumbnail_response.content_type, body.width.try_into().expect("all UInts are valid u32s"), body.height.try_into().expect("all UInts are valid u32s"), &get_thumbnail_response.file, - )?; + ).await?; Ok(get_thumbnail_response.into()) } else { diff --git a/src/database.rs b/src/database.rs index e3b954ef..04d37358 100644 --- a/src/database.rs +++ b/src/database.rs @@ -245,7 +245,7 @@ impl Database { db.globals.bump_database_version(1)?; - info!("Migration: 0 -> 1 finished"); + println!("Migration: 0 -> 1 finished"); } if db.globals.database_version()? < 2 { @@ -262,7 +262,7 @@ impl Database { db.globals.bump_database_version(2)?; - info!("Migration: 1 -> 2 finished"); + println!("Migration: 1 -> 2 finished"); } // This data is probably outdated diff --git a/src/database/abstraction.rs b/src/database/abstraction.rs index 5a2afd5a..ad032fb3 100644 --- a/src/database/abstraction.rs +++ b/src/database/abstraction.rs @@ -47,7 +47,7 @@ pub trait Tree: Send + Sync { fn scan_prefix<'a>( &'a self, prefix: Vec, - ) -> Box, Box<[u8]>)> + 'a>; + ) -> Box, Box<[u8]>)> + Send + 'a>; fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin + Send + 'a>>; @@ -142,7 +142,7 @@ impl Tree for SledEngineTree { fn scan_prefix<'a>( &'a self, prefix: Vec, - ) -> Box, Box<[u8]>)> + 'a> { + ) -> Box, Box<[u8]>)> + Send + 'a> { let iter = self .0 .scan_prefix(prefix) @@ -279,7 +279,7 @@ impl Tree for RocksDbEngineTree<'_> { fn scan_prefix<'a>( &'a self, prefix: Vec, - ) -> Box, Box<[u8]>)> + 'a> { + ) -> Box, Box<[u8]>)> + Send + 'a> { Box::new( self.db .0 diff --git a/src/database/globals.rs b/src/database/globals.rs index 37ebf136..1ca64de1 100644 --- a/src/database/globals.rs +++ b/src/database/globals.rs @@ -5,11 +5,7 @@ use ruma::{ EventId, MilliSecondsSinceUnixEpoch, ServerName, ServerSigningKeyId, }; use rustls::{ServerCertVerifier, WebPKIVerifier}; -use std::{ - collections::{BTreeMap, HashMap}, - sync::{Arc, RwLock}, - time::{Duration, Instant}, -}; +use std::{collections::{BTreeMap, HashMap}, path::{PathBuf}, sync::{Arc, RwLock}, time::{Duration, Instant}}; use tokio::sync::Semaphore; use trust_dns_resolver::TokioAsyncResolver; @@ -275,4 +271,19 @@ impl Globals { .insert(b"version", &new_version.to_be_bytes())?; Ok(()) } + + pub fn get_media_folder(&self) -> PathBuf { + let mut r = PathBuf::new(); + r.push(self.config.database_path.clone()); + r.push("media"); + r + } + + pub fn get_media_file(&self, key: &Vec) -> PathBuf { + let mut r = PathBuf::new(); + r.push(self.config.database_path.clone()); + r.push("media"); + r.push(base64::encode_config(key, base64::URL_SAFE_NO_PAD)); + r + } } diff --git a/src/database/media.rs b/src/database/media.rs index ca45484a..666a4943 100644 --- a/src/database/media.rs +++ b/src/database/media.rs @@ -1,9 +1,10 @@ use image::{imageops::FilterType, GenericImageView}; +use crate::database::globals::Globals; use crate::{utils, Error, Result}; use std::{mem, sync::Arc}; - use super::abstraction::Tree; +use tokio::{fs::{self, File}, io::AsyncWriteExt, io::AsyncReadExt}; pub struct FileMeta { pub content_disposition: Option, @@ -16,10 +17,11 @@ pub struct Media { } impl Media { - /// Uploads or replaces a file. - pub fn create( + /// Uploads a file. + pub async fn create( &self, mxc: String, + globals: &Globals, content_disposition: &Option<&str>, content_type: &Option<&str>, file: &[u8], @@ -43,15 +45,20 @@ impl Media { .unwrap_or_default(), ); - self.mediaid_file.insert(&key, file)?; + let path = globals.get_media_file(&key); + fs::create_dir_all(path.parent().unwrap()).await?; + let mut f = File::create(path).await?; + f.write_all(file).await?; + self.mediaid_file.insert(&key, &[])?; Ok(()) } /// Uploads or replaces a file thumbnail. - pub fn upload_thumbnail( + pub async fn upload_thumbnail( &self, mxc: String, + globals: &Globals, content_disposition: &Option, content_type: &Option, width: u32, @@ -77,20 +84,29 @@ impl Media { .unwrap_or_default(), ); - self.mediaid_file.insert(&key, file)?; + let path = globals.get_media_file(&key); + fs::create_dir_all(path.parent().unwrap()).await?; + let mut f = File::create(path).await?; + f.write_all(file).await?; + + self.mediaid_file.insert(&key, &[])?; Ok(()) } /// Downloads a file. - pub fn get(&self, mxc: &str) -> Result> { + pub async fn get(&self, globals: &Globals, mxc: &str) -> Result> { let mut prefix = mxc.as_bytes().to_vec(); prefix.push(0xff); prefix.extend_from_slice(&0_u32.to_be_bytes()); // Width = 0 if it's not a thumbnail prefix.extend_from_slice(&0_u32.to_be_bytes()); // Height = 0 if it's not a thumbnail prefix.push(0xff); - if let Some((key, file)) = self.mediaid_file.scan_prefix(prefix).next() { + let mut iter = self.mediaid_file.scan_prefix(prefix); + if let Some((key, _)) = iter.next() { + let path = globals.get_media_file(&key.to_vec()); + let mut file = vec![]; + File::open(path).await?.read_to_end(&mut file).await?; let mut parts = key.rsplit(|&b| b == 0xff); let content_type = parts @@ -121,7 +137,7 @@ impl Media { Ok(Some(FileMeta { content_disposition, content_type, - file: file.to_vec(), + file, })) } else { Ok(None) @@ -151,7 +167,7 @@ impl Media { /// - Server creates the thumbnail and sends it to the user /// /// For width,height <= 96 the server uses another thumbnailing algorithm which crops the image afterwards. - pub fn get_thumbnail(&self, mxc: String, width: u32, height: u32) -> Result> { + pub async fn get_thumbnail(&self, mxc: String, globals: &Globals, width: u32, height: u32) -> Result> { let (width, height, crop) = self .thumbnail_properties(width, height) .unwrap_or((0, 0, false)); // 0, 0 because that's the original file @@ -169,8 +185,11 @@ impl Media { original_prefix.extend_from_slice(&0_u32.to_be_bytes()); // Height = 0 if it's not a thumbnail original_prefix.push(0xff); - if let Some((key, file)) = self.mediaid_file.scan_prefix(thumbnail_prefix).next() { + if let Some((key, _)) = self.mediaid_file.scan_prefix(thumbnail_prefix).next() { // Using saved thumbnail + let path = globals.get_media_file(&key.to_vec()); + let mut file = vec![]; + File::open(path).await?.read_to_end(&mut file).await?; let mut parts = key.rsplit(|&b| b == 0xff); let content_type = parts @@ -201,8 +220,12 @@ impl Media { content_type, file: file.to_vec(), })) - } else if let Some((key, file)) = self.mediaid_file.scan_prefix(original_prefix).next() { + } else if let Some((key, _)) = self.mediaid_file.scan_prefix(original_prefix).next() { // Generate a thumbnail + let path = globals.get_media_file(&key.to_vec()); + let mut file = vec![]; + File::open(path).await?.read_to_end(&mut file).await?; + let mut parts = key.rsplit(|&b| b == 0xff); let content_type = parts @@ -299,19 +322,24 @@ impl Media { widthheight, ); - self.mediaid_file.insert(&thumbnail_key, &thumbnail_bytes)?; + let path = globals.get_media_file(&thumbnail_key); + fs::create_dir_all(path.parent().unwrap()).await?; + let mut f = File::create(path).await?; + f.write_all(&thumbnail_bytes).await?; + + self.mediaid_file.insert(&thumbnail_key, &[])?; Ok(Some(FileMeta { content_disposition, content_type, - file: thumbnail_bytes.to_vec(), + file: thumbnail_bytes.to_vec() })) } else { // Couldn't parse file to generate thumbnail, send original Ok(Some(FileMeta { content_disposition, content_type, - file: file.to_vec(), + file: file.to_vec() })) } } else { diff --git a/src/error.rs b/src/error.rs index 93c67c12..10a48b72 100644 --- a/src/error.rs +++ b/src/error.rs @@ -45,6 +45,11 @@ pub enum Error { }, #[error("{0}")] FederationError(Box, RumaError), + #[error("Could not do this io: {source}")] + IoError { + #[from] + source: std::io::Error, + }, #[error("{0}")] BadServerResponse(&'static str), #[error("{0}")] From 804105479c7ae275a79c9fdeeb32e7a20b22fa3b Mon Sep 17 00:00:00 2001 From: hamidreza kalbasi Date: Sun, 6 Jun 2021 16:58:32 +0430 Subject: [PATCH 3/8] fix fmt and clippy warnings --- src/client_server/media.rs | 85 +++++++++++++++++++++----------------- src/database/globals.rs | 9 +++- src/database/media.rs | 26 ++++++++---- 3 files changed, 73 insertions(+), 47 deletions(-) diff --git a/src/client_server/media.rs b/src/client_server/media.rs index 14ab6db0..0b1fbd7b 100644 --- a/src/client_server/media.rs +++ b/src/client_server/media.rs @@ -36,17 +36,20 @@ pub async fn create_content_route( db.globals.server_name(), utils::random_string(MXC_LENGTH) ); - db.media.create( - mxc.clone(), - &db.globals, - &body - .filename - .as_ref() - .map(|filename| "inline; filename=".to_owned() + filename) - .as_deref(), - &body.content_type.as_deref(), - &body.file, - ).await?; + + db.media + .create( + mxc.clone(), + &db.globals, + &body + .filename + .as_ref() + .map(|filename| "inline; filename=".to_owned() + filename) + .as_deref(), + &body.content_type.as_deref(), + &body.file, + ) + .await?; db.flush().await?; @@ -94,13 +97,15 @@ pub async fn get_content_route( ) .await?; - db.media.create( - mxc, - &db.globals, - &get_content_response.content_disposition.as_deref(), - &get_content_response.content_type.as_deref(), - &get_content_response.file, - ).await?; + db.media + .create( + mxc, + &db.globals, + &get_content_response.content_disposition.as_deref(), + &get_content_response.content_type.as_deref(), + &get_content_response.file, + ) + .await?; Ok(get_content_response.into()) } else { @@ -121,16 +126,20 @@ pub async fn get_content_thumbnail_route( if let Some(FileMeta { content_type, file, .. - }) = db.media.get_thumbnail( - mxc.clone(), - &db.globals, - body.width - .try_into() - .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Width is invalid."))?, - body.height - .try_into() - .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Width is invalid."))?, - ).await? { + }) = db + .media + .get_thumbnail( + mxc.clone(), + &db.globals, + body.width + .try_into() + .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Width is invalid."))?, + body.height + .try_into() + .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Width is invalid."))?, + ) + .await? + { Ok(get_content_thumbnail::Response { file, content_type }.into()) } else if &*body.server_name != db.globals.server_name() && body.allow_remote { let get_thumbnail_response = db @@ -149,15 +158,17 @@ pub async fn get_content_thumbnail_route( ) .await?; - db.media.upload_thumbnail( - mxc, - &db.globals, - &None, - &get_thumbnail_response.content_type, - body.width.try_into().expect("all UInts are valid u32s"), - body.height.try_into().expect("all UInts are valid u32s"), - &get_thumbnail_response.file, - ).await?; + db.media + .upload_thumbnail( + mxc, + &db.globals, + &None, + &get_thumbnail_response.content_type, + body.width.try_into().expect("all UInts are valid u32s"), + body.height.try_into().expect("all UInts are valid u32s"), + &get_thumbnail_response.file, + ) + .await?; Ok(get_thumbnail_response.into()) } else { diff --git a/src/database/globals.rs b/src/database/globals.rs index 1ca64de1..de97578b 100644 --- a/src/database/globals.rs +++ b/src/database/globals.rs @@ -5,7 +5,12 @@ use ruma::{ EventId, MilliSecondsSinceUnixEpoch, ServerName, ServerSigningKeyId, }; use rustls::{ServerCertVerifier, WebPKIVerifier}; -use std::{collections::{BTreeMap, HashMap}, path::{PathBuf}, sync::{Arc, RwLock}, time::{Duration, Instant}}; +use std::{ + collections::{BTreeMap, HashMap}, + path::PathBuf, + sync::{Arc, RwLock}, + time::{Duration, Instant}, +}; use tokio::sync::Semaphore; use trust_dns_resolver::TokioAsyncResolver; @@ -279,7 +284,7 @@ impl Globals { r } - pub fn get_media_file(&self, key: &Vec) -> PathBuf { + pub fn get_media_file(&self, key: &[u8]) -> PathBuf { let mut r = PathBuf::new(); r.push(self.config.database_path.clone()); r.push("media"); diff --git a/src/database/media.rs b/src/database/media.rs index 666a4943..6fd142d3 100644 --- a/src/database/media.rs +++ b/src/database/media.rs @@ -1,10 +1,14 @@ -use image::{imageops::FilterType, GenericImageView}; use crate::database::globals::Globals; +use image::{imageops::FilterType, GenericImageView}; +use super::abstraction::Tree; use crate::{utils, Error, Result}; use std::{mem, sync::Arc}; -use super::abstraction::Tree; -use tokio::{fs::{self, File}, io::AsyncWriteExt, io::AsyncReadExt}; +use tokio::{ + fs::{self, File}, + io::AsyncReadExt, + io::AsyncWriteExt, +}; pub struct FileMeta { pub content_disposition: Option, @@ -167,7 +171,13 @@ impl Media { /// - Server creates the thumbnail and sends it to the user /// /// For width,height <= 96 the server uses another thumbnailing algorithm which crops the image afterwards. - pub async fn get_thumbnail(&self, mxc: String, globals: &Globals, width: u32, height: u32) -> Result> { + pub async fn get_thumbnail( + &self, + mxc: String, + globals: &Globals, + width: u32, + height: u32, + ) -> Result> { let (width, height, crop) = self .thumbnail_properties(width, height) .unwrap_or((0, 0, false)); // 0, 0 because that's the original file @@ -225,7 +235,7 @@ impl Media { let path = globals.get_media_file(&key.to_vec()); let mut file = vec![]; File::open(path).await?.read_to_end(&mut file).await?; - + let mut parts = key.rsplit(|&b| b == 0xff); let content_type = parts @@ -326,20 +336,20 @@ impl Media { fs::create_dir_all(path.parent().unwrap()).await?; let mut f = File::create(path).await?; f.write_all(&thumbnail_bytes).await?; - + self.mediaid_file.insert(&thumbnail_key, &[])?; Ok(Some(FileMeta { content_disposition, content_type, - file: thumbnail_bytes.to_vec() + file: thumbnail_bytes.to_vec(), })) } else { // Couldn't parse file to generate thumbnail, send original Ok(Some(FileMeta { content_disposition, content_type, - file: file.to_vec() + file: file.to_vec(), })) } } else { From ff841b73c54fca06a604f7c99d48321965da053b Mon Sep 17 00:00:00 2001 From: hamidreza kalbasi Date: Tue, 8 Jun 2021 17:05:13 +0430 Subject: [PATCH 4/8] use .keys() and remove unneccery .to_vec() --- src/database.rs | 2 +- src/database/media.rs | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/database.rs b/src/database.rs index 04d37358..2870e2d1 100644 --- a/src/database.rs +++ b/src/database.rs @@ -16,7 +16,7 @@ pub mod users; use crate::{utils, Error, Result}; use abstraction::DatabaseEngine; use directories::ProjectDirs; -use log::{error, info}; +use log::error; use rocket::futures::{channel::mpsc, stream::FuturesUnordered, StreamExt}; use ruma::{DeviceId, ServerName, UserId}; use serde::Deserialize; diff --git a/src/database/media.rs b/src/database/media.rs index 6fd142d3..205f3a65 100644 --- a/src/database/media.rs +++ b/src/database/media.rs @@ -108,7 +108,7 @@ impl Media { let mut iter = self.mediaid_file.scan_prefix(prefix); if let Some((key, _)) = iter.next() { - let path = globals.get_media_file(&key.to_vec()); + let path = globals.get_media_file(&key); let mut file = vec![]; File::open(path).await?.read_to_end(&mut file).await?; let mut parts = key.rsplit(|&b| b == 0xff); @@ -197,7 +197,7 @@ impl Media { if let Some((key, _)) = self.mediaid_file.scan_prefix(thumbnail_prefix).next() { // Using saved thumbnail - let path = globals.get_media_file(&key.to_vec()); + let path = globals.get_media_file(&key); let mut file = vec![]; File::open(path).await?.read_to_end(&mut file).await?; let mut parts = key.rsplit(|&b| b == 0xff); @@ -232,7 +232,7 @@ impl Media { })) } else if let Some((key, _)) = self.mediaid_file.scan_prefix(original_prefix).next() { // Generate a thumbnail - let path = globals.get_media_file(&key.to_vec()); + let path = globals.get_media_file(&key); let mut file = vec![]; File::open(path).await?.read_to_end(&mut file).await?; From affa124864d30eed4453416077db40b63b98e479 Mon Sep 17 00:00:00 2001 From: hamidreza kalbasi Date: Tue, 8 Jun 2021 17:50:06 +0430 Subject: [PATCH 5/8] create media folder in init --- src/database/globals.rs | 9 +++++++-- src/database/media.rs | 9 +-------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/database/globals.rs b/src/database/globals.rs index de97578b..55256443 100644 --- a/src/database/globals.rs +++ b/src/database/globals.rs @@ -7,6 +7,7 @@ use ruma::{ use rustls::{ServerCertVerifier, WebPKIVerifier}; use std::{ collections::{BTreeMap, HashMap}, + fs, path::PathBuf, sync::{Arc, RwLock}, time::{Duration, Instant}, @@ -137,7 +138,7 @@ impl Globals { .as_ref() .map(|secret| jsonwebtoken::DecodingKey::from_secret(secret.as_bytes()).into_static()); - Ok(Self { + let s = Self { globals, config, keypair: Arc::new(keypair), @@ -152,7 +153,11 @@ impl Globals { bad_event_ratelimiter: Arc::new(RwLock::new(BTreeMap::new())), bad_signature_ratelimiter: Arc::new(RwLock::new(BTreeMap::new())), servername_ratelimiter: Arc::new(RwLock::new(BTreeMap::new())), - }) + }; + + fs::create_dir_all(s.get_media_folder())?; + + Ok(s) } /// Returns this server's keypair. diff --git a/src/database/media.rs b/src/database/media.rs index 205f3a65..b81e8633 100644 --- a/src/database/media.rs +++ b/src/database/media.rs @@ -4,11 +4,7 @@ use image::{imageops::FilterType, GenericImageView}; use super::abstraction::Tree; use crate::{utils, Error, Result}; use std::{mem, sync::Arc}; -use tokio::{ - fs::{self, File}, - io::AsyncReadExt, - io::AsyncWriteExt, -}; +use tokio::{fs::File, io::AsyncReadExt, io::AsyncWriteExt}; pub struct FileMeta { pub content_disposition: Option, @@ -50,7 +46,6 @@ impl Media { ); let path = globals.get_media_file(&key); - fs::create_dir_all(path.parent().unwrap()).await?; let mut f = File::create(path).await?; f.write_all(file).await?; @@ -89,7 +84,6 @@ impl Media { ); let path = globals.get_media_file(&key); - fs::create_dir_all(path.parent().unwrap()).await?; let mut f = File::create(path).await?; f.write_all(file).await?; @@ -333,7 +327,6 @@ impl Media { ); let path = globals.get_media_file(&thumbnail_key); - fs::create_dir_all(path.parent().unwrap()).await?; let mut f = File::create(path).await?; f.write_all(&thumbnail_bytes).await?; From 2385bd1cfd41fe14e8a73b9d6b556b5f8bb55ff1 Mon Sep 17 00:00:00 2001 From: hamidreza kalbasi Date: Tue, 8 Jun 2021 20:53:24 +0430 Subject: [PATCH 6/8] add migrations --- src/database.rs | 29 ++++++++++++++++++++++------- src/database/media.rs | 6 +++--- 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/src/database.rs b/src/database.rs index 2870e2d1..fa84881d 100644 --- a/src/database.rs +++ b/src/database.rs @@ -20,11 +20,7 @@ use log::error; use rocket::futures::{channel::mpsc, stream::FuturesUnordered, StreamExt}; use ruma::{DeviceId, ServerName, UserId}; use serde::Deserialize; -use std::{ - collections::HashMap, - fs::remove_dir_all, - sync::{Arc, RwLock}, -}; +use std::{collections::HashMap, fs::{self, remove_dir_all}, io::Write, sync::{Arc, RwLock}}; use tokio::sync::Semaphore; #[derive(Clone, Debug, Deserialize)] @@ -253,9 +249,11 @@ impl Database { for (userid, password) in db.users.userid_password.iter() { let password = utils::string_from_bytes(&password); - if password.map_or(false, |password| { + let empty_hashed_password = password.map_or(false, |password| { argon2::verify_encoded(&password, b"").unwrap_or(false) - }) { + }); + + if empty_hashed_password { db.users.userid_password.insert(&userid, b"")?; } } @@ -265,6 +263,23 @@ impl Database { println!("Migration: 1 -> 2 finished"); } + if db.globals.database_version()? < 3 { + // Move media to filesystem + for (key, content) in db.media.mediaid_file.iter() { + if content.len() == 0 { + continue; + } + + let path = db.globals.get_media_file(&key); + let mut file = fs::File::create(path)?; + file.write_all(&content)?; + db.media.mediaid_file.insert(&key, &[])?; + } + + db.globals.bump_database_version(3)?; + + println!("Migration: 2 -> 3 finished"); + } // This data is probably outdated db.rooms.edus.presenceid_presence.clear()?; diff --git a/src/database/media.rs b/src/database/media.rs index b81e8633..944c5bdc 100644 --- a/src/database/media.rs +++ b/src/database/media.rs @@ -103,7 +103,7 @@ impl Media { let mut iter = self.mediaid_file.scan_prefix(prefix); if let Some((key, _)) = iter.next() { let path = globals.get_media_file(&key); - let mut file = vec![]; + let mut file = Vec::new(); File::open(path).await?.read_to_end(&mut file).await?; let mut parts = key.rsplit(|&b| b == 0xff); @@ -192,7 +192,7 @@ impl Media { if let Some((key, _)) = self.mediaid_file.scan_prefix(thumbnail_prefix).next() { // Using saved thumbnail let path = globals.get_media_file(&key); - let mut file = vec![]; + let mut file = Vec::new(); File::open(path).await?.read_to_end(&mut file).await?; let mut parts = key.rsplit(|&b| b == 0xff); @@ -227,7 +227,7 @@ impl Media { } else if let Some((key, _)) = self.mediaid_file.scan_prefix(original_prefix).next() { // Generate a thumbnail let path = globals.get_media_file(&key); - let mut file = vec![]; + let mut file = Vec::new(); File::open(path).await?.read_to_end(&mut file).await?; let mut parts = key.rsplit(|&b| b == 0xff); From 2078af59d8b2a95fbaf522962265c222d3f17b1b Mon Sep 17 00:00:00 2001 From: hamidreza kalbasi Date: Tue, 8 Jun 2021 20:54:36 +0430 Subject: [PATCH 7/8] fix fmt problems --- src/database.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/database.rs b/src/database.rs index fa84881d..b5a25eac 100644 --- a/src/database.rs +++ b/src/database.rs @@ -20,7 +20,12 @@ use log::error; use rocket::futures::{channel::mpsc, stream::FuturesUnordered, StreamExt}; use ruma::{DeviceId, ServerName, UserId}; use serde::Deserialize; -use std::{collections::HashMap, fs::{self, remove_dir_all}, io::Write, sync::{Arc, RwLock}}; +use std::{ + collections::HashMap, + fs::{self, remove_dir_all}, + io::Write, + sync::{Arc, RwLock}, +}; use tokio::sync::Semaphore; #[derive(Clone, Debug, Deserialize)] From cd4bc520d8ef5410870cc82e967a3d480c5ca8c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6sters?= Date: Sat, 12 Jun 2021 15:04:28 +0200 Subject: [PATCH 8/8] improvement: feature flags for sled, rocksdb --- Cargo.toml | 8 ++-- src/client_server/membership.rs | 2 +- src/database.rs | 4 ++ src/database/abstraction.rs | 68 +++++++++++++++++++++------------ src/database/rooms.rs | 13 +++---- src/database/sending.rs | 2 - src/error.rs | 2 + src/ruma_wrapper.rs | 16 ++++---- src/server_server.rs | 2 +- src/utils.rs | 4 +- 10 files changed, 71 insertions(+), 50 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index eb43da5a..e7ebadf2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,8 +24,8 @@ ruma = { git = "https://github.com/ruma/ruma", rev = "b39537812c12caafcbf8b7bd74 # Used for long polling and federation sender, should be the same as rocket::tokio tokio = "1.2.0" # Used for storing data permanently -sled = { version = "0.34.6", features = ["compression", "no_metrics"] } -rocksdb = { version = "0.16.0", features = ["multi-threaded-cf"] } +sled = { version = "0.34.6", features = ["compression", "no_metrics"], optional = true } +rocksdb = { version = "0.16.0", features = ["multi-threaded-cf"], optional = true } #sled = { git = "https://github.com/spacejam/sled.git", rev = "e4640e0773595229f398438886f19bca6f7326a2", features = ["compression"] } # Used for the http request / response body type for Ruma endpoints used with reqwest @@ -75,7 +75,9 @@ opentelemetry-jaeger = "0.11.0" pretty_env_logger = "0.4.0" [features] -default = ["conduit_bin"] +default = ["conduit_bin", "backend_sled"] +backend_sled = ["sled"] +backend_rocksdb = ["rocksdb"] conduit_bin = [] # TODO: add rocket to this when it is optional [[bin]] diff --git a/src/client_server/membership.rs b/src/client_server/membership.rs index 92d7aced..a3f1389a 100644 --- a/src/client_server/membership.rs +++ b/src/client_server/membership.rs @@ -621,7 +621,7 @@ async fn join_room_by_id_helper( &pdu, utils::to_canonical_object(&pdu).expect("Pdu is valid canonical object"), count, - pdu_id.into(), + &pdu_id, &[pdu.event_id.clone()], db, )?; diff --git a/src/database.rs b/src/database.rs index b5a25eac..e00bdcd5 100644 --- a/src/database.rs +++ b/src/database.rs @@ -77,8 +77,12 @@ fn default_log() -> String { "info,state_res=warn,rocket=off,_=off,sled=off".to_owned() } +#[cfg(feature = "sled")] pub type Engine = abstraction::SledEngine; +#[cfg(feature = "rocksdb")] +pub type Engine = abstraction::RocksDbEngine; + pub struct Database { pub globals: globals::Globals, pub users: users::Users, diff --git a/src/database/abstraction.rs b/src/database/abstraction.rs index ad032fb3..f81c9def 100644 --- a/src/database/abstraction.rs +++ b/src/database/abstraction.rs @@ -1,21 +1,19 @@ -use std::{ - collections::BTreeMap, - future::Future, - pin::Pin, - sync::{Arc, RwLock}, -}; - -use log::warn; -use rocksdb::{ - BoundColumnFamily, ColumnFamilyDescriptor, DBWithThreadMode, Direction, MultiThreaded, Options, -}; - use super::Config; use crate::{utils, Result}; +use log::warn; +use std::{future::Future, pin::Pin, sync::Arc}; + +#[cfg(feature = "rocksdb")] +use std::{collections::BTreeMap, sync::RwLock}; +#[cfg(feature = "sled")] pub struct SledEngine(sled::Db); +#[cfg(feature = "sled")] pub struct SledEngineTree(sled::Tree); -pub struct RocksDbEngine(rocksdb::DBWithThreadMode); + +#[cfg(feature = "rocksdb")] +pub struct RocksDbEngine(rocksdb::DBWithThreadMode); +#[cfg(feature = "rocksdb")] pub struct RocksDbEngineTree<'a> { db: Arc, name: &'a str, @@ -60,6 +58,7 @@ pub trait Tree: Send + Sync { } } +#[cfg(feature = "sled")] impl DatabaseEngine for SledEngine { fn open(config: &Config) -> Result> { Ok(Arc::new(SledEngine( @@ -76,6 +75,7 @@ impl DatabaseEngine for SledEngine { } } +#[cfg(feature = "sled")] impl Tree for SledEngineTree { fn get(&self, key: &[u8]) -> Result>> { Ok(self.0.get(key)?.map(|v| v.to_vec())) @@ -165,29 +165,42 @@ impl Tree for SledEngineTree { } } +#[cfg(feature = "rocksdb")] impl DatabaseEngine for RocksDbEngine { fn open(config: &Config) -> Result> { - let mut db_opts = Options::default(); + let mut db_opts = rocksdb::Options::default(); db_opts.create_if_missing(true); + db_opts.set_max_open_files(16); + db_opts.set_compaction_style(rocksdb::DBCompactionStyle::Level); + db_opts.set_compression_type(rocksdb::DBCompressionType::Snappy); + db_opts.set_target_file_size_base(256 << 20); + db_opts.set_write_buffer_size(256 << 20); + + let mut block_based_options = rocksdb::BlockBasedOptions::default(); + block_based_options.set_block_size(512 << 10); + db_opts.set_block_based_table_factory(&block_based_options); - let cfs = DBWithThreadMode::::list_cf(&db_opts, &config.database_path) - .unwrap_or_default(); + let cfs = rocksdb::DBWithThreadMode::::list_cf( + &db_opts, + &config.database_path, + ) + .unwrap_or_default(); - let mut options = Options::default(); + let mut options = rocksdb::Options::default(); options.set_merge_operator_associative("increment", utils::increment_rocksdb); - let db = DBWithThreadMode::::open_cf_descriptors( + let db = rocksdb::DBWithThreadMode::::open_cf_descriptors( &db_opts, &config.database_path, cfs.iter() - .map(|name| ColumnFamilyDescriptor::new(name, options.clone())), + .map(|name| rocksdb::ColumnFamilyDescriptor::new(name, options.clone())), )?; Ok(Arc::new(RocksDbEngine(db))) } fn open_tree(self: &Arc, name: &'static str) -> Result> { - let mut options = Options::default(); + let mut options = rocksdb::Options::default(); options.set_merge_operator_associative("increment", utils::increment_rocksdb); // Create if it doesn't exist @@ -201,12 +214,14 @@ impl DatabaseEngine for RocksDbEngine { } } +#[cfg(feature = "rocksdb")] impl RocksDbEngineTree<'_> { - fn cf(&self) -> BoundColumnFamily<'_> { + fn cf(&self) -> rocksdb::BoundColumnFamily<'_> { self.db.0.cf_handle(self.name).unwrap() } } +#[cfg(feature = "rocksdb")] impl Tree for RocksDbEngineTree<'_> { fn get(&self, key: &[u8]) -> Result>> { Ok(self.db.0.get_cf(self.cf(), key)?) @@ -260,15 +275,20 @@ impl Tree for RocksDbEngineTree<'_> { rocksdb::IteratorMode::From( from, if backwards { - Direction::Reverse + rocksdb::Direction::Reverse } else { - Direction::Forward + rocksdb::Direction::Forward }, ), )) } fn increment(&self, key: &[u8]) -> Result> { + let stats = rocksdb::perf::get_memory_usage_stats(Some(&[&self.db.0]), None).unwrap(); + dbg!(stats.mem_table_total); + dbg!(stats.mem_table_unflushed); + dbg!(stats.mem_table_readers_total); + dbg!(stats.cache_total); // TODO: atomic? let old = self.get(key)?; let new = utils::increment(old.as_deref()).unwrap(); @@ -285,7 +305,7 @@ impl Tree for RocksDbEngineTree<'_> { .0 .iterator_cf( self.cf(), - rocksdb::IteratorMode::From(&prefix, Direction::Forward), + rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward), ) .take_while(move |(k, _)| k.starts_with(&prefix)), ) diff --git a/src/database/rooms.rs b/src/database/rooms.rs index 0a8239d4..736ff4d8 100644 --- a/src/database/rooms.rs +++ b/src/database/rooms.rs @@ -19,8 +19,6 @@ use ruma::{ state_res::{self, Event, RoomVersion, StateMap}, uint, EventId, RoomAliasId, RoomId, RoomVersionId, ServerName, UserId, }; -use sled::IVec; - use std::{ collections::{BTreeMap, HashMap, HashSet}, convert::{TryFrom, TryInto}, @@ -34,7 +32,7 @@ use super::{abstraction::Tree, admin::AdminCommand, pusher}; /// /// This is created when a state group is added to the database by /// hashing the entire state. -pub type StateHashId = IVec; +pub type StateHashId = Vec; pub struct Rooms { pub edus: edus::RoomEdus, @@ -665,7 +663,7 @@ impl Rooms { pdu: &PduEvent, mut pdu_json: CanonicalJsonObject, count: u64, - pdu_id: IVec, + pdu_id: &[u8], leaves: &[EventId], db: &Database, ) -> Result<()> { @@ -713,14 +711,13 @@ impl Rooms { self.reset_notification_counts(&pdu.sender, &pdu.room_id)?; self.pduid_pdu.insert( - &pdu_id, + pdu_id, &serde_json::to_vec(&pdu_json).expect("CanonicalJsonObject is always a valid"), )?; // This also replaces the eventid of any outliers with the correct // pduid, removing the place holder. - self.eventid_pduid - .insert(pdu.event_id.as_bytes(), &*pdu_id)?; + self.eventid_pduid.insert(pdu.event_id.as_bytes(), pdu_id)?; // See if the event matches any known pushers for user in db @@ -1360,7 +1357,7 @@ impl Rooms { &pdu, pdu_json, count, - pdu_id.clone().into(), + &pdu_id, // Since this PDU references all pdu_leaves we can update the leaves // of the room &[pdu.event_id.clone()], diff --git a/src/database/sending.rs b/src/database/sending.rs index 77f6ed78..ecf07618 100644 --- a/src/database/sending.rs +++ b/src/database/sending.rs @@ -91,8 +91,6 @@ enum TransactionStatus { impl Sending { pub fn start_handler(&self, db: Arc, mut receiver: mpsc::UnboundedReceiver>) { - let db = db.clone(); - tokio::spawn(async move { let mut futures = FuturesUnordered::new(); diff --git a/src/error.rs b/src/error.rs index 10a48b72..4f363fff 100644 --- a/src/error.rs +++ b/src/error.rs @@ -23,11 +23,13 @@ pub type Result = std::result::Result; #[derive(Error, Debug)] pub enum Error { + #[cfg(feature = "sled")] #[error("There was a problem with the connection to the sled database.")] SledError { #[from] source: sled::Error, }, + #[cfg(feature = "rocksdb")] #[error("There was a problem with the connection to the rocksdb database: {source}")] RocksDbError { #[from] diff --git a/src/ruma_wrapper.rs b/src/ruma_wrapper.rs index ba2c37ed..2912a578 100644 --- a/src/ruma_wrapper.rs +++ b/src/ruma_wrapper.rs @@ -1,14 +1,15 @@ -use crate::{Database, Error}; +use crate::Error; use ruma::{ api::OutgoingResponse, identifiers::{DeviceId, UserId}, - Outgoing, + signatures::CanonicalJsonValue, + Outgoing, ServerName, }; -use std::{ops::Deref, sync::Arc}; +use std::ops::Deref; #[cfg(feature = "conduit_bin")] use { - crate::server_server, + crate::{server_server, Database}, log::{debug, warn}, rocket::{ data::{self, ByteUnit, Data, FromData}, @@ -18,14 +19,11 @@ use { tokio::io::AsyncReadExt, Request, State, }, - ruma::{ - api::{AuthScheme, IncomingRequest}, - signatures::CanonicalJsonValue, - ServerName, - }, + ruma::api::{AuthScheme, IncomingRequest}, std::collections::BTreeMap, std::convert::TryFrom, std::io::Cursor, + std::sync::Arc, }; /// This struct converts rocket requests into ruma structs by converting them into http requests diff --git a/src/server_server.rs b/src/server_server.rs index 7a338dc5..2a445c2b 100644 --- a/src/server_server.rs +++ b/src/server_server.rs @@ -1681,7 +1681,7 @@ pub(crate) fn append_incoming_pdu( pdu, pdu_json, count, - pdu_id.clone().into(), + &pdu_id, &new_room_leaves.into_iter().collect::>(), &db, )?; diff --git a/src/utils.rs b/src/utils.rs index f59afb3a..0c8fb5ca 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,7 +1,6 @@ use argon2::{Config, Variant}; use cmp::Ordering; use rand::prelude::*; -use rocksdb::MergeOperands; use ruma::serde::{try_from_json_map, CanonicalJsonError, CanonicalJsonObject}; use std::{ cmp, @@ -16,10 +15,11 @@ pub fn millis_since_unix_epoch() -> u64 { .as_millis() as u64 } +#[cfg(feature = "rocksdb")] pub fn increment_rocksdb( _new_key: &[u8], old: Option<&[u8]>, - _operands: &mut MergeOperands, + _operands: &mut rocksdb::MergeOperands, ) -> Option> { increment(old) }