Correcly dump and recreate user db on --load

This commit is contained in:
Nadja Reitzenstein 2022-04-30 20:52:32 +02:00
parent cfaf4d509e
commit cae3b3a83e
5 changed files with 45 additions and 4 deletions

View File

@ -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()
}

View File

@ -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();

View File

@ -1,2 +1,6 @@
pub mod fabaccess;
pub mod fabaccess;
pub trait MachineModel {
}

View File

@ -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)?;

View File

@ -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(())
}
}