diff --git a/daemon/src/storage/mod.rs b/daemon/src/storage/mod.rs index 12438e1..0301050 100644 --- a/daemon/src/storage/mod.rs +++ b/daemon/src/storage/mod.rs @@ -1,3 +1,4 @@ pub(crate) mod credentials; pub(crate) mod pending; pub(crate) mod policy; +pub(crate) mod userdb; diff --git a/daemon/src/storage/userdb.rs b/daemon/src/storage/userdb.rs new file mode 100644 index 0000000..4c22d92 --- /dev/null +++ b/daemon/src/storage/userdb.rs @@ -0,0 +1,138 @@ +use authforge_common::types::Method; +use rusqlite::{params, Connection}; +use std::path::Path; +use thiserror::Error; + +#[derive(Debug, Error)] +pub(crate) enum UserDbError { + #[error("sqlite: {0}")] + Sqlite(#[from] rusqlite::Error), +} + +#[allow(dead_code)] // wired through AppState in Task 2.15. +pub(crate) struct UserDb { + conn: Connection, +} + +#[allow(dead_code)] // wired through AppState in Task 2.15. +impl UserDb { + pub fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).ok(); + } + let conn = Connection::open(path)?; + Self::init(&conn)?; + Ok(Self { conn }) + } + + pub fn open_in_memory() -> Result { + let conn = Connection::open_in_memory()?; + Self::init(&conn)?; + Ok(Self { conn }) + } + + fn init(conn: &Connection) -> Result<(), UserDbError> { + conn.execute_batch( + r#" + CREATE TABLE IF NOT EXISTS users ( + username TEXT PRIMARY KEY, + has_fido2 INTEGER NOT NULL DEFAULT 0, + has_totp INTEGER NOT NULL DEFAULT 0, + last_seen INTEGER NOT NULL DEFAULT 0 + ) WITHOUT ROWID; + "#, + )?; + Ok(()) + } + + pub fn record_enrollment(&self, user: &str, method: Method) -> Result<(), UserDbError> { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0); + let col = method_col(method); + let sql = format!( + "INSERT INTO users (username, {col}, last_seen) VALUES (?1, 1, ?2) + ON CONFLICT(username) DO UPDATE SET {col} = 1, last_seen = ?2" + ); + self.conn.execute(&sql, params![user, now])?; + Ok(()) + } + + pub fn drop_enrollment(&self, user: &str, method: Method) -> Result<(), UserDbError> { + let col = method_col(method); + let sql = format!("UPDATE users SET {col} = 0 WHERE username = ?1"); + self.conn.execute(&sql, params![user])?; + Ok(()) + } + + pub fn users_with(&self, method: Method) -> Result, UserDbError> { + let col = method_col(method); + let sql = format!("SELECT username FROM users WHERE {col} = 1"); + let mut stmt = self.conn.prepare(&sql)?; + let rows = stmt.query_map([], |r| r.get::<_, String>(0))?; + Ok(rows.filter_map(Result::ok).collect()) + } + + pub fn has_any(&self, user: &str) -> Result { + let row: Option = self + .conn + .query_row( + "SELECT 1 FROM users WHERE username = ?1 AND (has_fido2 = 1 OR has_totp = 1)", + params![user], + |r| r.get(0), + ) + .ok(); + Ok(row.is_some()) + } +} + +fn method_col(m: Method) -> &'static str { + match m { + Method::Fido2 => "has_fido2", + Method::Totp => "has_totp", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn open_in_memory_creates_schema() { + let db = UserDb::open_in_memory().unwrap(); + let row: i64 = db + .conn + .query_row( + "SELECT count(*) FROM sqlite_master WHERE type='table' AND name='users'", + [], + |r| r.get(0), + ) + .unwrap(); + assert_eq!(row, 1); + } + + #[test] + fn record_then_list_users_with_method() { + let db = UserDb::open_in_memory().unwrap(); + db.record_enrollment("alice", Method::Fido2).unwrap(); + db.record_enrollment("bob", Method::Totp).unwrap(); + db.record_enrollment("alice", Method::Totp).unwrap(); + + let mut fido2 = db.users_with(Method::Fido2).unwrap(); + fido2.sort(); + assert_eq!(fido2, vec!["alice".to_string()]); + + let mut totp = db.users_with(Method::Totp).unwrap(); + totp.sort(); + assert_eq!(totp, vec!["alice".to_string(), "bob".to_string()]); + + assert!(db.has_any("alice").unwrap()); + assert!(!db.has_any("nobody").unwrap()); + + db.drop_enrollment("alice", Method::Fido2).unwrap(); + assert!(db.users_with(Method::Fido2).unwrap().is_empty()); + // alice still has totp. + assert!(db.has_any("alice").unwrap()); + } +}