implement dumping and loading the full database

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<String, users::db::UserData>,
+    state: HashMap<String, resources::state::State>,
+}
+
 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])

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<String, State>) -> 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<std::collections::HashMap<String, State>> {
+        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)]

bffhd/users/mod.rs

@@ -173,6 +173,29 @@ impl Users {
         Ok(())
     }
 
+    pub fn load_map(&mut self, dump: &HashMap<String,UserData>) -> 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<HashMap<String, UserData>> {
+        return Ok(self.userdb.get_all()?)
+    }
     pub fn dump_file(&self, path_str: &str, force: bool) -> miette::Result<usize> {
         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()?;
 

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 {