From cae3b3a83ed693d1a31105dbb591dc3efcba7b99 Mon Sep 17 00:00:00 2001
From: Nadja Reitzenstein <me@dequbed.space>
Date: Sat, 30 Apr 2022 20:52:32 +0200
Subject: [PATCH] Correcly dump and recreate user db on --load

---
 bffhd/db/raw.rs                |  4 ++++
 bffhd/db/typed.rs              |  4 ++++
 bffhd/resources/modules/mod.rs |  6 +++++-
 bffhd/users/db.rs              | 24 +++++++++++++++++++++++-
 bffhd/users/mod.rs             | 11 +++++++++--
 5 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/bffhd/db/raw.rs b/bffhd/db/raw.rs
index 6302ed1..a05e2dc 100644
--- a/bffhd/db/raw.rs
+++ b/bffhd/db/raw.rs
@@ -52,6 +52,10 @@ impl RawDB {
         txn.del(self.db, key, value.map(AsRef::as_ref))
     }
 
+    pub fn clear(&self, txn: &mut RwTransaction) -> lmdb::Result<()> {
+        txn.clear_db(self.db)
+    }
+
     pub fn iter<'txn, C: lmdb::Cursor<'txn>>(&self, cursor: &'txn mut C) -> lmdb::Iter<'txn> {
         cursor.iter_start()
     }
diff --git a/bffhd/db/typed.rs b/bffhd/db/typed.rs
index fdfa892..0248d5d 100644
--- a/bffhd/db/typed.rs
+++ b/bffhd/db/typed.rs
@@ -142,6 +142,10 @@ impl<A: Adapter> DB<A> {
         self.db.del::<_, &[u8]>(txn, key, None)
     }
 
+    pub fn clear(&self, txn: &mut RwTransaction) -> Result<(), db::Error> {
+        self.db.clear(txn)
+    }
+
     pub fn get_all<'txn, T: Transaction>(&self, txn: &'txn T) -> Result<impl IntoIterator<Item=(&'txn [u8], A::Item)>, db::Error> {
         let mut cursor = self.db.open_ro_cursor(txn)?;
         let it = cursor.iter_start();
diff --git a/bffhd/resources/modules/mod.rs b/bffhd/resources/modules/mod.rs
index 4262264..29f4f9d 100644
--- a/bffhd/resources/modules/mod.rs
+++ b/bffhd/resources/modules/mod.rs
@@ -1,2 +1,6 @@
 
-pub mod fabaccess;
\ No newline at end of file
+pub mod fabaccess;
+
+pub trait MachineModel {
+
+}
\ No newline at end of file
diff --git a/bffhd/users/db.rs b/bffhd/users/db.rs
index 295aa18..85b2952 100644
--- a/bffhd/users/db.rs
+++ b/bffhd/users/db.rs
@@ -1,4 +1,4 @@
-use lmdb::{DatabaseFlags, Environment, Transaction, WriteFlags};
+use lmdb::{DatabaseFlags, Environment, RwTransaction, Transaction, WriteFlags};
 use std::collections::{HashMap};
 use rkyv::Infallible;
 
@@ -100,6 +100,12 @@ pub struct UserDB {
 }
 
 impl UserDB {
+    // TODO: Make an userdb-specific Transaction newtype to make this safe
+    pub unsafe fn get_rw_txn(&self) -> Result<RwTransaction, db::Error> {
+        // The returned transaction is only valid for *this* environment.
+        self.env.begin_rw_txn()
+    }
+
     pub unsafe fn new(env: Arc<Environment>, db: RawDB) -> Self {
         let db = DB::new(db);
         Self { env, db }
@@ -134,6 +140,17 @@ impl UserDB {
         Ok(())
     }
 
+    pub fn put_txn(&self, txn: &mut RwTransaction, uid: &str, user: &User) -> Result<(), db::Error> {
+        let mut serializer = AllocSerializer::<1024>::default();
+        serializer.serialize_value(user).expect("rkyv error");
+        let v = serializer.into_serializer().into_inner();
+        let value = ArchivedValue::new(v);
+
+        let flags = WriteFlags::empty();
+        self.db.put(txn, &uid.as_bytes(), &value, flags)?;
+        Ok(())
+    }
+
     pub fn delete(&self, uid: &str) -> Result<(), db::Error> {
         let mut txn = self.env.begin_rw_txn()?;
         self.db.del(&mut txn, &uid)?;
@@ -141,6 +158,11 @@ impl UserDB {
         Ok(())
     }
 
+    pub fn clear_txn(&self, txn: &mut RwTransaction) -> Result<(), db::Error> {
+        self.db.clear(txn);
+        Ok(())
+    }
+
     pub fn get_all(&self) -> Result<Vec<(String, User)>, db::Error> {
         let txn = self.env.begin_ro_txn()?;
         let iter = self.db.get_all(&txn)?;
diff --git a/bffhd/users/mod.rs b/bffhd/users/mod.rs
index 8941563..3f0f2f4 100644
--- a/bffhd/users/mod.rs
+++ b/bffhd/users/mod.rs
@@ -1,5 +1,5 @@
 use anyhow::Context;
-use lmdb::Environment;
+use lmdb::{Environment, Transaction};
 use once_cell::sync::OnceCell;
 use rkyv::{Archive, Deserialize, Infallible, Serialize};
 use std::collections::HashMap;
@@ -107,6 +107,10 @@ impl Users {
         let f = std::fs::read(path)?;
         let map: HashMap<String, UserData> = toml::from_slice(&f)?;
 
+        let mut txn = unsafe { self.userdb.get_rw_txn()? };
+
+        self.userdb.clear_txn(&mut txn)?;
+
         for (uid, mut userdata) in map {
             userdata.passwd = userdata.passwd.map(|pw| {
                 if !pw.starts_with("$argon2") {
@@ -126,9 +130,12 @@ impl Users {
                 userdata,
             };
             tracing::trace!(%uid, ?user, "Storing user object");
-            self.userdb.put(uid.as_str(), &user);
+            if let Err(e) = self.userdb.put_txn(&mut txn, uid.as_str(), &user) {
+                tracing::warn!(error=?e, "failed to add user")
+            }
         }
 
+        txn.commit()?;
         Ok(())
     }
 }