magmaus3/continuwuity
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merged current next

Browse source
This commit is contained in:
Torsten Flammiger 2022-01-21 14:32:59 +01:00
commit 960ba8bd99
38 changed files with 1895 additions and 894 deletions
Download patch file Download diff file Expand all files Collapse all files

22
src/database/abstraction.rs
View file

@ -12,10 +12,24 @@ pub mod sqlite;
#[cfg(feature = "heed")]
pub mod heed;
pub trait DatabaseEngine: Sized {
fn open(config: &Config) -> Result<Arc<Self>>;
fn open_tree(self: &Arc<Self>, name: &'static str) -> Result<Arc<dyn Tree>>;
fn flush(self: &Arc<Self>) -> Result<()>;
#[cfg(feature = "rocksdb")]
pub mod rocksdb;
#[cfg(any(feature = "sqlite", feature = "rocksdb", feature = "heed"))]
pub mod watchers;
pub trait DatabaseEngine: Send + Sync {
fn open(config: &Config) -> Result<Self>
where
Self: Sized;
fn open_tree(&self, name: &'static str) -> Result<Arc<dyn Tree>>;
fn flush(&self) -> Result<()>;
fn cleanup(&self) -> Result<()> {
Ok(())
}
fn memory_usage(&self) -> Result<String> {
Ok("Current database engine does not support memory usage reporting.".to_owned())
}
}
pub trait Tree: Send + Sync {

48
src/database/abstraction/heed.rs
View file

@ -1,15 +1,13 @@
use super::super::Config;
use super::{super::Config, watchers::Watchers};
use crossbeam::channel::{bounded, Sender as ChannelSender};
use threadpool::ThreadPool;
use crate::{Error, Result};
use std::{
collections::HashMap,
future::Future,
pin::Pin,
sync::{Arc, Mutex, RwLock},
sync::{Arc, Mutex},
};
use tokio::sync::oneshot::Sender;
use super::{DatabaseEngine, Tree};
@ -23,7 +21,7 @@ pub struct Engine {
pub struct EngineTree {
engine: Arc<Engine>,
tree: Arc<heed::UntypedDatabase>,
watchers: RwLock<HashMap<Vec<u8>, Vec<Sender<()>>>>,
watchers: Watchers,
}
fn convert_error(error: heed::Error) -> Error {
@ -60,7 +58,7 @@ impl DatabaseEngine for Engine {
.create_database(Some(name))
.map_err(convert_error)?,
),
watchers: RwLock::new(HashMap::new()),
watchers: Default::default(),
}))
}
@ -145,29 +143,7 @@ impl Tree for EngineTree {
.put(&mut txn, &key, &value)
.map_err(convert_error)?;
txn.commit().map_err(convert_error)?;
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(());
}
}
}
};
self.watchers.wake(key);
Ok(())
}
@ -223,18 +199,6 @@ impl Tree for EngineTree {
#[tracing::instrument(skip(self, prefix))]
fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin<Box<dyn Future<Output = ()> + 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();
})
self.watchers.watch(prefix)
}
}

231
src/database/abstraction/rocksdb.rs Normal file
View file

@ -0,0 +1,231 @@
use super::{super::Config, watchers::Watchers, DatabaseEngine, Tree};
use crate::{utils, Result};
use std::{future::Future, pin::Pin, sync::Arc, sync::RwLock};
pub struct Engine {
rocks: rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>,
max_open_files: i32,
cache: rocksdb::Cache,
old_cfs: Vec<String>,
}
pub struct RocksDbEngineTree<'a> {
db: Arc<Engine>,
name: &'a str,
watchers: Watchers,
write_lock: RwLock<()>,
}
fn db_options(max_open_files: i32, rocksdb_cache: &rocksdb::Cache) -> rocksdb::Options {
let mut block_based_options = rocksdb::BlockBasedOptions::default();
block_based_options.set_block_cache(rocksdb_cache);
// "Difference of spinning disk"
// https://zhangyuchi.gitbooks.io/rocksdbbook/content/RocksDB-Tuning-Guide.html
block_based_options.set_block_size(4 * 1024);
block_based_options.set_cache_index_and_filter_blocks(true);
let mut db_opts = rocksdb::Options::default();
db_opts.set_block_based_table_factory(&block_based_options);
db_opts.set_optimize_filters_for_hits(true);
db_opts.set_skip_stats_update_on_db_open(true);
db_opts.set_level_compaction_dynamic_level_bytes(true);
db_opts.set_target_file_size_base(256 * 1024 * 1024);
//db_opts.set_compaction_readahead_size(2 * 1024 * 1024);
//db_opts.set_use_direct_reads(true);
//db_opts.set_use_direct_io_for_flush_and_compaction(true);
db_opts.create_if_missing(true);
db_opts.increase_parallelism(num_cpus::get() as i32);
db_opts.set_max_open_files(max_open_files);
db_opts.set_compression_type(rocksdb::DBCompressionType::Zstd);
db_opts.set_compaction_style(rocksdb::DBCompactionStyle::Level);
db_opts.optimize_level_style_compaction(10 * 1024 * 1024);
let prefix_extractor = rocksdb::SliceTransform::create_fixed_prefix(1);
db_opts.set_prefix_extractor(prefix_extractor);
db_opts
}
impl DatabaseEngine for Arc<Engine> {
fn open(config: &Config) -> Result<Self> {
let cache_capacity_bytes = (config.db_cache_capacity_mb * 1024.0 * 1024.0) as usize;
let rocksdb_cache = rocksdb::Cache::new_lru_cache(cache_capacity_bytes).unwrap();
let db_opts = db_options(config.rocksdb_max_open_files, &rocksdb_cache);
let cfs = rocksdb::DBWithThreadMode::<rocksdb::MultiThreaded>::list_cf(
&db_opts,
&config.database_path,
)
.unwrap_or_default();
let db = rocksdb::DBWithThreadMode::<rocksdb::MultiThreaded>::open_cf_descriptors(
&db_opts,
&config.database_path,
cfs.iter().map(|name| {
rocksdb::ColumnFamilyDescriptor::new(
name,
db_options(config.rocksdb_max_open_files, &rocksdb_cache),
)
}),
)?;
Ok(Arc::new(Engine {
rocks: db,
max_open_files: config.rocksdb_max_open_files,
cache: rocksdb_cache,
old_cfs: cfs,
}))
}
fn open_tree(&self, name: &'static str) -> Result<Arc<dyn Tree>> {
if !self.old_cfs.contains(&name.to_owned()) {
// Create if it didn't exist
let _ = self
.rocks
.create_cf(name, &db_options(self.max_open_files, &self.cache));
}
Ok(Arc::new(RocksDbEngineTree {
name,
db: Arc::clone(self),
watchers: Watchers::default(),
write_lock: RwLock::new(()),
}))
}
fn flush(&self) -> Result<()> {
// TODO?
Ok(())
}
fn memory_usage(&self) -> Result<String> {
let stats =
rocksdb::perf::get_memory_usage_stats(Some(&[&self.rocks]), Some(&[&self.cache]))?;
Ok(format!(
"Approximate memory usage of all the mem-tables: {:.3} MB\n\
Approximate memory usage of un-flushed mem-tables: {:.3} MB\n\
Approximate memory usage of all the table readers: {:.3} MB\n\
Approximate memory usage by cache: {:.3} MB\n\
Approximate memory usage by cache pinned: {:.3} MB\n\
",
stats.mem_table_total as f64 / 1024.0 / 1024.0,
stats.mem_table_unflushed as f64 / 1024.0 / 1024.0,
stats.mem_table_readers_total as f64 / 1024.0 / 1024.0,
stats.cache_total as f64 / 1024.0 / 1024.0,
self.cache.get_pinned_usage() as f64 / 1024.0 / 1024.0,
))
}
}
impl RocksDbEngineTree<'_> {
fn cf(&self) -> Arc<rocksdb::BoundColumnFamily<'_>> {
self.db.rocks.cf_handle(self.name).unwrap()
}
}
impl Tree for RocksDbEngineTree<'_> {
fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
Ok(self.db.rocks.get_cf(&self.cf(), key)?)
}
fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> {
let lock = self.write_lock.read().unwrap();
self.db.rocks.put_cf(&self.cf(), key, value)?;
drop(lock);
self.watchers.wake(key);
Ok(())
}
fn insert_batch<'a>(&self, iter: &mut dyn Iterator<Item = (Vec<u8>, Vec<u8>)>) -> Result<()> {
for (key, value) in iter {
self.db.rocks.put_cf(&self.cf(), key, value)?;
}
Ok(())
}
fn remove(&self, key: &[u8]) -> Result<()> {
Ok(self.db.rocks.delete_cf(&self.cf(), key)?)
}
fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
Box::new(
self.db
.rocks
.iterator_cf(&self.cf(), rocksdb::IteratorMode::Start)
.map(|(k, v)| (Vec::from(k), Vec::from(v))),
)
}
fn iter_from<'a>(
&'a self,
from: &[u8],
backwards: bool,
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
Box::new(
self.db
.rocks
.iterator_cf(
&self.cf(),
rocksdb::IteratorMode::From(
from,
if backwards {
rocksdb::Direction::Reverse
} else {
rocksdb::Direction::Forward
},
),
)
.map(|(k, v)| (Vec::from(k), Vec::from(v))),
)
}
fn increment(&self, key: &[u8]) -> Result<Vec<u8>> {
let lock = self.write_lock.write().unwrap();
let old = self.db.rocks.get_cf(&self.cf(), &key)?;
let new = utils::increment(old.as_deref()).unwrap();
self.db.rocks.put_cf(&self.cf(), key, &new)?;
drop(lock);
Ok(new)
}
fn increment_batch<'a>(&self, iter: &mut dyn Iterator<Item = Vec<u8>>) -> Result<()> {
let lock = self.write_lock.write().unwrap();
for key in iter {
let old = self.db.rocks.get_cf(&self.cf(), &key)?;
let new = utils::increment(old.as_deref()).unwrap();
self.db.rocks.put_cf(&self.cf(), key, new)?;
}
drop(lock);
Ok(())
}
fn scan_prefix<'a>(
&'a self,
prefix: Vec<u8>,
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
Box::new(
self.db
.rocks
.iterator_cf(
&self.cf(),
rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward),
)
.map(|(k, v)| (Vec::from(k), Vec::from(v)))
.take_while(move |(k, _)| k.starts_with(&prefix)),
)
}
fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
self.watchers.watch(prefix)
}
}

60
src/database/abstraction/sqlite.rs
View file

@ -1,17 +1,15 @@
use super::{DatabaseEngine, Tree};
use super::{watchers::Watchers, DatabaseEngine, Tree};
use crate::{database::Config, Result};
use parking_lot::{Mutex, MutexGuard, RwLock};
use parking_lot::{Mutex, MutexGuard};
use rusqlite::{Connection, DatabaseName::Main, OptionalExtension};
use std::{
cell::RefCell,
collections::{hash_map, HashMap},
future::Future,
path::{Path, PathBuf},
pin::Pin,
sync::Arc,
};
use thread_local::ThreadLocal;
use tokio::sync::watch;
use tracing::debug;
thread_local! {
@ -82,8 +80,8 @@ impl Engine {
}
}
impl DatabaseEngine for Engine {
fn open(config: &Config) -> Result<Arc<Self>> {
impl DatabaseEngine for Arc<Engine> {
fn open(config: &Config) -> Result<Self> {
let path = Path::new(&config.database_path).join("conduit.db");
// calculates cache-size per permanent connection
@ -94,7 +92,7 @@ impl DatabaseEngine for Engine {
/ ((num_cpus::get().max(1) * 2) + 1) as f64)
as u32;
let writer = Mutex::new(Self::prepare_conn(&path, cache_size_per_thread)?);
let writer = Mutex::new(Engine::prepare_conn(&path, cache_size_per_thread)?);
let arc = Arc::new(Engine {
writer,
@ -107,26 +105,30 @@ impl DatabaseEngine for Engine {
Ok(arc)
}
fn open_tree(self: &Arc<Self>, name: &str) -> Result<Arc<dyn Tree>> {
fn open_tree(&self, name: &str) -> Result<Arc<dyn Tree>> {
self.write_lock().execute(&format!("CREATE TABLE IF NOT EXISTS {} ( \"key\" BLOB PRIMARY KEY, \"value\" BLOB NOT NULL )", name), [])?;
Ok(Arc::new(SqliteTable {
engine: Arc::clone(self),
name: name.to_owned(),
watchers: RwLock::new(HashMap::new()),
watchers: Watchers::default(),
}))
}
fn flush(self: &Arc<Self>) -> Result<()> {
fn flush(&self) -> Result<()> {
// we enabled PRAGMA synchronous=normal, so this should not be necessary
Ok(())
}
fn cleanup(&self) -> Result<()> {
self.flush_wal()
}
}
pub struct SqliteTable {
engine: Arc<Engine>,
name: String,
watchers: RwLock<HashMap<Vec<u8>, (watch::Sender<()>, watch::Receiver<()>)>>,
watchers: Watchers,
}
type TupleOfBytes = (Vec<u8>, Vec<u8>);
@ -200,27 +202,7 @@ impl Tree for SqliteTable {
let guard = self.engine.write_lock();
self.insert_with_guard(&guard, key, value)?;
drop(guard);
let watchers = self.watchers.read();
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();
for prefix in triggered {
if let Some(tx) = watchers.remove(prefix) {
let _ = tx.0.send(());
}
}
};
self.watchers.wake(key);
Ok(())
}
@ -365,19 +347,7 @@ impl Tree for SqliteTable {
#[tracing::instrument(skip(self, prefix))]
fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
let mut rx = match self.watchers.write().entry(prefix.to_vec()) {
hash_map::Entry::Occupied(o) => o.get().1.clone(),
hash_map::Entry::Vacant(v) => {
let (tx, rx) = tokio::sync::watch::channel(());
v.insert((tx, rx.clone()));
rx
}
};
Box::pin(async move {
// Tx is never destroyed
rx.changed().await.unwrap();
})
self.watchers.watch(prefix)
}
#[tracing::instrument(skip(self))]

54
src/database/abstraction/watchers.rs Normal file
View file

@ -0,0 +1,54 @@
use std::{
collections::{hash_map, HashMap},
future::Future,
pin::Pin,
sync::RwLock,
};
use tokio::sync::watch;
#[derive(Default)]
pub(super) struct Watchers {
watchers: RwLock<HashMap<Vec<u8>, (watch::Sender<()>, watch::Receiver<()>)>>,
}
impl Watchers {
pub(super) fn watch<'a>(
&'a self,
prefix: &[u8],
) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
let mut rx = match self.watchers.write().unwrap().entry(prefix.to_vec()) {
hash_map::Entry::Occupied(o) => o.get().1.clone(),
hash_map::Entry::Vacant(v) => {
let (tx, rx) = tokio::sync::watch::channel(());
v.insert((tx, rx.clone()));
rx
}
};
Box::pin(async move {
// Tx is never destroyed
rx.changed().await.unwrap();
})
}
pub(super) fn wake(&self, key: &[u8]) {
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(tx) = watchers.remove(prefix) {
let _ = tx.0.send(());
}
}
};
}
}

12
src/database/admin.rs
View file

@ -12,8 +12,10 @@ use tracing::warn;
pub enum AdminCommand {
RegisterAppservice(serde_yaml::Value),
UnregisterAppservice(String),
ListAppservices,
ListLocalUsers,
ShowMemoryUsage,
SendMessage(RoomMessageEventContent),
}
@ -109,6 +111,9 @@ impl Admin {
AdminCommand::RegisterAppservice(yaml) => {
guard.appservice.register_appservice(yaml).unwrap(); // TODO handle error
}
AdminCommand::UnregisterAppservice(service_name) => {
guard.appservice.unregister_appservice(&service_name).unwrap(); // TODO: see above
}
AdminCommand::ListAppservices => {
if let Ok(appservices) = guard.appservice.iter_ids().map(|ids| ids.collect::<Vec<_>>()) {
let count = appservices.len();
@ -122,6 +127,13 @@ impl Admin {
send_message(RoomMessageEventContent::text_plain("Failed to get appservices."), guard, &state_lock);
}
}
AdminCommand::ShowMemoryUsage => {
if let Ok(response) = guard._db.memory_usage() {
send_message(RoomMessageEventContent::text_plain(response), guard, &state_lock);
} else {
send_message(RoomMessageEventContent::text_plain("Failed to get database memory usage.".to_owned()), guard, &state_lock);
}
}
AdminCommand::SendMessage(message) => {
send_message(message, guard, &state_lock);
}

15
src/database/appservice.rs
View file

@ -27,6 +27,21 @@ impl Appservice {
Ok(())
}
/// Remove an appservice registration
///
/// # Arguments
///
/// * `service_name` - the name you send to register the service previously
pub fn unregister_appservice(&self, service_name: &str) -> Result<()> {
self.id_appserviceregistrations
.remove(service_name.as_bytes())?;
self.cached_registrations
.write()
.unwrap()
.remove(service_name);
Ok(())
}
pub fn get_registration(&self, id: &str) -> Result<Option<serde_yaml::Value>> {
self.cached_registrations
.read()

65
src/database/key_backups.rs
View file

@ -4,8 +4,10 @@ use ruma::{
error::ErrorKind,
r0::backup::{BackupAlgorithm, KeyBackupData, RoomKeyBackup},
},
serde::Raw,
RoomId, UserId,
};
use serde_json::json;
use std::{collections::BTreeMap, sync::Arc};
use super::abstraction::Tree;
@ -20,7 +22,7 @@ impl KeyBackups {
pub fn create_backup(
&self,
user_id: &UserId,
backup_metadata: &BackupAlgorithm,
backup_metadata: &Raw<BackupAlgorithm>,
globals: &super::globals::Globals,
) -> Result<String> {
let version = globals.next_count()?.to_string();
@ -59,7 +61,7 @@ impl KeyBackups {
&self,
user_id: &UserId,
version: &str,
backup_metadata: &BackupAlgorithm,
backup_metadata: &Raw<BackupAlgorithm>,
globals: &super::globals::Globals,
) -> Result<String> {
let mut key = user_id.as_bytes().to_vec();
@ -73,12 +75,8 @@ impl KeyBackups {
));
}
self.backupid_algorithm.insert(
&key,
serde_json::to_string(backup_metadata)
.expect("BackupAlgorithm::to_string always works")
.as_bytes(),
)?;
self.backupid_algorithm
.insert(&key, backup_metadata.json().get().as_bytes())?;
self.backupid_etag
.insert(&key, &globals.next_count()?.to_be_bytes())?;
Ok(version.to_owned())
@ -105,7 +103,10 @@ impl KeyBackups {
.transpose()
}
pub fn get_latest_backup(&self, user_id: &UserId) -> Result<Option<(String, BackupAlgorithm)>> {
pub fn get_latest_backup(
&self,
user_id: &UserId,
) -> Result<Option<(String, Raw<BackupAlgorithm>)>> {
let mut prefix = user_id.as_bytes().to_vec();
prefix.push(0xff);
let mut last_possible_key = prefix.clone();
@ -133,7 +134,11 @@ impl KeyBackups {
.transpose()
}
pub fn get_backup(&self, user_id: &UserId, version: &str) -> Result<Option<BackupAlgorithm>> {
pub fn get_backup(
&self,
user_id: &UserId,
version: &str,
) -> Result<Option<Raw<BackupAlgorithm>>> {
let mut key = user_id.as_bytes().to_vec();
key.push(0xff);
key.extend_from_slice(version.as_bytes());
@ -152,7 +157,7 @@ impl KeyBackups {
version: &str,
room_id: &RoomId,
session_id: &str,
key_data: &KeyBackupData,
key_data: &Raw<KeyBackupData>,
globals: &super::globals::Globals,
) -> Result<()> {
let mut key = user_id.as_bytes().to_vec();
@ -174,10 +179,8 @@ impl KeyBackups {
key.push(0xff);
key.extend_from_slice(session_id.as_bytes());
self.backupkeyid_backup.insert(
&key,
&serde_json::to_vec(&key_data).expect("KeyBackupData::to_vec always works"),
)?;
self.backupkeyid_backup
.insert(&key, key_data.json().get().as_bytes())?;
Ok(())
}
@ -209,13 +212,13 @@ impl KeyBackups {
&self,
user_id: &UserId,
version: &str,
) -> Result<BTreeMap<Box<RoomId>, RoomKeyBackup>> {
) -> Result<BTreeMap<Box<RoomId>, Raw<RoomKeyBackup>>> {
let mut prefix = user_id.as_bytes().to_vec();
prefix.push(0xff);
prefix.extend_from_slice(version.as_bytes());
prefix.push(0xff);
let mut rooms = BTreeMap::<Box<RoomId>, RoomKeyBackup>::new();
let mut rooms = BTreeMap::<Box<RoomId>, Raw<RoomKeyBackup>>::new();
for result in self
.backupkeyid_backup
@ -241,7 +244,7 @@ impl KeyBackups {
Error::bad_database("backupkeyid_backup room_id is invalid room id.")
})?;
let key_data = serde_json::from_slice(&value).map_err(|_| {
let key_data: serde_json::Value = serde_json::from_slice(&value).map_err(|_| {
Error::bad_database("KeyBackupData in backupkeyid_backup is invalid.")
})?;
@ -249,13 +252,25 @@ impl KeyBackups {
})
{
let (room_id, session_id, key_data) = result?;
rooms
.entry(room_id)
.or_insert_with(|| RoomKeyBackup {
let room_key_backup = rooms.entry(room_id).or_insert_with(|| {
Raw::new(&RoomKeyBackup {
sessions: BTreeMap::new(),
})
.sessions
.insert(session_id, key_data);
.expect("RoomKeyBackup serialization")
});
let mut object = room_key_backup
.deserialize_as::<serde_json::Map<String, serde_json::Value>>()
.map_err(|_| Error::bad_database("RoomKeyBackup is not an object"))?;
let sessions = object.entry("session").or_insert_with(|| json!({}));
if let serde_json::Value::Object(unsigned_object) = sessions {
unsigned_object.insert(session_id, key_data);
}
*room_key_backup = Raw::from_json(
serde_json::value::to_raw_value(&object).expect("Value => RawValue serialization"),
);
}
Ok(rooms)
@ -266,7 +281,7 @@ impl KeyBackups {
user_id: &UserId,
version: &str,
room_id: &RoomId,
) -> Result<BTreeMap<String, KeyBackupData>> {
) -> Result<BTreeMap<String, Raw<KeyBackupData>>> {
let mut prefix = user_id.as_bytes().to_vec();
prefix.push(0xff);
prefix.extend_from_slice(version.as_bytes());
@ -304,7 +319,7 @@ impl KeyBackups {
version: &str,
room_id: &RoomId,
session_id: &str,
) -> Result<Option<KeyBackupData>> {
) -> Result<Option<Raw<KeyBackupData>>> {
let mut key = user_id.as_bytes().to_vec();
key.push(0xff);
key.extend_from_slice(version.as_bytes());

117
src/database/rooms.rs
View file

@ -28,7 +28,7 @@ use ruma::{
push::{Action, Ruleset, Tweak},
serde::{CanonicalJsonObject, CanonicalJsonValue, Raw},
state_res::{self, RoomVersion, StateMap},
uint, EventId, RoomAliasId, RoomId, RoomVersionId, ServerName, UserId,
uint, DeviceId, EventId, RoomAliasId, RoomId, RoomVersionId, ServerName, UserId,
};
use serde::Deserialize;
use serde_json::value::to_raw_value;
@ -79,6 +79,8 @@ pub struct Rooms {
pub(super) userroomid_leftstate: Arc<dyn Tree>,
pub(super) roomuserid_leftcount: Arc<dyn Tree>,
pub(super) lazyloadedids: Arc<dyn Tree>, // LazyLoadedIds = UserId + DeviceId + RoomId + LazyLoadedUserId
pub(super) userroomid_notificationcount: Arc<dyn Tree>, // NotifyCount = u64
pub(super) userroomid_highlightcount: Arc<dyn Tree>, // HightlightCount = u64
@ -117,6 +119,8 @@ pub struct Rooms {
pub(super) shortstatekey_cache: Mutex<LruCache<u64, (EventType, String)>>,
pub(super) our_real_users_cache: RwLock<HashMap<Box<RoomId>, Arc<HashSet<Box<UserId>>>>>,
pub(super) appservice_in_room_cache: RwLock<HashMap<Box<RoomId>, HashMap<String, bool>>>,
pub(super) lazy_load_waiting:
Mutex<HashMap<(Box<UserId>, Box<DeviceId>, Box<RoomId>, u64), HashSet<Box<UserId>>>>,
pub(super) stateinfo_cache: Mutex<
LruCache<
u64,
@ -1528,6 +1532,19 @@ impl Rooms {
));
}
}
"unregister_appservice" => {
if args.len() == 1 {
db.admin.send(AdminCommand::UnregisterAppservice(
args[0].to_owned(),
));
} else {
db.admin.send(AdminCommand::SendMessage(
RoomMessageEventContent::text_plain(
"Missing appservice identifier",
),
));
}
}
"list_appservices" => {
db.admin.send(AdminCommand::ListAppservices);
}
@ -1679,6 +1696,9 @@ impl Rooms {
));
}
}
"database_memory_usage" => {
db.admin.send(AdminCommand::ShowMemoryUsage);
}
_ => {
db.admin.send(AdminCommand::SendMessage(
RoomMessageEventContent::text_plain(format!(
@ -2710,7 +2730,7 @@ impl Rooms {
let state_lock = mutex_state.lock().await;
let mut event: RoomMemberEventContent = serde_json::from_str(
self.room_state_get(room_id, &EventType::RoomMember, &user_id.to_string())?
self.room_state_get(room_id, &EventType::RoomMember, user_id.as_str())?
.ok_or(Error::BadRequest(
ErrorKind::BadState,
"Cannot leave a room you are not a member of.",
@ -3445,8 +3465,7 @@ impl Rooms {
&key[0].to_be_bytes(),
&chain
.iter()
.map(|s| s.to_be_bytes().to_vec())
.flatten()
.flat_map(|s| s.to_be_bytes().to_vec())
.collect::<Vec<u8>>(),
)?;
}
@ -3456,4 +3475,94 @@ impl Rooms {
Ok(())
}
#[tracing::instrument(skip(self))]
pub fn lazy_load_was_sent_before(
&self,
user_id: &UserId,
device_id: &DeviceId,
room_id: &RoomId,
ll_user: &UserId,
) -> Result<bool> {
let mut key = user_id.as_bytes().to_vec();
key.push(0xff);
key.extend_from_slice(device_id.as_bytes());
key.push(0xff);
key.extend_from_slice(room_id.as_bytes());
key.push(0xff);
key.extend_from_slice(ll_user.as_bytes());
Ok(self.lazyloadedids.get(&key)?.is_some())
}
#[tracing::instrument(skip(self))]
pub fn lazy_load_mark_sent(
&self,
user_id: &UserId,
device_id: &DeviceId,
room_id: &RoomId,
lazy_load: HashSet<Box<UserId>>,
count: u64,
) {
self.lazy_load_waiting.lock().unwrap().insert(
(
user_id.to_owned(),
device_id.to_owned(),
room_id.to_owned(),
count,
),
lazy_load,
);
}
#[tracing::instrument(skip(self))]
pub fn lazy_load_confirm_delivery(
&self,
user_id: &UserId,
device_id: &DeviceId,
room_id: &RoomId,
since: u64,
) -> Result<()> {
if let Some(user_ids) = self.lazy_load_waiting.lock().unwrap().remove(&(
user_id.to_owned(),
device_id.to_owned(),
room_id.to_owned(),
since,
)) {
let mut prefix = user_id.as_bytes().to_vec();
prefix.push(0xff);
prefix.extend_from_slice(device_id.as_bytes());
prefix.push(0xff);
prefix.extend_from_slice(room_id.as_bytes());
prefix.push(0xff);
for ll_id in user_ids {
let mut key = prefix.clone();
key.extend_from_slice(ll_id.as_bytes());
self.lazyloadedids.insert(&key, &[])?;
}
}
Ok(())
}
#[tracing::instrument(skip(self))]
pub fn lazy_load_reset(
&self,
user_id: &UserId,
device_id: &DeviceId,
room_id: &RoomId,
) -> Result<()> {
let mut prefix = user_id.as_bytes().to_vec();
prefix.push(0xff);
prefix.extend_from_slice(device_id.as_bytes());
prefix.push(0xff);
prefix.extend_from_slice(room_id.as_bytes());
prefix.push(0xff);
for (key, _) in self.lazyloadedids.scan_prefix(prefix) {
self.lazyloadedids.remove(&key)?;
}
Ok(())
}
}

10
src/database/sending.rs
View file

@ -524,7 +524,7 @@ impl Sending {
.unwrap(), // TODO: handle error
appservice::event::push_events::v1::Request {
events: &pdu_jsons,
txn_id: &base64::encode_config(
txn_id: (&*base64::encode_config(
Self::calculate_hash(
&events
.iter()
@ -534,7 +534,8 @@ impl Sending {
.collect::<Vec<_>>(),
),
base64::URL_SAFE_NO_PAD,
),
))
.into(),
},
)
.await
@ -682,7 +683,7 @@ impl Sending {
pdus: &pdu_jsons,
edus: &edu_jsons,
origin_server_ts: MilliSecondsSinceUnixEpoch::now(),
transaction_id: &base64::encode_config(
transaction_id: (&*base64::encode_config(
Self::calculate_hash(
&events
.iter()
@ -692,7 +693,8 @@ impl Sending {
.collect::<Vec<_>>(),
),
base64::URL_SAFE_NO_PAD,
),
))
.into(),
},
)
.await

6
src/database/transaction_ids.rs
View file

@ -1,7 +1,7 @@
use std::sync::Arc;
use crate::Result;
use ruma::{DeviceId, UserId};
use ruma::{identifiers::TransactionId, DeviceId, UserId};
use super::abstraction::Tree;
@ -14,7 +14,7 @@ impl TransactionIds {
&self,
user_id: &UserId,
device_id: Option<&DeviceId>,
txn_id: &str,
txn_id: &TransactionId,
data: &[u8],
) -> Result<()> {
let mut key = user_id.as_bytes().to_vec();
@ -32,7 +32,7 @@ impl TransactionIds {
&self,
user_id: &UserId,
device_id: Option<&DeviceId>,
txn_id: &str,
txn_id: &TransactionId,
) -> Result<Option<Vec<u8>>> {
let mut key = user_id.as_bytes().to_vec();
key.push(0xff);

43
src/database/uiaa.rs
View file

@ -1,4 +1,6 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::RwLock;
use crate::{client_server::SESSION_ID_LENGTH, utils, Error, Result};
use ruma::{
@ -18,7 +20,8 @@ use super::abstraction::Tree;
pub struct Uiaa {
pub(super) userdevicesessionid_uiaainfo: Arc<dyn Tree>, // User-interactive authentication
pub(super) userdevicesessionid_uiaarequest: Arc<dyn Tree>, // UiaaRequest = canonical json value
pub(super) userdevicesessionid_uiaarequest:
RwLock<BTreeMap<(Box<UserId>, Box<DeviceId>, String), CanonicalJsonValue>>,
}
impl Uiaa {
@ -147,16 +150,13 @@ impl Uiaa {
session: &str,
request: &CanonicalJsonValue,
) -> Result<()> {
let mut userdevicesessionid = user_id.as_bytes().to_vec();
userdevicesessionid.push(0xff);
userdevicesessionid.extend_from_slice(device_id.as_bytes());
userdevicesessionid.push(0xff);
userdevicesessionid.extend_from_slice(session.as_bytes());
self.userdevicesessionid_uiaarequest.insert(
&userdevicesessionid,
&serde_json::to_vec(request).expect("json value to vec always works"),
)?;
self.userdevicesessionid_uiaarequest
.write()
.unwrap()
.insert(
(user_id.to_owned(), device_id.to_owned(), session.to_owned()),
request.to_owned(),
);
Ok(())
}
@ -166,23 +166,12 @@ impl Uiaa {
user_id: &UserId,
device_id: &DeviceId,
session: &str,
) -> Result<Option<CanonicalJsonValue>> {
let mut userdevicesessionid = user_id.as_bytes().to_vec();
userdevicesessionid.push(0xff);
userdevicesessionid.extend_from_slice(device_id.as_bytes());
userdevicesessionid.push(0xff);
userdevicesessionid.extend_from_slice(session.as_bytes());
) -> Option<CanonicalJsonValue> {
self.userdevicesessionid_uiaarequest
.get(&userdevicesessionid)?
.map(|bytes| {
serde_json::from_str::<CanonicalJsonValue>(
&utils::string_from_bytes(&bytes)
.map_err(|_| Error::bad_database("Invalid uiaa request bytes in db."))?,
)
.map_err(|_| Error::bad_database("Invalid uiaa request in db."))
})
.transpose()
.read()
.unwrap()
.get(&(user_id.to_owned(), device_id.to_owned(), session.to_owned()))
.map(|j| j.to_owned())
}
fn update_uiaa_session(

179
src/database/users.rs
View file

@ -1,6 +1,9 @@
use crate::{utils, Error, Result};
use ruma::{
api::client::{error::ErrorKind, r0::device::Device},
api::client::{
error::ErrorKind,
r0::{device::Device, filter::IncomingFilterDefinition},
},
encryption::{CrossSigningKey, DeviceKeys, OneTimeKey},
events::{AnyToDeviceEvent, EventType},
identifiers::MxcUri,
@ -8,7 +11,12 @@ use ruma::{
DeviceId, DeviceKeyAlgorithm, DeviceKeyId, MilliSecondsSinceUnixEpoch, RoomAliasId, UInt,
UserId,
};
use std::{collections::BTreeMap, convert::TryInto, mem, sync::Arc};
use std::{
collections::BTreeMap,
convert::{TryFrom, TryInto},
mem,
sync::Arc,
};
use tracing::warn;
use super::abstraction::Tree;
@ -31,6 +39,8 @@ pub struct Users {
pub(super) userid_selfsigningkeyid: Arc<dyn Tree>,
pub(super) userid_usersigningkeyid: Arc<dyn Tree>,
pub(super) userfilterid_filter: Arc<dyn Tree>, // UserFilterId = UserId + FilterId
pub(super) todeviceid_events: Arc<dyn Tree>, // ToDeviceId = UserId + DeviceId + Count
}
@ -395,7 +405,7 @@ impl Users {
user_id: &UserId,
device_id: &DeviceId,
one_time_key_key: &DeviceKeyId,
one_time_key_value: &OneTimeKey,
one_time_key_value: &Raw<OneTimeKey>,
globals: &super::globals::Globals,
) -> Result<()> {
let mut key = user_id.as_bytes().to_vec();
@ -445,7 +455,7 @@ impl Users {
device_id: &DeviceId,
key_algorithm: &DeviceKeyAlgorithm,
globals: &super::globals::Globals,
) -> Result<Option<(Box<DeviceKeyId>, OneTimeKey)>> {
) -> Result<Option<(Box<DeviceKeyId>, Raw<OneTimeKey>)>> {
let mut prefix = user_id.as_bytes().to_vec();
prefix.push(0xff);
prefix.extend_from_slice(device_id.as_bytes());
@ -516,7 +526,7 @@ impl Users {
&self,
user_id: &UserId,
device_id: &DeviceId,
device_keys: &DeviceKeys,
device_keys: &Raw<DeviceKeys>,
rooms: &super::rooms::Rooms,
globals: &super::globals::Globals,
) -> Result<()> {
@ -545,9 +555,9 @@ impl Users {
pub fn add_cross_signing_keys(
&self,
user_id: &UserId,
master_key: &CrossSigningKey,
self_signing_key: &Option<CrossSigningKey>,
user_signing_key: &Option<CrossSigningKey>,
master_key: &Raw<CrossSigningKey>,
self_signing_key: &Option<Raw<CrossSigningKey>>,
user_signing_key: &Option<Raw<CrossSigningKey>>,
rooms: &super::rooms::Rooms,
globals: &super::globals::Globals,
) -> Result<()> {
@ -557,7 +567,12 @@ impl Users {
prefix.push(0xff);
// Master key
let mut master_key_ids = master_key.keys.values();
let mut master_key_ids = master_key
.deserialize()
.map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Invalid master key"))?
.keys
.into_values();
let master_key_id = master_key_ids.next().ok_or(Error::BadRequest(
ErrorKind::InvalidParam,
"Master key contained no key.",
@ -573,17 +588,22 @@ impl Users {
let mut master_key_key = prefix.clone();
master_key_key.extend_from_slice(master_key_id.as_bytes());
self.keyid_key.insert(
&master_key_key,
&serde_json::to_vec(&master_key).expect("CrossSigningKey::to_vec always works"),
)?;
self.keyid_key
.insert(&master_key_key, master_key.json().get().as_bytes())?;
self.userid_masterkeyid
.insert(user_id.as_bytes(), &master_key_key)?;
// Self-signing key
if let Some(self_signing_key) = self_signing_key {
let mut self_signing_key_ids = self_signing_key.keys.values();
let mut self_signing_key_ids = self_signing_key
.deserialize()
.map_err(|_| {
Error::BadRequest(ErrorKind::InvalidParam, "Invalid self signing key")
})?
.keys
.into_values();
let self_signing_key_id = self_signing_key_ids.next().ok_or(Error::BadRequest(
ErrorKind::InvalidParam,
"Self signing key contained no key.",
@ -601,8 +621,7 @@ impl Users {
self.keyid_key.insert(
&self_signing_key_key,
&serde_json::to_vec(&self_signing_key)
.expect("CrossSigningKey::to_vec always works"),
self_signing_key.json().get().as_bytes(),
)?;
self.userid_selfsigningkeyid
@ -611,7 +630,14 @@ impl Users {
// User-signing key
if let Some(user_signing_key) = user_signing_key {
let mut user_signing_key_ids = user_signing_key.keys.values();
let mut user_signing_key_ids = user_signing_key
.deserialize()
.map_err(|_| {
Error::BadRequest(ErrorKind::InvalidParam, "Invalid user signing key")
})?
.keys
.into_values();
let user_signing_key_id = user_signing_key_ids.next().ok_or(Error::BadRequest(
ErrorKind::InvalidParam,
"User signing key contained no key.",
@ -629,8 +655,7 @@ impl Users {
self.keyid_key.insert(
&user_signing_key_key,
&serde_json::to_vec(&user_signing_key)
.expect("CrossSigningKey::to_vec always works"),
user_signing_key.json().get().as_bytes(),
)?;
self.userid_usersigningkeyid
@ -763,7 +788,7 @@ impl Users {
&self,
user_id: &UserId,
device_id: &DeviceId,
) -> Result<Option<DeviceKeys>> {
) -> Result<Option<Raw<DeviceKeys>>> {
let mut key = user_id.as_bytes().to_vec();
key.push(0xff);
key.extend_from_slice(device_id.as_bytes());
@ -780,25 +805,19 @@ impl Users {
&self,
user_id: &UserId,
allowed_signatures: F,
) -> Result<Option<CrossSigningKey>> {
) -> Result<Option<Raw<CrossSigningKey>>> {
self.userid_masterkeyid
.get(user_id.as_bytes())?
.map_or(Ok(None), |key| {
self.keyid_key.get(&key)?.map_or(Ok(None), |bytes| {
let mut cross_signing_key = serde_json::from_slice::<CrossSigningKey>(&bytes)
.map_err(|_| {
Error::bad_database("CrossSigningKey in db is invalid.")
})?;
let mut cross_signing_key = serde_json::from_slice::<serde_json::Value>(&bytes)
.map_err(|_| Error::bad_database("CrossSigningKey in db is invalid."))?;
clean_signatures(&mut cross_signing_key, user_id, allowed_signatures)?;
// A user is not allowed to see signatures from users other than himself and
// the target user
cross_signing_key.signatures = cross_signing_key
.signatures
.into_iter()
.filter(|(user, _)| allowed_signatures(user))
.collect();
Ok(Some(cross_signing_key))
Ok(Some(Raw::from_json(
serde_json::value::to_raw_value(&cross_signing_key)
.expect("Value to RawValue serialization"),
)))
})
})
}
@ -808,31 +827,25 @@ impl Users {
&self,
user_id: &UserId,
allowed_signatures: F,
) -> Result<Option<CrossSigningKey>> {
) -> Result<Option<Raw<CrossSigningKey>>> {
self.userid_selfsigningkeyid
.get(user_id.as_bytes())?
.map_or(Ok(None), |key| {
self.keyid_key.get(&key)?.map_or(Ok(None), |bytes| {
let mut cross_signing_key = serde_json::from_slice::<CrossSigningKey>(&bytes)
.map_err(|_| {
Error::bad_database("CrossSigningKey in db is invalid.")
})?;
let mut cross_signing_key = serde_json::from_slice::<serde_json::Value>(&bytes)
.map_err(|_| Error::bad_database("CrossSigningKey in db is invalid."))?;
clean_signatures(&mut cross_signing_key, user_id, allowed_signatures)?;
// A user is not allowed to see signatures from users other than himself and
// the target user
cross_signing_key.signatures = cross_signing_key
.signatures
.into_iter()
.filter(|(user, _)| user == user_id || allowed_signatures(user))
.collect();
Ok(Some(cross_signing_key))
Ok(Some(Raw::from_json(
serde_json::value::to_raw_value(&cross_signing_key)
.expect("Value to RawValue serialization"),
)))
})
})
}
#[tracing::instrument(skip(self, user_id))]
pub fn get_user_signing_key(&self, user_id: &UserId) -> Result<Option<CrossSigningKey>> {
pub fn get_user_signing_key(&self, user_id: &UserId) -> Result<Option<Raw<CrossSigningKey>>> {
self.userid_usersigningkeyid
.get(user_id.as_bytes())?
.map_or(Ok(None), |key| {
@ -1026,4 +1039,72 @@ impl Users {
// TODO: Unhook 3PID
Ok(())
}
/// Creates a new sync filter. Returns the filter id.
#[tracing::instrument(skip(self))]
pub fn create_filter(
&self,
user_id: &UserId,
filter: &IncomingFilterDefinition,
) -> Result<String> {
let filter_id = utils::random_string(4);
let mut key = user_id.as_bytes().to_vec();
key.push(0xff);
key.extend_from_slice(filter_id.as_bytes());
self.userfilterid_filter.insert(
&key,
&serde_json::to_vec(&filter).expect("filter is valid json"),
)?;
Ok(filter_id)
}
#[tracing::instrument(skip(self))]
pub fn get_filter(
&self,
user_id: &UserId,
filter_id: &str,
) -> Result<Option<IncomingFilterDefinition>> {
let mut key = user_id.as_bytes().to_vec();
key.push(0xff);
key.extend_from_slice(filter_id.as_bytes());
let raw = self.userfilterid_filter.get(&key)?;
if let Some(raw) = raw {
serde_json::from_slice(&raw)
.map_err(|_| Error::bad_database("Invalid filter event in db."))
} else {
Ok(None)
}
}
}
/// Ensure that a user only sees signatures from themselves and the target user
fn clean_signatures<F: Fn(&UserId) -> bool>(
cross_signing_key: &mut serde_json::Value,
user_id: &UserId,
allowed_signatures: F,
) -> Result<(), Error> {
if let Some(signatures) = cross_signing_key
.get_mut("signatures")
.and_then(|v| v.as_object_mut())
{
// Don't allocate for the full size of the current signatures, but require
// at most one resize if nothing is dropped
let new_capacity = signatures.len() / 2;
for (user, signature) in
mem::replace(signatures, serde_json::Map::with_capacity(new_capacity))
{
let id = <&UserId>::try_from(user.as_str())
.map_err(|_| Error::bad_database("Invalid user ID in database."))?;
if id == user_id || allowed_signatures(id) {
signatures.insert(user, signature);
}
}
}
Ok(())
}