add port forwarding (#53)

Panaetius · web-flow · commit 42653ac64f12 · 2026-01-20T12:57:26.000+01:00
* refactor cli

* add port forwarding on remote

* add local port forwarding command
diff --git a/coman/.config/config.toml b/coman/.config/config.toml
@@ -50,6 +50,9 @@ workdir = "{{container_workdir}}"
 {% if iroh_secret%}
 COMAN_IROH_SECRET="{{iroh_secret}}"
 {% endif %}
+{% if port_forward %}
+COMAN_FORWARDED_PORTS="{{port_forward}}"
+{% endif %}
 
 [annotations]
 {% if ssh_public_key %}
diff --git a/coman/src/cli/app.rs b/coman/src/cli/app.rs
@@ -1,13 +1,9 @@
-use std::{error::Error, path::PathBuf, str::FromStr, thread, time::Duration};
+use std::{error::Error, path::PathBuf, str::FromStr, thread};
 
-use base64::prelude::*;
 use clap::{Args, Command, Parser, Subcommand, ValueHint, builder::TypedValueParser};
 use clap_complete::{ArgValueCompleter, CompletionCandidate, Generator, Shell, generate};
-use color_eyre::{Report, Result, eyre::eyre};
-use iroh_ssh::IrohSsh;
+use color_eyre::{Report, Result};
 use itertools::Itertools;
-use pid1::Pid1Settings;
-use rust_supervisor::{ChildType, Supervisor, SupervisorConfig};
 use strum::VariantNames;
 use tokio::sync::mpsc;
 
@@ -16,9 +12,9 @@ use crate::{
     cscs::{
         api_client::{
             client::{EdfSpec as EdfSpecEnum, ScriptSpec as ScriptSpecEnum},
-            types::{JobStatus, PathType},
+            types::PathType,
         },
-        handlers::{cscs_file_list, cscs_job_details, cscs_job_list, file_system_roots},
+        handlers::{cscs_file_list, cscs_job_list, file_system_roots},
     },
     util::types::DockerImageUrl,
 };
@@ -133,6 +129,18 @@ pub enum CscsCommands {
         #[command(subcommand)]
         command: CscsSystemCommands,
     },
+    #[clap(
+        alias("pf"),
+        about = "Forward a local port to a remote port for a job. Note that the port needs to have been exposed with the -P flag on job submission [aliases: pf]"
+    )]
+    PortForward {
+        #[arg(short, long, help = "Local port to forward from")]
+        source_port: u16,
+        #[arg(short, long, help = "Remote port to forward to")]
+        destination_port: u16,
+        #[arg(help="id or name of the job (name uses newest job of that name)", add = ArgValueCompleter::new(job_id_or_name_completer))]
+        job: JobIdOrName,
+    },
 }
 
 #[derive(Args, Clone, Debug)]
@@ -257,6 +265,11 @@ pub enum CscsJobCommands {
             help="Environment variables to set in the container",
             value_hint=ValueHint::Other)]
         env: Vec<(String, String)>,
+        #[clap(short='P',
+            value_name="TARGET",
+            help="Ports to forward from the container",
+            value_hint=ValueHint::Other)]
+        port_forward: Vec<u16>,
         #[clap(short='M',
             value_name="PATH:CONTAINER_PATH",
             value_parser=parse_key_val_colon::<String,String>,
@@ -543,79 +556,3 @@ fn is_bare_string(value_str: &str) -> bool {
 pub fn print_completions<G: Generator>(generator: G, cmd: &mut Command) {
     generate(generator, cmd, cmd.get_name().to_string(), &mut std::io::stdout());
 }
-
-/// Runs a wrapped command in a container-safe way and potentially runs background processes like iroh-ssh
-pub(crate) async fn cli_exec_command(command: Vec<String>) -> Result<()> {
-    // Pid1 takes care of proper terminating of processes and signal handling when running in a container
-    Pid1Settings::new()
-        .enable_log(true)
-        .timeout(Duration::from_secs(2))
-        .launch()
-        .expect("Launch failed");
-
-    let mut supervisor = Supervisor::new(SupervisorConfig::default());
-    supervisor.add_process("iroh-ssh", ChildType::Permanent, || {
-        thread::spawn(|| {
-            let rt = tokio::runtime::Builder::new_current_thread()
-                .enable_all()
-                .build()
-                .expect("couldn't start tokio");
-
-            // Call the asynchronous connect method using the runtime.
-            rt.block_on(async move {
-                let mut builder = IrohSsh::builder().accept_incoming(true).accept_port(15263);
-                if let Ok(secret) = std::env::var("COMAN_IROH_SECRET") {
-                    let secret_key = BASE64_STANDARD.decode(secret).unwrap();
-                    let secret_key: &[u8; 32] = secret_key[0..32].try_into().unwrap();
-                    builder = builder.secret_key(secret_key);
-                }
-
-                let server = builder.build().await.expect("couldn't create iroh server");
-                println!("{}@{}", whoami::username(), server.node_id());
-                loop {
-                    tokio::time::sleep(Duration::from_secs(60)).await;
-                }
-            });
-        })
-    });
-    supervisor.add_process("main-process", ChildType::Temporary, move || {
-        let command = command.clone();
-        thread::spawn(move || {
-            let mut child = std::process::Command::new(command[0].clone())
-                .args(&command[1..])
-                .spawn()
-                .expect("Failed to start compute job");
-            child.wait().expect("Failed to wait on compute job");
-        })
-    });
-
-    let supervisor = supervisor.start_monitoring();
-    loop {
-        thread::sleep(Duration::from_secs(1));
-
-        if let Some(rust_supervisor::ProcessState::Failed | rust_supervisor::ProcessState::Stopped) =
-            supervisor.get_process_state("main-process")
-        {
-            break;
-        }
-    }
-    Ok(())
-}
-
-/// Thin wrapper around iroh proxy
-pub(crate) async fn cli_proxy_command(system: String, job_id: i64) -> Result<()> {
-    let data_dir = get_data_dir();
-    let job_info = cscs_job_details(job_id, Some(system.clone()), None).await?;
-    if job_info.is_none() {
-        return Err(eyre!("remote job does not exist!"));
-    } else if let Some(job_info) = job_info
-        && job_info.status != JobStatus::Running
-    {
-        return Err(eyre!("remote job is not in running state, connection not available"));
-    }
-    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
-        .map_err(|e| eyre!("couldn't proxy ssh connection: {:?}", e))
-}
diff --git a/coman/src/cli/exec.rs b/coman/src/cli/exec.rs
@@ -0,0 +1,136 @@
+use std::{thread, time::Duration};
+
+use base64::prelude::*;
+use color_eyre::Result;
+use iroh::{Endpoint, SecretKey};
+use iroh_ssh::IrohSsh;
+use pid1::Pid1Settings;
+use rust_supervisor::{ChildType, Supervisor, SupervisorConfig};
+use tokio::{net::TcpStream, task::JoinSet};
+
+const SECRET_KEY_ENV: &str = "COMAN_IROH_SECRET";
+const PORT_FORWARD_ENV: &str = "COMAN_FORWARDED_PORTS";
+
+fn get_secret_key() -> Option<Vec<u8>> {
+    if let Ok(secret) = std::env::var(SECRET_KEY_ENV) {
+        let secret_key = BASE64_STANDARD.decode(secret).unwrap();
+        Some(secret_key)
+    } else {
+        None
+    }
+}
+
+#[tokio::main]
+async fn run_ssh() -> Result<()> {
+    let mut builder = IrohSsh::builder().accept_incoming(true).accept_port(15263);
+    if let Some(secret_key) = get_secret_key() {
+        let secret_key: &[u8; 32] = secret_key[0..32].try_into().unwrap();
+        builder = builder.secret_key(secret_key);
+    }
+    let server = builder.build().await.expect("couldn't create iroh server");
+    println!("{}@{}", whoami::username(), server.node_id());
+    loop {
+        tokio::time::sleep(Duration::from_secs(60)).await;
+    }
+}
+
+#[tokio::main]
+async fn port_forward() -> Result<()> {
+    let Some(secret_key) = get_secret_key() else {
+        return Ok(());
+    };
+    let secret_key: &[u8; 32] = secret_key[0..32].try_into().unwrap();
+    let secret_key = SecretKey::from_bytes(secret_key);
+    if let Ok(forwarded_ports) = std::env::var(PORT_FORWARD_ENV) {
+        let mut join_set = JoinSet::new();
+        for port in forwarded_ports.split(',') {
+            let alpn: Vec<u8> = format!("/coman/{port}").into_bytes();
+            let endpoint = Endpoint::builder()
+                .secret_key(secret_key.clone())
+                .alpns(vec![alpn])
+                .bind()
+                .await?;
+            let port = port.to_owned();
+            join_set.spawn(async move {
+                while let Some(incoming) = endpoint.accept().await {
+                    let connection = incoming.await.unwrap();
+                    match connection.accept_bi().await {
+                        Ok((mut iroh_send, mut iroh_recv)) => {
+                            match TcpStream::connect(format!("127.0.0.1:{port}")).await {
+                                Ok(mut stream) => {
+                                    let (mut local_read, mut local_write) = stream.split();
+                                    let a_to_b = async move { tokio::io::copy(&mut local_read, &mut iroh_send).await };
+                                    let b_to_a = async move { tokio::io::copy(&mut iroh_recv, &mut local_write).await };
+
+                                    tokio::select! {
+                                        result = a_to_b => {
+                                            println!("{port}->Iroh stream ended: {result:?}");
+                                        },
+                                        result = b_to_a => {
+                                            println!("Iroh->{port} stream ended: {result:?}");
+                                        },
+                                    };
+                                }
+                                Err(e) => {
+                                    println!("Failed to connect to {port}: {e:?}");
+                                }
+                            }
+                        }
+                        Err(e) => {
+                            println!("Failed to accept stream to {port}: {e:?}");
+                        }
+                    }
+                }
+            });
+        }
+        while let Some(res) = join_set.join_next().await {
+            println!("Task joined: {res:?}");
+        }
+    }
+
+    Ok(())
+}
+
+/// Runs a wrapped command in a container-safe way and potentially runs background processes like iroh-ssh
+pub(crate) async fn cli_exec_command(command: Vec<String>) -> Result<()> {
+    // Pid1 takes care of proper terminating of processes and signal handling when running in a container
+    Pid1Settings::new()
+        .enable_log(true)
+        .timeout(Duration::from_secs(2))
+        .launch()
+        .expect("Launch failed");
+
+    let mut supervisor = Supervisor::new(SupervisorConfig::default());
+    supervisor.add_process("iroh-ssh", ChildType::Permanent, || {
+        thread::spawn(|| {
+            let _ = run_ssh();
+        })
+    });
+    supervisor.add_process("port-forward", ChildType::Permanent, || {
+        thread::spawn(|| {
+            let _ = port_forward();
+        })
+    });
+    supervisor.add_process("main-process", ChildType::Temporary, move || {
+        let command = command.clone();
+        thread::spawn(move || {
+            let mut child = std::process::Command::new(command[0].clone())
+                .args(&command[1..])
+                .spawn()
+                .expect("Failed to start compute job");
+            child.wait().expect("Failed to wait on compute job");
+        })
+    });
+
+    let supervisor = supervisor.start_monitoring();
+    loop {
+        thread::sleep(Duration::from_secs(1));
+
+        if let Some(rust_supervisor::ProcessState::Failed | rust_supervisor::ProcessState::Stopped) =
+            supervisor.get_process_state("main-process")
+        {
+            break;
+        }
+    }
+    Ok(())
+}
diff --git a/coman/src/cli/mod.rs b/coman/src/cli/mod.rs
@@ -0,0 +1,4 @@
+pub mod app;
+pub mod exec;
+pub mod port_forward;
+pub mod proxy;
diff --git a/coman/src/cli/port_forward.rs b/coman/src/cli/port_forward.rs
@@ -0,0 +1 @@
+
diff --git a/coman/src/cli/proxy.rs b/coman/src/cli/proxy.rs
@@ -0,0 +1,24 @@
+use color_eyre::{Result, eyre::eyre};
+
+use crate::{
+    config::get_data_dir,
+    cscs::{api_client::types::JobStatus, handlers::cscs_job_details},
+};
+
+/// Thin wrapper around iroh proxy
+pub(crate) async fn cli_proxy_command(system: String, job_id: i64) -> Result<()> {
+    let data_dir = get_data_dir();
+    let job_info = cscs_job_details(job_id, Some(system.clone()), None).await?;
+    if job_info.is_none() {
+        return Err(eyre!("remote job does not exist!"));
+    } else if let Some(job_info) = job_info
+        && job_info.status != JobStatus::Running
+    {
+        return Err(eyre!("remote job is not in running state, connection not available"));
+    }
+    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
+        .map_err(|e| eyre!("couldn't proxy ssh connection: {:?}", e))
+}
diff --git a/coman/src/config.rs b/coman/src/config.rs
@@ -58,6 +58,8 @@ pub struct CscsConfig {
     #[serde(default)]
     pub env: HashMap<String, String>,
     #[serde(default)]
+    pub port_forward: Vec<u16>,
+    #[serde(default)]
     pub image: Option<String>,
     #[serde(default)]
     pub edf_file_template: String,
diff --git a/coman/src/cscs/api_client/client.rs b/coman/src/cscs/api_client/client.rs
@@ -47,6 +47,7 @@ pub struct JobStartOptions {
     pub stderr: Option<PathBuf>,
     pub container_workdir: Option<String>,
     pub env: Vec<(String, String)>,
+    pub port_forward: Vec<u16>,
     pub mount: Vec<(String, String)>,
     pub edf_spec: EdfSpec,
     pub script_spec: ScriptSpec,
diff --git a/coman/src/cscs/cli.rs b/coman/src/cscs/cli.rs
@@ -16,13 +16,14 @@ use tokio::{
 };
 
 use crate::{
-    cli::JobIdOrName,
+    cli::app::JobIdOrName,
     config::ComputePlatform,
     cscs::{
         api_client::{client::JobStartOptions, types::JobStatus},
         handlers::{
             cscs_file_delete, cscs_file_download, cscs_file_list, cscs_file_upload, cscs_job_cancel, cscs_job_details,
-            cscs_job_list, cscs_job_log, cscs_job_start, cscs_login, cscs_system_list, cscs_system_set,
+            cscs_job_list, cscs_job_log, cscs_job_start, cscs_login, cscs_port_forward, cscs_system_list,
+            cscs_system_set,
         },
     },
 };
@@ -123,6 +124,17 @@ pub(crate) async fn cli_cscs_job_log(
     }
 }
 
+pub(crate) async fn cli_cscs_port_forward(
+    source_port: u16,
+    destination_port: u16,
+    job: JobIdOrName,
+    system: Option<String>,
+    platform: Option<ComputePlatform>,
+) -> Result<()> {
+    let job_id = maybe_job_id_from_name(job, system.clone(), platform.clone()).await?;
+    cscs_port_forward(job_id, source_port, destination_port, system).await
+}
+
 #[allow(clippy::too_many_arguments)]
 pub(crate) async fn cli_cscs_job_start(
     name: Option<String>,
diff --git a/coman/src/cscs/handlers.rs b/coman/src/cscs/handlers.rs
diff --git a/coman/src/main.rs b/coman/src/main.rs
