Skip to main content

dragoon_server/
users_repo.rs

1//! User + pubkey CRUD. Mirrors `python/.../server/users_repo.py`.
2
3use anyhow::Result;
4use chrono::Utc;
5use rusqlite::{params, Connection, OptionalExtension};
6use serde_json::Value as JsonValue;
7
8use dragoon_proto::pubkey::{fingerprint_pubkey_blob, parse_pubkey_blob};
9
10use crate::auth;
11
12#[derive(Debug, Clone)]
13pub struct UserRow {
14    pub id: i64,
15    pub username: String,
16    pub password_hash: String,
17    pub totp_secret: String,
18    pub recovery_codes_hash: Vec<String>,
19    pub revoked_at: Option<String>,
20}
21
22fn iso_now() -> String {
23    Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Micros, true)
24}
25
26/// Create a new user. Returns `(user_id, totp_secret, plain_recovery_codes)`.
27pub fn create_user(
28    conn: &Connection,
29    username: &str,
30    password: &str,
31) -> Result<(i64, String, Vec<String>)> {
32    let pw_hash = auth::hash_password(password)?;
33    let totp_secret = auth::generate_totp_secret();
34    let (plain, hashes) = auth::generate_recovery_codes(10);
35    let codes_json = serde_json::to_string(&hashes)?;
36    conn.execute(
37        "INSERT INTO users
38         (username, password_hash, totp_secret_enc, recovery_codes_hash, created_at)
39         VALUES (?,?,?,?,?)",
40        params![username, pw_hash, totp_secret, codes_json, iso_now()],
41    )?;
42    Ok((conn.last_insert_rowid(), totp_secret, plain))
43}
44
45pub fn get_user(conn: &Connection, username: &str) -> Result<Option<UserRow>> {
46    let row = conn
47        .query_row(
48            "SELECT id, username, password_hash, totp_secret_enc, recovery_codes_hash, revoked_at
49             FROM users WHERE username=?",
50            [username],
51            |r| {
52                Ok((
53                    r.get::<_, i64>(0)?,
54                    r.get::<_, String>(1)?,
55                    r.get::<_, String>(2)?,
56                    r.get::<_, String>(3)?,
57                    r.get::<_, Option<String>>(4)?,
58                    r.get::<_, Option<String>>(5)?,
59                ))
60            },
61        )
62        .optional()?;
63    let Some((id, username, password_hash, totp_secret, codes_json, revoked_at)) = row else {
64        return Ok(None);
65    };
66    let recovery_codes_hash: Vec<String> =
67        serde_json::from_str(codes_json.as_deref().unwrap_or("[]"))?;
68    Ok(Some(UserRow {
69        id,
70        username,
71        password_hash,
72        totp_secret,
73        recovery_codes_hash,
74        revoked_at,
75    }))
76}
77
78pub fn revoke_user(conn: &Connection, username: &str) -> Result<()> {
79    conn.execute(
80        "UPDATE users SET revoked_at=? WHERE username=?",
81        params![iso_now(), username],
82    )?;
83    Ok(())
84}
85
86pub fn set_password(conn: &Connection, username: &str, new_password: &str) -> Result<()> {
87    let h = auth::hash_password(new_password)?;
88    conn.execute(
89        "UPDATE users SET password_hash=? WHERE username=?",
90        params![h, username],
91    )?;
92    Ok(())
93}
94
95#[derive(Debug, Clone, serde::Serialize)]
96pub struct PubkeyRow {
97    pub fingerprint: String,
98    pub alg: String,
99    pub label: Option<String>,
100    pub added_at: String,
101    pub revoked_at: Option<String>,
102}
103
104pub fn add_pubkey(
105    conn: &Connection,
106    user_id: i64,
107    pubkey_blob: &[u8],
108    label: Option<&str>,
109) -> Result<String> {
110    let parsed = parse_pubkey_blob(pubkey_blob)?;
111    let fp = fingerprint_pubkey_blob(pubkey_blob);
112    conn.execute(
113        "INSERT INTO user_pubkeys (user_id, fingerprint, alg, pubkey_blob, label, added_at)
114         VALUES (?,?,?,?,?,?)",
115        params![user_id, fp, parsed.alg(), pubkey_blob, label, iso_now()],
116    )?;
117    Ok(fp)
118}
119
120pub fn list_pubkeys(conn: &Connection, user_id: i64) -> Result<Vec<PubkeyRow>> {
121    let mut stmt = conn.prepare(
122        "SELECT fingerprint, alg, label, added_at, revoked_at FROM user_pubkeys
123         WHERE user_id=? ORDER BY added_at ASC",
124    )?;
125    let rows = stmt
126        .query_map([user_id], |r| {
127            Ok(PubkeyRow {
128                fingerprint: r.get(0)?,
129                alg: r.get(1)?,
130                label: r.get(2)?,
131                added_at: r.get(3)?,
132                revoked_at: r.get(4)?,
133            })
134        })?
135        .collect::<rusqlite::Result<Vec<_>>>()?;
136    Ok(rows)
137}
138
139pub fn revoke_pubkey(conn: &Connection, user_id: i64, fingerprint: &str) -> Result<bool> {
140    let n = conn.execute(
141        "UPDATE user_pubkeys SET revoked_at=? WHERE user_id=? AND fingerprint=?",
142        params![iso_now(), user_id, fingerprint],
143    )?;
144    Ok(n > 0)
145}
146
147/// Helper used by routes/messages — strict username -> id resolution.
148pub fn lookup_user_id_by_name(conn: &Connection, username: &str) -> Result<Option<i64>> {
149    let id: Option<i64> = conn
150        .query_row(
151            "SELECT id FROM users WHERE username=? AND revoked_at IS NULL",
152            [username],
153            |r| r.get(0),
154        )
155        .optional()?;
156    Ok(id)
157}
158
159// keep clippy quiet about an unused import on builds without tests
160#[allow(dead_code)]
161fn _silence() -> JsonValue {
162    JsonValue::Null
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168    use ed25519_dalek::SigningKey;
169
170    fn fresh() -> Connection {
171        let c = crate::db::connect_in_memory().unwrap();
172        crate::db::bootstrap(&c).unwrap();
173        c
174    }
175
176    fn ed25519_pub_blob(seed: u8) -> Vec<u8> {
177        let sk = SigningKey::from_bytes(&[seed; 32]);
178        dragoon_proto::task_sig::ed25519_public_blob(&sk)
179    }
180
181    #[test]
182    fn create_user_returns_totp_and_codes() {
183        let c = fresh();
184        let (uid, secret, codes) = create_user(&c, "alice", "hunter2").unwrap();
185        assert!(uid > 0);
186        assert!(!secret.is_empty());
187        assert_eq!(codes.len(), 10);
188    }
189
190    #[test]
191    fn get_user_round_trip() {
192        let c = fresh();
193        let (uid, _, _) = create_user(&c, "alice", "x").unwrap();
194        let row = get_user(&c, "alice").unwrap().unwrap();
195        assert_eq!(row.id, uid);
196        assert!(auth::verify_password("x", &row.password_hash));
197        assert_eq!(row.recovery_codes_hash.len(), 10);
198    }
199
200    #[test]
201    fn add_list_revoke_pubkey() {
202        let c = fresh();
203        let (uid, _, _) = create_user(&c, "alice", "pw").unwrap();
204        let blob = ed25519_pub_blob(11);
205        let fp = add_pubkey(&c, uid, &blob, Some("laptop")).unwrap();
206        assert!(fp.starts_with("SHA256:"));
207        let keys = list_pubkeys(&c, uid).unwrap();
208        assert_eq!(keys.len(), 1);
209        assert_eq!(keys[0].fingerprint, fp);
210        assert_eq!(keys[0].label.as_deref(), Some("laptop"));
211        assert!(revoke_pubkey(&c, uid, &fp).unwrap());
212        let keys = list_pubkeys(&c, uid).unwrap();
213        assert!(keys[0].revoked_at.is_some());
214    }
215
216    #[test]
217    fn lookup_user_id_helper() {
218        let c = fresh();
219        let (uid, _, _) = create_user(&c, "bob", "x").unwrap();
220        assert_eq!(lookup_user_id_by_name(&c, "bob").unwrap(), Some(uid));
221        assert_eq!(lookup_user_id_by_name(&c, "ghost").unwrap(), None);
222    }
223}