make port forwarding wait instead of exit

Author: Panaetius

coman/src/cli/exec.rs

@@ -6,13 +6,13 @@ use iroh::{
     Endpoint, SecretKey,
     protocol::{ProtocolHandler, Router},
 };
-use iroh_ssh::IrohSsh;
 use pid1::Pid1Settings;
 use rust_supervisor::{ChildType, Supervisor, SupervisorConfig};
-use tokio::{net::TcpStream, task::JoinSet};
+use tokio::net::TcpStream;
 
 const SECRET_KEY_ENV: &str = "COMAN_IROH_SECRET";
 const PORT_FORWARD_ENV: &str = "COMAN_FORWARDED_PORTS";
+const SSH_PORT: u16 = 15263;
 
 fn get_secret_key() -> Option<Vec<u8>> {
     if let Ok(secret) = std::env::var(SECRET_KEY_ENV) {
@@ -23,19 +23,6 @@ fn get_secret_key() -> Option<Vec<u8>> {
     }
 }
 
-#[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());
-    tokio::signal::ctrl_c().await?;
-    Ok(())
-}
-
 #[derive(Debug)]
 struct PortForwardHandler {
     port: u16,
@@ -88,30 +75,30 @@ async fn port_forward() -> Result<()> {
     };
     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) {
-        println!("setting up port forwarding...");
-        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.clone()])
-                .bind()
-                .await?;
-
-            let port = port.to_owned();
-            join_set.spawn(async move {
-                let handler = PortForwardHandler {
-                    port: port.parse::<u16>().expect("couldn't parse port"),
-                };
-                Router::builder(endpoint.clone()).accept(&alpn, handler).spawn();
-            });
-        }
-        while let Some(res) = join_set.join_next().await {
-            println!("Task joined: {res:?}");
-        }
+    let mut forwarded_ports = vec!["ssh".to_owned()];
+    if let Ok(env_ports) = std::env::var(PORT_FORWARD_ENV) {
+        forwarded_ports.extend(env_ports.split(',').map(|p| p.to_owned()).collect::<Vec<String>>());
+    }
+    let endpoint = Endpoint::builder().secret_key(secret_key.clone()).bind().await?;
+
+    println!("setting up port forwarding...");
+    let mut builder = Router::builder(endpoint.clone());
+    for port in forwarded_ports {
+        let alpn: Vec<u8> = format!("/coman/{port}").into_bytes();
+        let port = if port == "ssh" {
+            SSH_PORT
+        } else {
+            port.parse::<u16>().expect("couldn't parse port")
+        };
+
+        let handler = PortForwardHandler { port };
+        builder = builder.accept(alpn, handler);
+        println!("set up port forwarding for port {port}");
     }
+    builder.spawn();
+    println!("port forwarding started");
 
+    let _ = tokio::signal::ctrl_c().await;
     Ok(())
 }
 
@@ -125,11 +112,6 @@ pub(crate) async fn cli_exec_command(command: Vec<String>) -> Result<()> {
         .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();

coman/src/cscs/cli.rs

@@ -132,6 +132,7 @@ pub(crate) async fn cli_cscs_port_forward(
     platform: Option<ComputePlatform>,
 ) -> Result<()> {
     let job_id = maybe_job_id_from_name(job, system.clone(), platform.clone()).await?;
+    println!("running port forward for job {job_id}");
     cscs_port_forward(job_id, source_port, destination_port, system).await
 }