support global ssh keys

Cargo.lock

@@ -1812,10 +1812,13 @@ dependencies = [
  "figment",
  "futures",
  "hickory-proto",
+ "lazy_static",
  "log",
+ "regex",
  "serde",
  "sqlx",
  "tarpc",
+ "thiserror",
  "tokio",
  "uuid",
 ]

client/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 nzr-api = { path = "../nzr-api" }
 clap = { version = "4.0.26", features = ["derive"] }
 home = "0.5.4"
-tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
+tokio = { version = "1.0", features = ["fs", "macros", "rt-multi-thread"] }
 tokio-serde = { version = "0.9", features = ["bincode"] }
 tabled = "0.15"
 serde_json = "1"

client/src/main.rs

@@ -123,6 +123,16 @@ enum NetCmd {
     Dump { name: Option<String> },
 }
 
+#[derive(Debug, Subcommand)]
+enum KeyCmd {
+    /// Add a new SSH key
+    Add { path: PathBuf },
+    /// List SSH keys
+    List,
+    /// Delete an SSH key
+    Delete { id: i32 },
+}
+
 #[derive(Debug, Subcommand)]
 enum Commands {
     /// Commands for managing instances
@@ -135,6 +145,11 @@ enum Commands {
         #[command(subcommand)]
         command: NetCmd,
     },
+    /// Commands for managing SSH public keys
+    SshKey {
+        #[command(subcommand)]
+        command: KeyCmd,
+    },
 }
 
 #[derive(Parser, Debug)]
@@ -215,39 +230,6 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                 }
             }
             InstanceCmd::New(args) => {
-                /*
-                let ssh_keys: Vec<String> = {
-                    let key_file = args.sshkey_file.map_or_else(
-                        || {
-                            home::home_dir().map_or_else(
-                                || {
-                                    Err(CommandError::from(
-                                        "SSH keyfile not defined, and couldn't find home directory",
-                                    ))
-                                },
-                                |hd| Ok(hd.join(".ssh/authorized_keys")),
-                            )
-                        },
-                        Ok,
-                    )?;
-
-                    if !key_file.exists() {
-                        Err("SSH keyfile doesn't exist".into())
-                    } else {
-                        match std::fs::read_to_string(&key_file) {
-                            Ok(data) => {
-                                let keys: Vec<String> =
-                                    data.split('\n').map(|s| s.trim().to_owned()).collect();
-                                Ok(keys)
-                            }
-                            Err(err) => Err(CommandError::new(
-                                format!("Couldn't read {} for SSH keys", &key_file.display()),
-                                err,
-                            )),
-                        }
-                    }
-                }?;
-                */
                 let ci_userdata = {
                     if let Some(path) = &args.ci_userdata {
                         if !path.exists() {
@@ -323,7 +305,9 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                 }
             }
             InstanceCmd::Delete { name } => {
-                (client.delete_instance(nzr_api::default_ctx(), name).await?)?;
+                client
+                    .delete_instance(nzr_api::default_ctx(), name)
+                    .await??;
             }
             InstanceCmd::List => {
                 let instances = client.get_instances(nzr_api::default_ctx(), true).await?;
@@ -426,6 +410,30 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                 println!("{}", table.with(tabled::settings::Style::psql()));
             }
         },
+        Commands::SshKey { command } => match command {
+            KeyCmd::Add { path } => {
+                if !path.exists() {
+                    return Err("Provided path doesn't exist".into());
+                }
+
+                let keyfile = tokio::fs::read_to_string(&path).await?;
+                let res = client
+                    .add_ssh_pubkey(nzr_api::default_ctx(), keyfile)
+                    .await??;
+                println!("Key #{} added.", res.id.unwrap_or(-1));
+            }
+            KeyCmd::List => {
+                let keys = client.get_ssh_pubkeys(nzr_api::default_ctx()).await??;
+                let tabular = keys.iter().map(table::SshKey::from);
+                let mut table = tabled::Table::new(tabular);
+                println!("{}", table.with(tabled::settings::Style::psql()));
+            }
+            KeyCmd::Delete { id } => {
+                client
+                    .delete_ssh_pubkey(nzr_api::default_ctx(), id)
+                    .await??;
+            }
+        },
     };
     Ok(())
 }

client/src/table.rs

@@ -40,3 +40,23 @@ impl From<&model::Subnet> for Subnet {
         }
     }
 }
+
+#[derive(Tabled)]
+pub struct SshKey {
+    #[tabled(rename = "ID")]
+    id: i32,
+    #[tabled(rename = "Comment")]
+    comment: String,
+    #[tabled(rename = "Key data")]
+    key_data: String,
+}
+
+impl From<&model::SshPubkey> for SshKey {
+    fn from(value: &model::SshPubkey) -> Self {
+        Self {
+            id: value.id.unwrap_or(-1),
+            comment: value.comment.clone().unwrap_or_default(),
+            key_data: format!("{} {}", value.algorithm, value.key_data),
+        }
+    }
+}

nzr-api/Cargo.toml

@@ -19,6 +19,9 @@ log = "0.4.17"
 sqlx = "0.8"
 diesel = { version = "2.2", optional = true }
 futures = { version = "0.3", optional = true }
+thiserror = "1"
+regex = "1"
+lazy_static = "1"
 
 [dev-dependencies]
 uuid = { version = "1.2.2", features = ["serde", "v4"] }

nzr-api/src/lib.rs

@@ -1,6 +1,6 @@
 use std::net::Ipv4Addr;
 
-use model::{CreateStatus, Instance, Subnet};
+use model::{CreateStatus, Instance, SshPubkey, Subnet};
 
 pub mod args;
 pub mod config;
@@ -51,6 +51,12 @@ pub trait Nazrin {
     async fn delete_subnet(interface: String) -> Result<(), String>;
     /// Gets the cloud-init user-data for the given instance.
     async fn get_instance_userdata(id: i32) -> Result<Vec<u8>, String>;
+    /// Gets all SSH keys stored in the database.
+    async fn get_ssh_pubkeys() -> Result<Vec<SshPubkey>, String>;
+    /// Adds a new SSH public key to the database.
+    async fn add_ssh_pubkey(pub_key: String) -> Result<SshPubkey, String>;
+    /// Deletes an SSH public key from the database.
+    async fn delete_ssh_pubkey(id: i32) -> Result<(), String>;
 }
 
 /// Create a new NazrinClient.

nzr-api/src/mock/mod.rs

@@ -2,7 +2,7 @@ pub mod client;
 #[cfg(test)]
 mod test;
 
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, str::FromStr, sync::Arc};
 
 use tarpc::server::{BaseChannel, Channel as _};
 
@@ -36,6 +36,7 @@ struct MockDb {
     subnet_lease: HashMap<i32, u32>,
     ci_userdatas: HashMap<String, Vec<u8>>,
     create_tasks: HashMap<uuid::Uuid, (model::Instance, bool)>,
+    ssh_keys: Vec<Option<model::SshPubkey>>,
 }
 
 /// Mock Nazrin RPC server for testing, where the full server isn't required.
@@ -252,6 +253,41 @@ impl Nazrin for MockServer {
     async fn garbage_collect(self, _: tarpc::context::Context) -> Result<(), String> {
         todo!()
     }
+
+    async fn get_ssh_pubkeys(
+        self,
+        _: tarpc::context::Context,
+    ) -> Result<Vec<model::SshPubkey>, String> {
+        let db = self.db.read().await;
+
+        Ok(db
+            .ssh_keys
+            .iter()
+            .filter_map(|key| key.as_ref().cloned())
+            .collect())
+    }
+
+    async fn add_ssh_pubkey(
+        self,
+        _: tarpc::context::Context,
+        pub_key: String,
+    ) -> Result<model::SshPubkey, String> {
+        let mut key_model = model::SshPubkey::from_str(&pub_key).map_err(|e| e.to_string())?;
+        let mut db = self.db.write().await;
+        key_model.id = Some(db.ssh_keys.len() as i32);
+        db.ssh_keys.push(Some(key_model.clone()));
+        Ok(key_model)
+    }
+
+    async fn delete_ssh_pubkey(self, _: tarpc::context::Context, id: i32) -> Result<(), String> {
+        let mut db = self.db.write().await;
+        if let Some(key) = db.ssh_keys.get_mut(id as usize) {
+            key.take();
+            Ok(())
+        } else {
+            Err("No such key".into())
+        }
+    }
 }
 
 /// Generates a MockServer task and connected client.

nzr-api/src/model.rs

@@ -1,6 +1,9 @@
 use hickory_proto::rr::Name;
+use lazy_static::lazy_static;
+use regex::Regex;
 use serde::{Deserialize, Serialize};
 use std::{fmt, net::Ipv4Addr};
+use thiserror::Error;
 
 use crate::net::{cidr::CidrV4, mac::MacAddr};
 
@@ -127,3 +130,58 @@ impl SubnetData {
         self.network.host_bits(&self.end_host)
     }
 }
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct SshPubkey {
+    pub id: Option<i32>,
+    pub algorithm: String,
+    pub key_data: String,
+    pub comment: Option<String>,
+}
+
+impl fmt::Display for SshPubkey {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if let Some(comment) = &self.comment {
+            write!(f, "{} {} {}", &self.algorithm, &self.key_data, comment)
+        } else {
+            write!(f, "{} {}", &self.algorithm, &self.key_data)
+        }
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum SshPubkeyParseError {
+    #[error("Key file is not of the expected format")]
+    MissingField,
+    #[error("Key data must be base64-encoded")]
+    InvalidKeyData,
+}
+
+lazy_static! {
+    static ref BASE64_RE: Regex = Regex::new(r"^[A-Za-z0-9+/=]+$").unwrap();
+}
+
+impl std::str::FromStr for SshPubkey {
+    type Err = SshPubkeyParseError;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let mut pieces = s.split(' ');
+        let Some(algorithm) = pieces.next() else {
+            return Err(SshPubkeyParseError::MissingField);
+        };
+        let Some(key_data) = pieces.next() else {
+            return Err(SshPubkeyParseError::MissingField);
+        };
+        // Validate key data
+        if !BASE64_RE.is_match(key_data) {
+            return Err(SshPubkeyParseError::InvalidKeyData);
+        }
+        let comment = pieces.next().map(|s| s.to_owned());
+
+        Ok(Self {
+            id: None,
+            algorithm: algorithm.to_owned(),
+            key_data: key_data.to_owned(),
+            comment,
+        })
+    }
+}

nzrd/src/model/mod.rs

@@ -488,6 +488,19 @@ impl SshPubkey {
         Ok(res)
     }
 
+    pub async fn get(ctx: &Context, id: i32) -> Result<Option<Self>, ModelError> {
+        Ok(ctx
+            .spawn_db(move |mut db| {
+                Self::table()
+                    .find(id)
+                    .select(Self::as_select())
+                    .load::<Self>(&mut db)
+            })
+            .await??
+            .into_iter()
+            .next())
+    }
+
     pub async fn insert(
         ctx: &Context,
         algorithm: impl AsRef<str>,
@@ -513,6 +526,15 @@ impl SshPubkey {
         Ok(ent)
     }
 
+    pub fn api_model(&self) -> nzr_api::model::SshPubkey {
+        nzr_api::model::SshPubkey {
+            id: Some(self.id),
+            algorithm: self.algorithm.clone(),
+            key_data: self.algorithm.clone(),
+            comment: self.comment.clone(),
+        }
+    }
+
     pub async fn delete(self, ctx: &Context) -> Result<(), ModelError> {
         ctx.spawn_db(move |mut db| diesel::delete(&self).execute(&mut db))
             .await??;

nzrd/src/rpc.rs

@@ -1,5 +1,6 @@
 use futures::{future, StreamExt};
 use nzr_api::{args, model, InstanceQuery, Nazrin};
+use std::str::FromStr;
 use std::sync::Arc;
 use tarpc::server::{BaseChannel, Channel};
 use tarpc::tokio_serde::formats::Bincode;
@@ -11,7 +12,7 @@ use uuid::Uuid;
 
 use crate::cmd;
 use crate::ctx::Context;
-use crate::model::{Instance, Subnet};
+use crate::model::{Instance, SshPubkey, Subnet};
 use log::*;
 use std::collections::HashMap;
 
@@ -252,6 +253,41 @@ impl Nazrin for NzrServer {
 
         Ok(db_model.ci_userdata.unwrap_or_default())
     }
+
+    async fn get_ssh_pubkeys(
+        self,
+        _: tarpc::context::Context,
+    ) -> Result<Vec<model::SshPubkey>, String> {
+        SshPubkey::all(&self.ctx).await.map_or_else(
+            |e| Err(e.to_string()),
+            |k| Ok(k.iter().map(|k| k.api_model()).collect()),
+        )
+    }
+
+    async fn add_ssh_pubkey(
+        self,
+        _: tarpc::context::Context,
+        pub_key: String,
+    ) -> Result<model::SshPubkey, String> {
+        let pubkey = model::SshPubkey::from_str(&pub_key).map_err(|e| e.to_string())?;
+
+        SshPubkey::insert(&self.ctx, pubkey.algorithm, pubkey.key_data, pubkey.comment)
+            .await
+            .map_err(|e| e.to_string())
+            .map(|k| k.api_model())
+    }
+
+    async fn delete_ssh_pubkey(self, _: tarpc::context::Context, id: i32) -> Result<(), String> {
+        let Some(key) = SshPubkey::get(&self.ctx, id)
+            .await
+            .map_err(|e| e.to_string())?
+        else {
+            return Err("SSH key with ID doesn't exist".into());
+        };
+
+        key.delete(&self.ctx).await.map_err(|e| e.to_string())?;
+        Ok(())
+    }
 }
 
 #[derive(Debug)]