add garbage collection for ssh connections

Author: Panaetius

Cargo.lock

@@ -17,6 +17,7 @@ Table of contents
       * [Logging in](#logging-in)
       * [CLI](#cli)
       * [Terminal UI](#tui)
+      * [SSH](#ssh)
       * [coman.toml config file](#comantoml-config-file)
         * [Editing the config](#editing-the-config)
    * [Development](#development)
@@ -161,7 +162,8 @@ To execute a job on CSCS, run a command like
 coman cscs job submit -i ubuntu:latest -- echo test
 ```
 This will run the command `echo test` using the `ubuntu:latest` docker image and default settings.
-See `coman cscs job submit -h` for more options.
+See `coman cscs job submit -h` for more options. This will also automatically set up an ssh connection for
+the job (use `--no-ssh` to prevent this), see the [SSH](#ssh) section for more details.
 
 You can list your jobs using
 
@@ -270,7 +272,7 @@ name = "myproject" # the name of the project, used to generate job names
 current_system = "daint" # what system/cluster to execute commands on
 current_platform = "HPC" # what platform to execute commands on (valid: HPC, ML or CW)
 account = "..." # the project/group account to use on cscs
-
+ssh_key = "path/to/ssh/public/key.pub" # To use a different public key for SSH connections, other than the default auto-detected id_dsa, id_rsa or id_ecdsa
 
 image = "ubuntu" # default docker image to use
 
@@ -336,6 +338,27 @@ coman config get cscs.current_system
 coman config set cscs.current_system "daint"
 ```
 
+### SSH
+
+`coman cscs job submit` will automatically create an SSH connection for the job. It will search for an
+`id_dsa.pub`, `id_rsa.pub` or `id_ecdsa.pub` file in your `.ssh` folder and use that for the connection,
+unless you specify another key using the `--ssh-key` argument or the `cscs.ssh_key` setting in the config file.
+
+Creating the ssh connection involves several steps, all handled by coman:
+
+- Uploading your ssh public key into the remote coman project folder
+- Setting the public key in the [CSCS SSH hook](https://docs.cscs.ch/software/container-engine/resource-hook/#ssh-hook)
+- Uploading a squash file containing the coman executable to the remote coman project folder
+- Mounting the coman squash file into the container so the coman executable is available in the container
+- Creating an [iroh](https://github.com/n0-computer/iroh) secret key to use for the [QUIC tunnel](https://en.wikipedia.org/wiki/QUIC)
+- Using the coman executable as the entrypoint of the container (wrapping the original command), which
+  allows coman to create an iroh/QUIC tunnel for remote connections, as well as properly handling pid1
+  process signals in the container
+- Creating a local SSH config in the coman data dir (`~/.local/share/coman` by default) containing connection
+  information and the correct iroh proxy command
+- Including the SSH config in `.ssh/config` so it's accessible in other tools
+- Garbage collecting old SSH connections for jobs that are not running anymore
+
 ## Development
 
 ### Prerequisites

README.md

@@ -82,6 +82,7 @@ base64 = "0.22.1"
 dirs = "6.0.0"
 iroh = "0.95.1"
 rand = "0.9.2"
+regex = "1.12.2"
 
 [build-dependencies]
 anyhow = "1.0.90"

coman/Cargo.toml

@@ -73,7 +73,7 @@ pub enum CliCommands {
         command: Vec<String>,
     },
     #[clap(hide = true)]
-    Proxy { job_id: i64 },
+    Proxy { system: String, job_id: i64 },
 }
 
 #[derive(Subcommand, Debug)]
@@ -449,9 +449,9 @@ pub(crate) async fn cli_exec_command(command: Vec<String>) -> Result<()> {
 }
 
 /// Thin wrapper around iroh proxy
-pub(crate) async fn cli_proxy_command(job_id: i64) -> Result<()> {
+pub(crate) async fn cli_proxy_command(system: String, job_id: i64) -> Result<()> {
     let data_dir = get_data_dir();
-    let endpoint_id = std::fs::read_to_string(data_dir.join(format!("{}.endpoint", job_id)))?;
+    let endpoint_id = std::fs::read_to_string(data_dir.join(format!("{}_{}.endpoint", system, job_id)))?;
     println!("{}", endpoint_id);
     iroh_ssh::api::proxy_mode(iroh_ssh::ProxyArgs { node_id: endpoint_id })
         .await

coman/src/cli.rs

@@ -62,6 +62,8 @@ pub struct CscsConfig {
     #[serde(default)]
     pub edf_file_template: String,
     #[serde(default)]
+    pub ssh_key: Option<PathBuf>,
+    #[serde(default)]
     pub command: Vec<String>,
 
     #[serde(default)]

coman/src/config.rs

@@ -3,7 +3,7 @@ use std::os::unix::fs::MetadataExt;
 #[cfg(target_family = "windows")]
 use std::os::windows::fs::MetadataExt;
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     io::{BufWriter, Read, Write},
     path::{Path, PathBuf},
 };
@@ -14,6 +14,7 @@ use eyre::Context;
 use futures::StreamExt;
 use iroh::SecretKey;
 use itertools::Itertools;
+use regex::Regex;
 use reqwest::Url;
 use tokio::{fs::File, io::AsyncWriteExt};
 
@@ -172,13 +173,14 @@ async fn setup_ssh(
     base_path: &Path,
     current_system: &str,
     options: &JobStartOptions,
+    config: &Config,
 ) -> Result<Option<(PathBuf, SecretKey)>> {
     if options.no_ssh {
         return Ok(None);
     }
     let secret = SecretKey::generate(&mut rand::rng());
 
-    let ssh_key = if let Some(path) = options.ssh_key.clone() {
+    let ssh_key = if let Some(path) = options.ssh_key.clone().or(config.values.cscs.ssh_key.clone()) {
         path.canonicalize().map(Some).wrap_err("couldn't get ssh key path")?
     } else {
         // try to figure our ssh key
@@ -202,13 +204,115 @@ async fn setup_ssh(
 
             api_client
                 .upload(current_system, remote_path.clone(), public_key.into_bytes())
-                .await?;
+                .await
+                .wrap_err(eyre!("couldn't upload ssh public key"))?;
             Ok(Some((remote_path, secret)))
         }
         None => Err(eyre!("couldn't find ssh public key, use `--ssh_key` to specify it")),
     }
 }
 
+async fn garbage_collect_ssh(api_client: &CscsApi, current_system: &str) -> Result<()> {
+    let data_dir = get_data_dir();
+    if !data_dir.exists() {
+        return Ok(());
+    }
+    let jobs = api_client.list_jobs(current_system, None).await?;
+    let job_entries: HashSet<_> = jobs.iter().map(|j| format!("{}_{}", current_system, j.id)).collect();
+    let outdated_endpoints: Vec<_> = std::fs::read_dir(&data_dir)?
+        .filter(|d| {
+            d.as_ref().is_ok_and(|e| {
+                e.path().is_file()
+                    && e.file_name().to_string_lossy().ends_with(".endpoint")
+                    && e.file_name().to_string_lossy().starts_with(current_system)
+                    && !job_entries.contains(e.file_name().to_string_lossy().split_once('.').unwrap().0)
+            })
+        })
+        .map(|d| d.unwrap())
+        .collect();
+
+    // delete connection files
+    for d in outdated_endpoints.iter() {
+        std::fs::remove_file(d.path())?;
+    }
+
+    // cleanup ssh config
+    let coman_ssh_config_path = data_dir.join("ssh_config");
+    if !coman_ssh_config_path.exists() {
+        return Ok(());
+    }
+    let mut ssh_content = std::fs::read_to_string(&coman_ssh_config_path)?;
+    for d in outdated_endpoints {
+        let re = Regex::new(
+            format!(
+                r"(?ms)#Start {0}_[^\s]_{1}.*?#End {0}_[^\s]_{1}\n",
+                current_system,
+                d.file_name()
+                    .to_string_lossy()
+                    .split_once('.')
+                    .unwrap()
+                    .0
+                    .rsplit('_')
+                    .next()
+                    .unwrap()
+            )
+            .as_str(),
+        )?;
+        ssh_content = re.replace(&ssh_content, "").to_string();
+    }
+
+    std::fs::write(coman_ssh_config_path, ssh_content)?;
+
+    Ok(())
+}
+
+async fn store_ssh_information(
+    current_system: &str,
+    user_info: &UserInfo,
+    job_id: &i64,
+    job_name: &str,
+    secret_key: &SecretKey,
+) -> Result<String> {
+    let data_dir = get_data_dir();
+    std::fs::write(
+        data_dir.join(format!("{}_{}.endpoint", current_system, job_id)),
+        format!("{}", secret_key.public()),
+    )?;
+    let coman_ssh_config_path = data_dir.join("ssh_config");
+    let coman_ssh_config = std::fs::OpenOptions::new()
+        .create(true)
+        .append(true)
+        .open(coman_ssh_config_path.clone())?;
+    let connection_name = format!("{}-{}-{}", current_system, job_name, job_id);
+    let mut writer = BufWriter::new(coman_ssh_config);
+    write!(
+        writer,
+        "\n#Start {0}\nHost {0}\n    Hostname {1}\n    User {2}\n    ProxyCommand coman proxy {3}{4}\n#End {0}",
+        connection_name,
+        secret_key.public(),
+        user_info.name,
+        current_system,
+        job_id
+    )?;
+    let ssh_dir = dirs::home_dir().ok_or(eyre!("couldn't find home dir"))?.join(".ssh");
+    let ssh_config_path = ssh_dir.join("config");
+    let mut ssh_config = std::fs::OpenOptions::new()
+        .read(true)
+        .append(true)
+        .open(ssh_config_path)?;
+    let mut content = String::new();
+    ssh_config.read_to_string(&mut content)?;
+    if !content.contains(&format!("Include {}", coman_ssh_config_path.clone().display())) {
+        let mut writer = BufWriter::new(ssh_config);
+        write!(
+            writer,
+            "\n\n#coman include\nMatch all\nInclude {}",
+            coman_ssh_config_path.display()
+        )?;
+    }
+    Ok(connection_name)
+}
+
 async fn inject_coman_squash(
     api_client: &CscsApi,
     base_path: &Path,
@@ -272,10 +376,19 @@ async fn inject_coman_squash(
     #[cfg(target_family = "windows")]
     let size = file_meta.file_size() as usize;
 
+    let existing = api_client.list_path(current_system, target.clone()).await?;
+    if !existing.is_empty() {
+        //squash file already present on remote, don't upload if it's the same
+        let entry = existing.first().unwrap();
+        if entry.size.unwrap_or_default() == size {
+            return Ok(Some(target));
+        }
+    }
     //upload squash file
     let transfer_data = api_client
         .transfer_upload(current_system, config.values.cscs.account, target.clone(), size as i64)
-        .await?;
+        .await
+        .wrap_err(eyre!("couldn't upload coman squash file"))?;
     let mut etags: Vec<String> = Vec::new();
     let client = reqwest::Client::new();
     let num_parts = transfer_data.1.num_parts;
@@ -452,11 +565,11 @@ pub async fn cscs_job_start(
         Ok(access_token) => {
             let api_client = CscsApi::new(access_token.0, platform).unwrap();
             let config = Config::new()?;
-            let current_system = &system.unwrap_or(config.values.cscs.current_system);
-            let account = account.or(config.values.cscs.account);
+            let current_system = &system.unwrap_or(config.values.cscs.current_system.clone());
+            let account = account.or(config.values.cscs.account.clone());
             let user_info = api_client.get_userinfo(current_system).await?;
             let job_name = name
-                .or(config.values.name)
+                .or(config.values.name.clone())
                 .unwrap_or(format!("{}-coman", user_info.name));
             let current_system_info = api_client.get_system(current_system).await?;
             let scratch = match current_system_info {
@@ -476,15 +589,16 @@ pub async fn cscs_job_start(
             let container_workdir = options
                 .container_workdir
                 .clone()
-                .unwrap_or(config.values.cscs.workdir.unwrap_or("/scratch".to_owned()));
+                .unwrap_or(config.values.cscs.workdir.clone().unwrap_or("/scratch".to_owned()));
             let base_path = scratch.join(user_info.name.clone()).join(&job_name);
 
             let mut envvars = config.values.cscs.env.clone();
             envvars.extend(options.env.clone());
 
-            let (ssh_public_key_path, secret_key) = setup_ssh(&api_client, &base_path, current_system, &options)
-                .await?
-                .unzip();
+            let (ssh_public_key_path, secret_key) =
+                setup_ssh(&api_client, &base_path, current_system, &options, &config)
+                    .await?
+                    .unzip();
             let coman_squash = inject_coman_squash(&api_client, &base_path, current_system, &options).await?;
 
             let environment_path = handle_edf(
@@ -520,46 +634,10 @@ pub async fn cscs_job_start(
 
             if let Some(secret_key) = secret_key {
                 // store connection information in data dir and set up ssh connection
-                let data_dir = get_data_dir();
-                std::fs::write(
-                    data_dir.join(format!("{}.endpoint", job_id)),
-                    format!("{}", secret_key.public()),
-                )?;
-                let coman_ssh_config_path = data_dir.join("ssh_config");
-                let coman_ssh_config = std::fs::OpenOptions::new()
-                    .create(true)
-                    .append(true)
-                    .open(coman_ssh_config_path.clone())?;
-                let mut writer = BufWriter::new(coman_ssh_config);
-                write!(
-                    writer,
-                    "\nHost {}-{}\n    Hostname {}\n    User {}\n    ProxyCommand coman proxy {}\n",
-                    job_name,
-                    job_id,
-                    secret_key.public(),
-                    user_info.name,
-                    job_id
-                )?;
-                let ssh_dir = dirs::home_dir().ok_or(eyre!("couldn't find home dir"))?.join(".ssh");
-                let ssh_config_path = ssh_dir.join("config");
-                let mut ssh_config = std::fs::OpenOptions::new()
-                    .read(true)
-                    .append(true)
-                    .open(ssh_config_path)?;
-                let mut content = String::new();
-                ssh_config.read_to_string(&mut content)?;
-                if !content.contains(&format!("Include {}", coman_ssh_config_path.clone().display())) {
-                    let mut writer = BufWriter::new(ssh_config);
-                    write!(
-                        writer,
-                        "\n\n#coman include\nMatch all\nInclude {}",
-                        coman_ssh_config_path.display()
-                    )?;
-                }
-                println!(
-                    "Use ssh {}@{}-{} to connect to the job",
-                    user_info.name, job_name, job_id
-                );
+                garbage_collect_ssh(&api_client, current_system).await?;
+                let connection_name =
+                    store_ssh_information(current_system, &user_info, &job_id, &job_name, &secret_key).await?;
+                println!("Use ssh {}@{} to connect to the job", user_info.name, connection_name);
             }
 
             Ok(())

coman/src/cscs/handlers.rs

@@ -142,7 +142,7 @@ async fn main() -> Result<()> {
             },
             cli::CliCommands::Init { destination, name } => Config::create_project_config(destination, name)?,
             cli::CliCommands::Exec { command } => cli_exec_command(command).await?,
-            cli::CliCommands::Proxy { job_id } => cli_proxy_command(job_id).await?,
+            cli::CliCommands::Proxy { system, job_id } => cli_proxy_command(system, job_id).await?,
         },
         None => run_tui(args.tick_rate)?,
     }