Merge remote-tracking branch 'refs/remotes/origin/next' into command-refactor

Resolved conflict for the new list_local_users command
This commit is contained in:
Andrei Vasiliu 2022-01-22 14:29:50 +02:00
commit 7505548b94
10 changed files with 472 additions and 178 deletions

View file

@ -15,7 +15,15 @@ pub mod heed;
#[cfg(feature = "rocksdb")]
pub mod rocksdb;
#[cfg(any(feature = "sqlite", feature = "rocksdb", feature = "heed"))]
#[cfg(feature = "persy")]
pub mod persy;
#[cfg(any(
feature = "sqlite",
feature = "rocksdb",
feature = "heed",
feature = "persy"
))]
pub mod watchers;
pub trait DatabaseEngine: Send + Sync {

View file

@ -0,0 +1,207 @@
use crate::{
database::{
abstraction::{watchers::Watchers, DatabaseEngine, Tree},
Config,
},
Result,
};
use persy::{ByteVec, OpenOptions, Persy, Transaction, TransactionConfig, ValueMode};
use std::{future::Future, pin::Pin, sync::Arc};
use tracing::warn;
pub struct Engine {
persy: Persy,
}
impl DatabaseEngine for Arc<Engine> {
fn open(config: &Config) -> Result<Self> {
let mut cfg = persy::Config::new();
cfg.change_cache_size((config.db_cache_capacity_mb * 1024.0 * 1024.0) as u64);
let persy = OpenOptions::new()
.create(true)
.config(cfg)
.open(&format!("{}/db.persy", config.database_path))?;
Ok(Arc::new(Engine { persy }))
}
fn open_tree(&self, name: &'static str) -> Result<Arc<dyn Tree>> {
// Create if it doesn't exist
if !self.persy.exists_index(name)? {
let mut tx = self.persy.begin()?;
tx.create_index::<ByteVec, ByteVec>(name, ValueMode::Replace)?;
tx.prepare()?.commit()?;
}
Ok(Arc::new(PersyTree {
persy: self.persy.clone(),
name: name.to_owned(),
watchers: Watchers::default(),
}))
}
fn flush(&self) -> Result<()> {
Ok(())
}
}
pub struct PersyTree {
persy: Persy,
name: String,
watchers: Watchers,
}
impl PersyTree {
fn begin(&self) -> Result<Transaction> {
Ok(self
.persy
.begin_with(TransactionConfig::new().set_background_sync(true))?)
}
}
impl Tree for PersyTree {
#[tracing::instrument(skip(self, key))]
fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
let result = self
.persy
.get::<ByteVec, ByteVec>(&self.name, &ByteVec::from(key))?
.next()
.map(|v| (*v).to_owned());
Ok(result)
}
#[tracing::instrument(skip(self, key, value))]
fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> {
self.insert_batch(&mut Some((key.to_owned(), value.to_owned())).into_iter())?;
self.watchers.wake(key);
Ok(())
}
#[tracing::instrument(skip(self, iter))]
fn insert_batch<'a>(&self, iter: &mut dyn Iterator<Item = (Vec<u8>, Vec<u8>)>) -> Result<()> {
let mut tx = self.begin()?;
for (key, value) in iter {
tx.put::<ByteVec, ByteVec>(
&self.name,
ByteVec::from(key.clone()),
ByteVec::from(value),
)?;
}
tx.prepare()?.commit()?;
Ok(())
}
#[tracing::instrument(skip(self, iter))]
fn increment_batch<'a>(&self, iter: &mut dyn Iterator<Item = Vec<u8>>) -> Result<()> {
let mut tx = self.begin()?;
for key in iter {
let old = tx
.get::<ByteVec, ByteVec>(&self.name, &ByteVec::from(key.clone()))?
.next()
.map(|v| (*v).to_owned());
let new = crate::utils::increment(old.as_deref()).unwrap();
tx.put::<ByteVec, ByteVec>(&self.name, ByteVec::from(key), ByteVec::from(new))?;
}
tx.prepare()?.commit()?;
Ok(())
}
#[tracing::instrument(skip(self, key))]
fn remove(&self, key: &[u8]) -> Result<()> {
let mut tx = self.begin()?;
tx.remove::<ByteVec, ByteVec>(&self.name, ByteVec::from(key), None)?;
tx.prepare()?.commit()?;
Ok(())
}
#[tracing::instrument(skip(self))]
fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
let iter = self.persy.range::<ByteVec, ByteVec, _>(&self.name, ..);
match iter {
Ok(iter) => Box::new(iter.filter_map(|(k, v)| {
v.into_iter()
.map(|val| ((*k).to_owned().into(), (*val).to_owned().into()))
.next()
})),
Err(e) => {
warn!("error iterating {:?}", e);
Box::new(std::iter::empty())
}
}
}
#[tracing::instrument(skip(self, from, backwards))]
fn iter_from<'a>(
&'a self,
from: &[u8],
backwards: bool,
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
let range = if backwards {
self.persy
.range::<ByteVec, ByteVec, _>(&self.name, ..=ByteVec::from(from))
} else {
self.persy
.range::<ByteVec, ByteVec, _>(&self.name, ByteVec::from(from)..)
};
match range {
Ok(iter) => {
let map = iter.filter_map(|(k, v)| {
v.into_iter()
.map(|val| ((*k).to_owned().into(), (*val).to_owned().into()))
.next()
});
if backwards {
Box::new(map.rev())
} else {
Box::new(map)
}
}
Err(e) => {
warn!("error iterating with prefix {:?}", e);
Box::new(std::iter::empty())
}
}
}
#[tracing::instrument(skip(self, key))]
fn increment(&self, key: &[u8]) -> Result<Vec<u8>> {
self.increment_batch(&mut Some(key.to_owned()).into_iter())?;
Ok(self.get(key)?.unwrap())
}
#[tracing::instrument(skip(self, prefix))]
fn scan_prefix<'a>(
&'a self,
prefix: Vec<u8>,
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
let range_prefix = ByteVec::from(prefix.clone());
let range = self
.persy
.range::<ByteVec, ByteVec, _>(&self.name, range_prefix..);
match range {
Ok(iter) => {
let owned_prefix = prefix.clone();
Box::new(
iter.take_while(move |(k, _)| (*k).starts_with(&owned_prefix))
.filter_map(|(k, v)| {
v.into_iter()
.map(|val| ((*k).to_owned().into(), (*val).to_owned().into()))
.next()
}),
)
}
Err(e) => {
warn!("error scanning prefix {:?}", e);
Box::new(std::iter::empty())
}
}
}
#[tracing::instrument(skip(self, prefix))]
fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
self.watchers.watch(prefix)
}
}

View file

@ -23,6 +23,7 @@ pub enum AdminCommand {
RegisterAppservice(serde_yaml::Value),
UnregisterAppservice(String),
ListAppservices,
ListLocalUsers,
ShowMemoryUsage,
SendMessage(RoomMessageEventContent),
}
@ -104,6 +105,18 @@ impl Admin {
let state_lock = mutex_state.lock().await;
match event {
AdminCommand::ListLocalUsers => {
match guard.users.list_local_users() {
Ok(users) => {
let mut msg: String = format!("Found {} local user account(s):\n", users.len());
msg += &users.join("\n");
send_message(RoomMessageEventContent::text_plain(&msg), guard, &state_lock);
}
Err(e) => {
send_message(RoomMessageEventContent::text_plain(e.to_string()), guard, &state_lock);
}
}
}
AdminCommand::RegisterAppservice(yaml) => {
guard.appservice.register_appservice(yaml).unwrap(); // TODO handle error
}
@ -226,6 +239,9 @@ enum AdminCommands {
/// List all the currently registered appservices
ListAppservices,
/// List users in the database
ListLocalUsers,
/// Get the auth_chain of a PDU
GetAuthChain {
/// An event ID (the $ character followed by the base64 reference hash)
@ -289,6 +305,7 @@ pub fn try_parse_admin_command(
appservice_identifier,
} => AdminCommand::UnregisterAppservice(appservice_identifier),
AdminCommands::ListAppservices => AdminCommand::ListAppservices,
AdminCommands::ListLocalUsers => AdminCommand::ListLocalUsers,
AdminCommands::GetAuthChain { event_id } => {
let event_id = Arc::<EventId>::from(event_id);
if let Some(event) = db.rooms.get_pdu_json(&event_id)? {

View file

@ -131,6 +131,42 @@ impl Users {
})
}
/// Returns a list of local users as list of usernames.
///
/// A user account is considered `local` if the length of it's password is greater then zero.
#[tracing::instrument(skip(self))]
pub fn list_local_users(&self) -> Result<Vec<String>> {
let users: Vec<String> = self
.userid_password
.iter()
.filter_map(|(username, pw)| self.get_username_with_valid_password(&username, &pw))
.collect();
Ok(users)
}
/// Will only return with Some(username) if the password was not empty and the
/// username could be successfully parsed.
/// If utils::string_from_bytes(...) returns an error that username will be skipped
/// and the error will be logged.
#[tracing::instrument(skip(self))]
fn get_username_with_valid_password(&self, username: &[u8], password: &[u8]) -> Option<String> {
// A valid password is not empty
if password.is_empty() {
None
} else {
match utils::string_from_bytes(username) {
Ok(u) => Some(u),
Err(e) => {
warn!(
"Failed to parse username while calling get_local_users(): {}",
e.to_string()
);
None
}
}
}
}
/// Returns the password hash for the given user.
#[tracing::instrument(skip(self, user_id))]
pub fn password_hash(&self, user_id: &UserId) -> Result<Option<String>> {