add option to add ssh key to the server

config.sample.toml

@@ -8,6 +8,7 @@ config_mode = "6v6" # 6v6 or 9v9, defaults to "6v6"
 name = "Spire" # server name. optional, defaults to "Spire"
 tv_name = "SpireTV" # stv name. optional, defaults to "SpireTV"
 image = "spiretf/docker-spire-server" # docker image for the tf2 server. optional, defaults to "spiretf/docker-spire-server"
+ssh_key = "ssh-rsa AAAA..." # ssh key to add to the server. optional
 
 [vultr]
 api_key = "xxx"

src/cloud/mod.rs

@@ -63,11 +63,13 @@ pub trait Cloud: Send + Sync + 'static {
     /// List all running servers on this cloud
     async fn list(&self) -> Result<Vec<Server>>;
     /// Create a new server with the given parameter
-    async fn spawn(&self) -> Result<Created>;
+    async fn spawn(&self, ssh_key_id: Option<&str>) -> Result<Created>;
     /// Destroy a given server
     async fn kill(&self, id: &str) -> Result<()>;
     /// Wait until the server has an ip
     async fn wait_for_ip(&self, id: &str) -> Result<Server>;
+    /// Get the id for the given ssh key
+    async fn get_ssh_key_id(&self, key: &str) -> Result<String>;
 }
 
 #[derive(Debug)]

src/cloud/vultr.rs

@@ -48,7 +48,12 @@ impl Cloud for Vultr {
             .collect())
     }
 
-    async fn spawn(&self) -> Result<Created> {
+    async fn spawn(&self, ssh_key_id: Option<&str>) -> Result<Created> {
+        let key_ids = if let Some(key) = ssh_key_id {
+            vec![key]
+        } else {
+            vec![]
+        };
         let response = self
             .client
             .post("https://api.vultr.com/v2/instances")
@@ -59,6 +64,7 @@ impl Cloud for Vultr {
                 tag: "spire",
                 label: petname(2, "-"),
                 app_id: self.get_app_id("docker").await?,
+                sshkey_id: &key_ids,
             })
             .send()
             .await
@@ -96,6 +102,49 @@ impl Cloud for Vultr {
         };
         Ok(instance.into())
     }
+
+    async fn get_ssh_key_id(&self, ssh_key: &str) -> Result<String> {
+        let response = self
+            .client
+            .get("https://api.vultr.com/v2/ssh-keys")
+            .bearer_auth(&self.token)
+            .send()
+            .await
+            .map_err(NetworkError::from)?;
+        CloudError::from_status_code(response.status())?;
+
+        if !response.status().is_success() {
+            return Err(
+                ResponseError::Other(response.text().await.map_err(NetworkError::from)?).into(),
+            );
+        }
+
+        let response: VultrSshListResponse = response.json().await.map_err(ResponseError::from)?;
+        if let Some(key) = response
+            .ssh_keys
+            .into_iter()
+            .find(|key| key.ssh_key == ssh_key)
+        {
+            Ok(key.id)
+        } else {
+            let response = self
+                .client
+                .post("https://api.vultr.com/v2/ssh-keys")
+                .bearer_auth(&self.token)
+                .json(&VultrCreateSshKeyParams {
+                    name: "Dispenser Key",
+                    ssh_key,
+                })
+                .send()
+                .await
+                .map_err(NetworkError::from)?;
+            CloudError::from_status_code(response.status())?;
+            let response: VultrSshCreateResponse =
+                response.json().await.map_err(ResponseError::from)?;
+
+            Ok(response.ssh_key.id)
+        }
+    }
 }
 
 impl Vultr {
@@ -139,6 +188,7 @@ struct VultrCreateParams<'a> {
     tag: &'a str,
     label: String,
     app_id: u16,
+    sshkey_id: &'a [&'a str],
 }
 
 #[derive(Debug, Deserialize)]
@@ -203,3 +253,26 @@ struct VultrApplicationResponse {
     id: u16,
     short_name: String,
 }
+
+#[derive(Debug, Deserialize)]
+struct VultrSshListResponse {
+    ssh_keys: Vec<VultrSshKeyResponse>,
+}
+
+#[derive(Debug, Deserialize)]
+struct VultrSshCreateResponse {
+    ssh_key: VultrSshKeyResponse,
+}
+
+#[derive(Debug, Deserialize)]
+struct VultrSshKeyResponse {
+    id: String,
+    name: String,
+    ssh_key: String,
+}
+
+#[derive(Serialize)]
+struct VultrCreateSshKeyParams<'a> {
+    name: &'a str,
+    ssh_key: &'a str,
+}

src/config.rs

@@ -71,6 +71,7 @@ pub struct ServerConfig {
     pub name: String,
     #[serde(default = "server_default_tv_name")]
     pub tv_name: String,
+    pub ssh_key: Option<String>,
 }
 
 fn server_default_image() -> String {

src/main.rs

@@ -187,7 +187,14 @@ async fn start(cloud: &dyn Cloud, config: &Config) -> Result<String, Error> {
     if !list.is_empty() {
         return Err(Error::AlreadyRunning);
     }
-    let created = cloud.spawn().await?;
+
+    let ssh_key = if let Some(key) = config.server.ssh_key.as_ref() {
+        Some(cloud.get_ssh_key_id(key).await?)
+    } else {
+        None
+    };
+
+    let created = cloud.spawn(ssh_key.as_deref()).await?;
     let server = cloud.wait_for_ip(&created.id).await?;
 
     println!("Server is booting");