api: cloud-init userdata and ssh keys

client/src/main.rs

@@ -1,13 +1,11 @@
 use clap::{CommandFactory, FromArgMatches, Parser, Subcommand};
+use nzr_api::config;
 use nzr_api::hickory_proto::rr::Name;
 use nzr_api::model;
 use nzr_api::net::cidr::CidrV4;
-use nzr_api::{config, NazrinClient};
 use std::any::{Any, TypeId};
 use std::path::PathBuf;
 use std::str::FromStr;
-use tarpc::tokio_serde::formats::Bincode;
-use tarpc::tokio_util::codec::LengthDelimitedCodec;
 use tokio::net::UnixStream;
 
 mod table;
@@ -35,11 +33,11 @@ pub struct NewInstanceArgs {
     #[arg(short, long, default_value_t = 20)]
     primary_size: u32,
     /// Secndary HDD size, in GiB
-    #[arg(short, long)]
-    secondary_size: Option<u32>,
-    /// File containing a list of SSH keys to use
     #[arg(long)]
-    sshkey_file: Option<PathBuf>,
+    secondary_size: Option<u32>,
+    /// Path to cloud-init userdata, if any
+    #[arg(long)]
+    ci_userdata: Option<PathBuf>,
 }
 
 #[derive(Debug, Subcommand)]
@@ -202,18 +200,12 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
     let cli = Args::from_arg_matches_mut(&mut matches)?;
     let config: config::Config = nzr_api::config::Config::figment().extract()?;
     let conn = UnixStream::connect(&config.rpc.socket_path).await?;
-    let framed_io = LengthDelimitedCodec::builder()
+    let client = nzr_api::new_client(conn);
-        .length_field_type::<u32>()
-        .new_framed(conn);
-    let transport = tarpc::serde_transport::new(framed_io, Bincode::default());
-    let client = NazrinClient::new(Default::default(), transport).spawn();
 
     match cli.command {
         Commands::Instance { command } => match command {
             InstanceCmd::Dump { name, quick } => {
-                let instances = (client
+                let instances = (client.get_instances(nzr_api::default_ctx(), !quick).await?)?;
-                    .get_instances(tarpc::context::current(), !quick)
-                    .await?)?;
                 if let Some(name) = name {
                     if let Some(inst) = instances.iter().find(|f| f.name == name) {
                         println!("{}", serde_json::to_string(inst)?);
@@ -223,6 +215,7 @@ 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(
                         || {
@@ -254,6 +247,21 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                         }
                     }
                 }?;
+                */
+                let ci_userdata = {
+                    if let Some(path) = &args.ci_userdata {
+                        if !path.exists() {
+                            return Err("cloud-init userdata file doesn't exist".into());
+                        } else {
+                            Some(
+                                std::fs::read(path)
+                                    .map_err(|e| format!("Couldn't read userdata file: {e}"))?,
+                            )
+                        }
+                    } else {
+                        None
+                    }
+                };
 
                 let build_args = nzr_api::args::NewInstance {
                     name: args.name,
@@ -264,10 +272,10 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                     cores: args.cores,
                     memory: args.mem,
                     disk_sizes: (args.primary_size, args.secondary_size),
-                    ssh_keys,
+                    ci_userdata,
                 };
                 let task_id = (client
-                    .new_instance(tarpc::context::current(), build_args)
+                    .new_instance(nzr_api::default_ctx(), build_args)
                     .await?)?;
 
                 const MAX_RETRIES: i32 = 5;
@@ -275,7 +283,7 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                 let mut current_pct: f32 = 0.0;
                 loop {
                     let status = client
-                        .poll_new_instance(tarpc::context::current(), task_id)
+                        .poll_new_instance(nzr_api::default_ctx(), task_id)
                         .await;
                     match status {
                         Ok(Some(status)) => {
@@ -315,21 +323,17 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                 }
             }
             InstanceCmd::Delete { name } => {
-                (client
+                (client.delete_instance(nzr_api::default_ctx(), name).await?)?;
-                    .delete_instance(tarpc::context::current(), name)
-                    .await?)?;
             }
             InstanceCmd::List => {
-                let instances = client
+                let instances = client.get_instances(nzr_api::default_ctx(), true).await?;
-                    .get_instances(tarpc::context::current(), true)
-                    .await?;
 
                 let tabular: Vec<table::Instance> =
                     instances?.iter().map(table::Instance::from).collect();
                 let mut table = tabled::Table::new(tabular);
                 println!("{}", table.with(tabled::settings::Style::psql()));
             }
-            InstanceCmd::Prune => (client.garbage_collect(tarpc::context::current()).await?)?,
+            InstanceCmd::Prune => (client.garbage_collect(nzr_api::default_ctx()).await?)?,
         },
         Commands::Net { command } => match command {
             NetCmd::Add(args) => {
@@ -350,12 +354,12 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                     },
                 };
                 (client
-                    .new_subnet(tarpc::context::current(), build_args)
+                    .new_subnet(nzr_api::default_ctx(), build_args)
                     .await?)?;
             }
             NetCmd::Edit(args) => {
                 let mut net = client
-                    .get_subnets(tarpc::context::current())
+                    .get_subnets(nzr_api::default_ctx())
                     .await
                     .map_err(|e| e.to_string())
                     .and_then(|res| {
@@ -391,7 +395,7 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
 
                 // run the update
                 client
-                    .modify_subnet(tarpc::context::current(), net)
+                    .modify_subnet(nzr_api::default_ctx(), net)
                     .await
                     .map_err(|err| format!("RPC error: {}", err))
                     .and_then(|res| {
@@ -401,7 +405,7 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                     })?;
             }
             NetCmd::Dump { name } => {
-                let subnets = (client.get_subnets(tarpc::context::current()).await?)?;
+                let subnets = (client.get_subnets(nzr_api::default_ctx()).await?)?;
                 if let Some(name) = name {
                     if let Some(net) = subnets.iter().find(|s| s.name == name) {
                         println!("{}", serde_json::to_string(net)?);
@@ -411,12 +415,10 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
                 }
             }
             NetCmd::Delete { name } => {
-                (client
+                (client.delete_subnet(nzr_api::default_ctx(), name).await?)?;
-                    .delete_subnet(tarpc::context::current(), name)
-                    .await?)?;
             }
             NetCmd::List => {
-                let subnets = client.get_subnets(tarpc::context::current()).await?;
+                let subnets = client.get_subnets(nzr_api::default_ctx()).await?;
 
                 let tabular: Vec<table::Subnet> =
                     subnets?.iter().map(table::Subnet::from).collect();
@@ -431,7 +433,7 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     if let Err(err) = handle_command().await {
-        if std::any::Any::type_id(&*err).type_id() == TypeId::of::<tarpc::client::RpcError>() {
+        if std::any::Any::type_id(&*err).type_id() == TypeId::of::<nzr_api::RpcError>() {
             log::error!("Error communicating with server: {}", err);
         } else {
             log::error!("{}", err);

nzr-api/src/args.rs

@@ -13,7 +13,7 @@ pub struct NewInstance {
     pub cores: u8,
     pub memory: u32,
     pub disk_sizes: (u32, Option<u32>),
-    pub ssh_keys: Vec<String>,
+    pub ci_userdata: Option<Vec<u8>>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

nzrd/migrations/2024081401_ssh_keys/down.sql

@@ -0,0 +1 @@
+DROP TABLE ssh_keys;

nzrd/migrations/2024081401_ssh_keys/up.sql

@@ -0,0 +1,7 @@
+CREATE TABLE ssh_keys (
+    id INTEGER PRIMARY KEY NOT NULL,
+    algorithm TEXT NOT NULL,
+    key_data TEXT NOT NULL,
+    comment TEXT,
+    UNIQUE(key_data)
+);

nzrd/src/model/mod.rs

@@ -54,6 +54,15 @@ diesel::table! {
     }
 }
 
+diesel::table! {
+    ssh_keys {
+        id -> Integer,
+        algorithm -> Text,
+        key_data -> Text,
+        comment -> Nullable<Text>,
+    }
+}
+
 #[derive(
     AsChangeset,
     Clone,
@@ -456,3 +465,58 @@ impl Transactable for Subnet {
         self.delete(ctx).await
     }
 }
+
+#[derive(Clone, Insertable, Identifiable, Selectable, Queryable)]
+#[diesel(table_name = ssh_keys)]
+pub struct SshPubkey {
+    pub id: i32,
+    pub algorithm: String,
+    pub key_data: String,
+    pub comment: Option<String>,
+}
+
+impl SshPubkey {
+    pub async fn all(ctx: &Context) -> Result<Vec<Self>, ModelError> {
+        let res = ctx
+            .spawn_db(move |mut db| {
+                Self::table()
+                    .select(Self::as_select())
+                    .load::<Self>(&mut db)
+            })
+            .await??;
+
+        Ok(res)
+    }
+
+    pub async fn insert(
+        ctx: &Context,
+        algorithm: impl AsRef<str>,
+        key_data: impl AsRef<str>,
+        comment: Option<impl AsRef<str>>,
+    ) -> Result<Self, ModelError> {
+        use self::ssh_keys::columns;
+        let values = (
+            columns::algorithm.eq(algorithm.as_ref().to_owned()),
+            columns::key_data.eq(key_data.as_ref().to_owned()),
+            columns::comment.eq(comment.map(|s| s.as_ref().to_owned())),
+        );
+
+        let ent = ctx
+            .spawn_db(move |mut db| {
+                diesel::insert_into(Self::table())
+                    .values(values)
+                    .returning(ssh_keys::table::all_columns())
+                    .get_result::<Self>(&mut db)
+            })
+            .await??;
+
+        Ok(ent)
+    }
+
+    pub async fn delete(self, ctx: &Context) -> Result<(), ModelError> {
+        ctx.spawn_db(move |mut db| diesel::delete(&self).execute(&mut db))
+            .await??;
+
+        Ok(())
+    }
+}