diff --git a/src/components/backend-core/src/cli.rs b/src/components/backend-core/src/cli.rs index ce7b0491..b536fcbe 100644 --- a/src/components/backend-core/src/cli.rs +++ b/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) -> 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) -> 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) -> 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) -> anyhow::Result<()> { ) .await; exit_code = 255; + break; } _ => {} }