fix: exit ssh broker on session termination

src/components/backend-core/src/cli.rs

@@ -1,4 +1,5 @@
 use std::env;
+use std::time::Duration;
 
 use anyhow::bail;
 use base64::Engine;
@@ -150,8 +151,44 @@ async fn ssh_broker_cli(args: Vec<String>) -> anyhow::Result<()> {
             .await;
     });
 
+    let open_timeout_ms = payload
+        .get("openTimeoutMs")
+        .and_then(Value::as_u64)
+        .unwrap_or(15_000)
+        .clamp(1_000, 300_000);
+    let runtime_timeout_ms = payload
+        .get("runtimeTimeoutMs")
+        .and_then(Value::as_u64)
+        .unwrap_or(60_000)
+        .clamp(1_000, 3_600_000);
+    let open_timer = tokio::time::sleep(Duration::from_millis(open_timeout_ms));
+    let runtime_timer = tokio::time::sleep(Duration::from_millis(runtime_timeout_ms));
+    tokio::pin!(open_timer);
+    tokio::pin!(runtime_timer);
+
+    let mut opened = false;
     let mut exit_code = 255_i32;
-    while let Some(message) = ws_read.next().await {
+    loop {
+        let message = tokio::select! {
+            _ = &mut runtime_timer => {
+                let _ = tokio::io::stderr()
+                    .write_all(format!("unidesk ssh bridge runtime timeout after {runtime_timeout_ms}ms\n").as_bytes())
+                    .await;
+                exit_code = 124;
+                break;
+            }
+            _ = &mut open_timer, if !opened => {
+                let _ = tokio::io::stderr()
+                    .write_all(format!("unidesk ssh bridge timed out waiting for provider session after {open_timeout_ms}ms\n").as_bytes())
+                    .await;
+                exit_code = 255;
+                break;
+            }
+            message = ws_read.next() => message,
+        };
+        let Some(message) = message else {
+            break;
+        };
         let message = message?;
         let text = match message {
             Message::Text(text) => text,
@@ -160,7 +197,11 @@ async fn ssh_broker_cli(args: Vec<String>) -> anyhow::Result<()> {
         };
         let parsed: Value = serde_json::from_str(&text).unwrap_or_else(|_| json!({}));
         match parsed.get("type").and_then(Value::as_str).unwrap_or("") {
+            "ssh.opened" | "ssh.dispatched" => {
+                opened = true;
+            }
             "ssh.data" => {
+                opened = true;
                 let data = parsed.get("data").and_then(Value::as_str).unwrap_or("");
                 let bytes = base64::engine::general_purpose::STANDARD
                     .decode(data)
@@ -172,10 +213,12 @@ async fn ssh_broker_cli(args: Vec<String>) -> anyhow::Result<()> {
                 }
             }
             "ssh.exit" => {
+                opened = true;
                 exit_code = parsed
                     .get("exitCode")
                     .and_then(Value::as_i64)
                     .unwrap_or(255) as i32;
+                break;
             }
             "ssh.error" => {
                 let _ = tokio::io::stderr()
@@ -191,6 +234,7 @@ async fn ssh_broker_cli(args: Vec<String>) -> anyhow::Result<()> {
                     )
                     .await;
                 exit_code = 255;
+                break;
             }
             _ => {}
         }