1 files changed, 150 insertions, 0 deletions
diff --git a/server/src/authorized_keys.rs b/server/src/authorized_keys.rs
new file mode 100644
index 0000000..21651ef
--- /dev/null
+++ b/server/src/authorized_keys.rs
@@ -0,0 +1,150 @@
+use futures::stream::TryStreamExt;
+use rocket::fairing::{self, AdHoc};
+use rocket::serde::Deserialize;
+use rocket::{Build, Rocket};
+use rocket_db_pools::{sqlx, Database};
+use std::borrow::Cow;
+use std::path::{Path, PathBuf};
+use std::sync::Mutex;
+use tokio::fs;
+use tokio::sync::Semaphore;
+
+use crate::Db;
+
+#[derive(Debug, Deserialize)]
+pub struct Config<'a> {
+    authorized_keys: Cow<'a, str>,
+}
+
+#[derive(Clone)]
+struct Key {
+    id: u64,
+    user: String,
+    kind: String,
+    data: String,
+}
+
+struct AuthorizedKeysData {
+    keys: Vec<Key>,
+}
+
+pub struct AuthorizedKeys {
+    data: Mutex<AuthorizedKeysData>,
+    update_lock: Semaphore,
+}
+
+impl AuthorizedKeys {
+    pub async fn new_user_key(
+        &self,
+        config: &Config<'_>,
+        id: u64,
+        user: &str,
+        kind: &str,
+        data: &str,
+    ) {
+        let key = Key {
+            id,
+            user: user.to_string(),
+            kind: kind.to_string(),
+            data: data.to_string(),
+        };
+
+        let path = PathBuf::from(config.authorized_keys.to_string());
+
+        let keys = {
+            let mut data = self.data.lock().unwrap();
+            data.keys.push(key);
+            data.keys.clone()
+        };
+
+        self.update(path.as_path(), &keys).await.unwrap();
+    }
+
+    pub async fn del_user_key(&self, config: &Config<'_>, id: u64, user: &str) {
+        let path = PathBuf::from(config.authorized_keys.to_string());
+
+        let keys = {
+            let mut data = self.data.lock().unwrap();
+
+            if let Some(index) = data.keys.iter().position(|x| x.id == id && x.user == user) {
+                data.keys.remove(index);
+            }
+            data.keys.clone()
+        };
+
+        self.update(path.as_path(), &keys).await.unwrap();
+    }
+
+    async fn update(&self, path: &Path, keys: &Vec<Key>) -> std::io::Result<()> {
+        let mut content = String::new();
+
+        for key in keys {
+            content.push_str(format!("{} {} {}\n", key.kind, key.data, key.user).as_str())
+        }
+
+        let _one_at_a_time = self.update_lock.acquire().await;
+
+        let tmp = path.with_extension("new");
+        fs::write(&tmp, content.as_bytes()).await?;
+        fs::rename(tmp, path).await?;
+
+        Ok(())
+    }
+}
+
+async fn setup_users_keys(
+    authorized_keys: &AuthorizedKeys,
+    config: &Config<'_>,
+    db: &Db,
+) -> anyhow::Result<()> {
+    let keys = sqlx::query!("SELECT id,user,kind,data FROM user_keys ORDER BY id")
+        .fetch(&**db)
+        .map_ok(|r| Key {
+            id: r.id,
+            user: r.user,
+            kind: r.kind,
+            data: r.data,
+        })
+        .try_collect::<Vec<_>>()
+        .await
+        .unwrap();
+
+    let path = PathBuf::from(config.authorized_keys.to_string());
+
+    {
+        let mut data = authorized_keys.data.lock().unwrap();
+        data.keys = keys.clone();
+    }
+
+    authorized_keys.update(path.as_path(), &keys).await?;
+
+    Ok(())
+}
+
+async fn setup_users(rocket: Rocket<Build>) -> fairing::Result {
+    match rocket.state::<Config>() {
+        Some(config) => match rocket.state::<AuthorizedKeys>() {
+            Some(roots) => match Db::fetch(&rocket) {
+                Some(db) => match setup_users_keys(roots, config, db).await {
+                    Ok(_) => Ok(rocket),
+                    Err(_) => Err(rocket),
+                },
+                None => Err(rocket),
+            },
+            None => Err(rocket),
+        },
+        None => Err(rocket),
+    }
+}
+
+pub fn stage() -> AdHoc {
+    AdHoc::on_ignite("Authorized Keys Stage", |rocket| async {
+        rocket
+            .manage(AuthorizedKeys {
+                data: Mutex::new(AuthorizedKeysData { keys: Vec::new() }),
+                update_lock: Semaphore::new(1),
+            })
+            .attach(AdHoc::config::<Config>())
+            .attach(AdHoc::try_on_ignite("Users setup", setup_users))
+    })
+}