From 1d2ba9eddc75a8fadecc4961896e2ff33ff32b43 Mon Sep 17 00:00:00 2001 From: Jonathan Krebs Date: Fri, 24 Jan 2025 04:18:14 +0100 Subject: [PATCH] implement dumping and loading the full database --- bffhd/lib.rs | 27 +++++++++++++++++- bffhd/resources/state/db.rs | 26 +++++++++++++++++ bffhd/users/mod.rs | 25 +++++++++++++++- bin/bffhd/main.rs | 57 ++++++++++++++++++++++++++++--------- 4 files changed, 120 insertions(+), 15 deletions(-) diff --git a/bffhd/lib.rs b/bffhd/lib.rs index 6ce37b1..6461f26 100644 --- a/bffhd/lib.rs +++ b/bffhd/lib.rs @@ -8,7 +8,7 @@ //! This is the capnp component of the FabAccess project. //! The entry point of bffhd can be found in [bin/bffhd/main.rs](../bin/bffhd/main.rs) -use miette::Diagnostic; +use miette::{Diagnostic, IntoDiagnostic}; use thiserror::Error; pub mod config; @@ -67,6 +67,8 @@ use lightproc::recoverable_handle::RecoverableHandle; use signal_hook::consts::signal::*; use tracing::Span; +use std::collections::HashMap; + pub struct Difluoroborane { config: Config, executor: Executor<'static>, @@ -135,6 +137,12 @@ pub enum BFFHError { ), } +#[derive(serde::Serialize, serde::Deserialize)] +struct DatabaseDump { + users: HashMap, + state: HashMap, +} + impl Difluoroborane { pub fn setup() {} @@ -197,6 +205,23 @@ impl Difluoroborane { }) } + pub fn dump_db(&mut self, file: &str) -> Result<(), miette::Error> { + let users = self.users.dump_map()?; + let state = self.statedb.dump_map()?; + let dump = DatabaseDump{users, state}; + let data = toml::ser::to_vec(&dump).map_err(|e| miette::Error::msg(format!("Serializing database dump failed: {}", e)))?; + std::fs::write(file, &data).map_err(|e| miette::Error::msg(format!("writing database dump failed: {}", e)))?; + Ok(()) + } + + pub fn load_db(&mut self, file: &str) -> Result<(), miette::Error> { + let data = std::fs::read(file).into_diagnostic()?; + let dump: DatabaseDump = toml::de::from_slice(&data).into_diagnostic()?; + self.users.load_map(&dump.users)?; + self.statedb.load_map(&dump.state)?; + Ok(()) + } + pub fn run(&mut self) -> Result<(), BFFHError> { let _guard = self.span.enter(); let mut signals = signal_hook_async_std::Signals::new(&[SIGINT, SIGQUIT, SIGTERM]) diff --git a/bffhd/resources/state/db.rs b/bffhd/resources/state/db.rs index e8ad67c..d2237aa 100644 --- a/bffhd/resources/state/db.rs +++ b/bffhd/resources/state/db.rs @@ -1,3 +1,5 @@ +use rkyv::ser::Serializer; +use rkyv::ser::serializers::AllocSerializer; use thiserror::Error; use crate::db; @@ -97,6 +99,30 @@ impl StateDB { self.db.put(&mut txn, key, val, flags)?; Ok(txn.commit()?) } + + pub fn load_map(&self, map: &std::collections::HashMap) -> miette::Result<()> { + use miette::IntoDiagnostic; + let mut txn = self.env.begin_rw_txn().into_diagnostic()?; + let flags = WriteFlags::empty(); + for (key, val) in map { + let mut serializer = AllocSerializer::<1024>::default(); + serializer.serialize_value(val).into_diagnostic()?; + let serialized = ArchivedValue::new(serializer.into_serializer().into_inner()); + self.db.put(&mut txn, &key.as_bytes(), &serialized, flags)?; + } + txn.commit().into_diagnostic()?; + Ok(()) + } + + pub fn dump_map(&self) -> miette::Result> { + let mut map = std::collections::HashMap::new(); + for (key, val) in self.get_all(&self.begin_ro_txn()?)? { + let key_str = core::str::from_utf8(&key).map_err(|_e| miette::Error::msg("state key not UTF8"))?.to_string(); + let val_state: State = rkyv::Deserialize::deserialize(val.as_ref(), &mut rkyv::Infallible).unwrap(); + map.insert(key_str, val_state); + } + Ok(map) + } } #[cfg(test)] diff --git a/bffhd/users/mod.rs b/bffhd/users/mod.rs index 0f6760b..ee4a6c5 100644 --- a/bffhd/users/mod.rs +++ b/bffhd/users/mod.rs @@ -173,6 +173,29 @@ impl Users { Ok(()) } + pub fn load_map(&mut self, dump: &HashMap) -> miette::Result<()> { + let mut txn = unsafe { self.userdb.get_rw_txn() }?; + + self.userdb.clear_txn(&mut txn)?; + + for (uid, data) in dump { + let user = db::User { + id: uid.clone(), + userdata: data.clone(), + }; + + tracing::trace!(%uid, ?user, "Storing user object"); + if let Err(e) = self.userdb.put_txn(&mut txn, uid.as_str(), &user) { + tracing::warn!(error=?e, "failed to add user") + } + } + txn.commit().map_err(crate::db::Error::from)?; + Ok(()) + } + + pub fn dump_map(&self) -> miette::Result> { + return Ok(self.userdb.get_all()?) + } pub fn dump_file(&self, path_str: &str, force: bool) -> miette::Result { let path = Path::new(path_str); let exists = path.exists(); @@ -203,7 +226,7 @@ impl Users { } let mut file = fs::File::create(path).into_diagnostic()?; - let users = self.userdb.get_all()?; + let users = self.dump_map()?; let encoded = toml::ser::to_vec(&users).into_diagnostic()?; file.write_all(&encoded[..]).into_diagnostic()?; diff --git a/bin/bffhd/main.rs b/bin/bffhd/main.rs index 727a1a8..a3e64b3 100644 --- a/bin/bffhd/main.rs +++ b/bin/bffhd/main.rs @@ -57,10 +57,18 @@ fn main() -> miette::Result<()> { .help("Check config for validity") .long("check")) .arg( - Arg::new("dump") + Arg::new("dump-db") .help("Dump all internal databases") - .long("dump") - .conflicts_with("load")) + .long("dump-db") + .alias("dump") + .conflicts_with("dump-users") + .conflicts_with("load-users") + .conflicts_with("load-db") + .takes_value(true) + .value_name("FILE") + .value_hint(ValueHint::AnyPath) + .default_missing_value("bffh-db.toml") + ) .arg( Arg::new("dump-users") .help("Dump the users db to the given file as TOML") @@ -69,18 +77,33 @@ fn main() -> miette::Result<()> { .value_name("FILE") .value_hint(ValueHint::AnyPath) .default_missing_value("users.toml") - .conflicts_with("load")) + .conflicts_with("load-users") + .conflicts_with("load-db") + .conflicts_with("dump-db") + ) .arg( Arg::new("force") .help("force ops that may clobber") .long("force") ) .arg( - Arg::new("load") - .help("Load values into the internal databases") - .long("load") + Arg::new("load-users") + .help("Load users into the internal databases") + .long("load-users") + .alias("load") .takes_value(true) - .conflicts_with("dump")) + .conflicts_with("dump-db") + .conflicts_with("load-db") + .conflicts_with("dump-users") + ) + .arg( + Arg::new("load-db") + .help("Load values into the internal databases") + .long("load-db") + .takes_value(true) + .conflicts_with("dump-db") + .conflicts_with("load-users") + .conflicts_with("dump-users")) .arg(Arg::new("keylog") .help("log TLS keys into PATH. If no path is specified the value of the envvar SSLKEYLOGFILE is used.") .long("tls-key-log") @@ -137,8 +160,16 @@ fn main() -> miette::Result<()> { let mut config = config::read(&PathBuf::from_str(configpath).unwrap())?; - if matches.is_present("dump") { - return Err(miette::miette!("DB Dumping is currently not implemented, except for the users db, using `--dump-users`")); + if matches.is_present("dump-db") { + let mut bffh = Difluoroborane::new(config)?; + let fname = matches.value_of("dump-db").unwrap(); + bffh.dump_db(fname)?; + return Ok(()); + } else if matches.is_present("load-db") { + let mut bffh = Difluoroborane::new(config)?; + let fname = matches.value_of("load-db").unwrap(); + bffh.load_db(fname)?; + return Ok(()); } else if matches.is_present("dump-users") { let bffh = Difluoroborane::new(config)?; @@ -150,12 +181,12 @@ fn main() -> miette::Result<()> { tracing::info!("successfully dumped {} users", number); return Ok(()); - } else if matches.is_present("load") { + } else if matches.is_present("load-users") { let bffh = Difluoroborane::new(config)?; - bffh.users.load_file(matches.value_of("load").unwrap())?; + bffh.users.load_file(matches.value_of("load-users").unwrap())?; - tracing::info!("loaded users from {}", matches.value_of("load").unwrap()); + tracing::info!("loaded users from {}", matches.value_of("load-users").unwrap()); return Ok(()); } else {