diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/cli/machine.rs | 124 | ||||
| -rw-r--r-- | src/cli/spec/machine.rs | 2 | ||||
| -rw-r--r-- | src/remote/attach.rs | 269 |
3 files changed, 320 insertions, 75 deletions
diff --git a/src/cli/machine.rs b/src/cli/machine.rs index 52828f75..2459805a 100644 --- a/src/cli/machine.rs +++ b/src/cli/machine.rs @@ -1,3 +1,8 @@ +use std::io::IsTerminal as _; + +use crossterm::event::{Event, KeyCode, KeyEventKind, KeyModifiers}; +use crossterm::style::{Attribute, SetAttribute}; +use crossterm::{cursor, execute, terminal}; use serde::Serialize; use crate::client::endpoint::{EndpointCatalog, ProfileId, MAX_LABEL_BYTES}; @@ -206,7 +211,7 @@ fn reconnect(args: &[String]) -> std::io::Result<i32> { struct AddArgs { target: String, label: Option<String>, - session: String, + session: Option<String>, } fn parse_add_args(args: &[String]) -> Result<AddArgs, String> { @@ -249,7 +254,6 @@ fn parse_add_args(args: &[String]) -> Result<AddArgs, String> { "usage: herdr machine add <ssh-target> [--label <label>] [--remote-session <name>]" .to_owned() })?; - let session = session.unwrap_or_else(|| crate::session::DEFAULT_SESSION_NAME.to_owned()); Ok(AddArgs { target, label, @@ -304,6 +308,27 @@ fn add(args: &[String]) -> std::io::Result<i32> { return Ok(2); } }; + let mut setup = None; + let session = + if session.is_none() && std::io::stdin().is_terminal() && std::io::stderr().is_terminal() { + let discovered = (|| { + let connection = crate::remote::SavedSshSetup::connect(&target)?; + let sessions = connection.running_sessions()?; + let session = select_remote_session(&sessions, &target)?; + setup = Some(connection); + Ok::<_, std::io::Error>(session) + })(); + match discovered { + Ok(session) => session, + Err(error) => { + eprintln!("error: {error}; machine was not saved"); + crate::remote::print_saved_ssh_error_hint(&error, &target); + return Ok(1); + } + } + } else { + session.unwrap_or_else(|| crate::session::DEFAULT_SESSION_NAME.to_owned()) + }; let label_is_default = label.is_none(); let label = label.unwrap_or_else(|| default_label(&target, &session)); let mut catalog = load_catalog()?; @@ -320,7 +345,11 @@ fn add(args: &[String]) -> std::io::Result<i32> { return Ok(2); } } - let metadata = match crate::remote::prepare_saved_ssh(&target, &session) { + let metadata = match setup + .map(Ok) + .unwrap_or_else(|| crate::remote::SavedSshSetup::connect(&target)) + .and_then(|setup| setup.prepare(&session)) + { Ok(metadata) => metadata, Err(error) => { eprintln!("error: {error}; machine was not saved"); @@ -361,6 +390,83 @@ fn add(args: &[String]) -> std::io::Result<i32> { Ok(0) } +fn select_remote_session(sessions: &[String], target: &str) -> std::io::Result<String> { + match sessions { + [] => return Ok(crate::session::DEFAULT_SESSION_NAME.to_owned()), + [session] => return Ok(session.clone()), + _ => {} + } + + let _raw_mode = RawModeGuard::enable()?; + let mut output = std::io::stderr(); + let mut selected = 0; + render_remote_session_picker(&mut output, target, sessions, selected, false)?; + loop { + let Event::Key(key) = crossterm::event::read()? else { + continue; + }; + if !matches!(key.kind, KeyEventKind::Press | KeyEventKind::Repeat) { + continue; + } + match key.code { + KeyCode::Up => selected = selected.checked_sub(1).unwrap_or(sessions.len() - 1), + KeyCode::Down => selected = (selected + 1) % sessions.len(), + KeyCode::Enter => return Ok(sessions[selected].clone()), + KeyCode::Esc => break, + KeyCode::Char('c') if key.modifiers.contains(KeyModifiers::CONTROL) => break, + _ => continue, + } + render_remote_session_picker(&mut output, target, sessions, selected, true)?; + } + Err(std::io::Error::new( + std::io::ErrorKind::Interrupted, + "remote session selection cancelled", + )) +} + +struct RawModeGuard; + +impl RawModeGuard { + fn enable() -> std::io::Result<Self> { + terminal::enable_raw_mode()?; + Ok(Self) + } +} + +impl Drop for RawModeGuard { + fn drop(&mut self) { + let _ = terminal::disable_raw_mode(); + } +} + +fn render_remote_session_picker( + output: &mut impl std::io::Write, + target: &str, + sessions: &[String], + selected: usize, + redraw: bool, +) -> std::io::Result<()> { + if redraw { + execute!(output, cursor::MoveUp((sessions.len() + 2) as u16))?; + } + execute!(output, terminal::Clear(terminal::ClearType::CurrentLine))?; + write!(output, "Running sessions on {target}:\r\n")?; + for (index, session) in sessions.iter().enumerate() { + execute!(output, terminal::Clear(terminal::ClearType::CurrentLine))?; + if index == selected { + execute!(output, SetAttribute(Attribute::Bold))?; + write!(output, "> {session}")?; + execute!(output, SetAttribute(Attribute::Reset))?; + write!(output, "\r\n")?; + } else { + write!(output, " {session}\r\n")?; + } + } + execute!(output, terminal::Clear(terminal::ClearType::CurrentLine))?; + write!(output, "↑/↓ select · Enter confirm · Esc cancel\r\n")?; + output.flush() +} + fn rename(args: &[String]) -> std::io::Result<i32> { let args = super::expand_equals_args(args, &["--label"]); let [raw_id, flag, label] = args.as_slice() else { @@ -480,8 +586,8 @@ mod tests { #[test] fn add_parser_preserves_values_across_argument_orders() { for (args, session) in [ - (vec!["--label", "coder", "workstation.coder"], "default"), - (vec!["workstation.coder", "--label", "coder"], "default"), + (vec!["--label", "coder", "workstation.coder"], None), + (vec!["workstation.coder", "--label", "coder"], None), ( vec![ "--remote-session", @@ -490,7 +596,7 @@ mod tests { "--label", "coder", ], - "agents", + Some("agents"), ), ( vec![ @@ -498,7 +604,7 @@ mod tests { "--remote-session=agents", "workstation.coder", ], - "agents", + Some("agents"), ), ] { let args = args.into_iter().map(str::to_owned).collect::<Vec<_>>(); @@ -507,7 +613,7 @@ mod tests { AddArgs { target: "workstation.coder".into(), label: Some("coder".into()), - session: session.into(), + session: session.map(str::to_owned), }, "{args:?}" ); @@ -518,7 +624,7 @@ mod tests { fn add_parser_leaves_label_unset_without_flag() { let parsed = parse_add_args(&["workstation.coder".to_owned()]).unwrap(); assert_eq!(parsed.label, None); - assert_eq!(parsed.session, "default"); + assert_eq!(parsed.session, None); } #[test] diff --git a/src/cli/spec/machine.rs b/src/cli/spec/machine.rs index 19a34088..dc4eefa5 100644 --- a/src/cli/spec/machine.rs +++ b/src/cli/spec/machine.rs @@ -36,7 +36,7 @@ pub(super) fn command() -> Command { ) .arg( option("remote-session", "NAME") - .help("Set the explicit Herdr session on the remote machine"), + .help("Select a session explicitly (default without an interactive terminal)"), ), ) .subcommand( diff --git a/src/remote/attach.rs b/src/remote/attach.rs index a4a00dea..82caeb9a 100644 --- a/src/remote/attach.rs +++ b/src/remote/attach.rs @@ -119,72 +119,145 @@ pub(crate) fn check_saved_ssh(target: &str, session: &str) -> io::Result<()> { } } -pub(crate) fn prepare_saved_ssh( - target: &str, - session_name: &str, -) -> io::Result<Option<crate::client::endpoint::SshMachineMetadata>> { - super::validate_remote_target(target) - .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?; - crate::session::validate_name(session_name) - .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?; - let manage_ssh_config = crate::config::Config::load() - .config - .remote - .manage_ssh_config; - let ssh = RemoteSsh::new( - target.to_owned(), - manage_ssh_config, - session_name.to_owned(), - ); - let prepared = prepare_remote_herdr(&ssh, false, true)?; - ensure_remote_server_ready( - &ssh, - &prepared.remote_herdr, - prepared.stop_after_install_approved, - false, - true, - )?; +pub(crate) struct SavedSshSetup { + ssh: RemoteSsh, + remote_herdr: RemoteHerdr, + candidates: Vec<RemoteHerdr>, +} + +impl SavedSshSetup { + pub(crate) fn connect(target: &str) -> io::Result<Self> { + super::validate_remote_target(target) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?; + let manage = crate::config::Config::load() + .config + .remote + .manage_ssh_config; + let ssh = RemoteSsh::new( + target.to_owned(), + manage, + crate::session::DEFAULT_SESSION_NAME.to_owned(), + ); + let remote_herdr = RemoteHerdr::for_platform(detect_remote_platform(&ssh)?); + let candidates = remote_binary_candidates(&ssh, &remote_herdr)?; + Ok(Self { + ssh, + remote_herdr, + candidates, + }) + } - // The bridge already owns daemon startup. EOF closes only this temporary attachment, - // leaving the named server running even when no local TUI is open yet. - let command = prepared - .remote_herdr - .executable - .saved_bridge_command(session_name); - let output = ssh.shell_output(&prepared.remote_herdr.platform, &command)?; - if !output.status.success() { - return Err(command_failed("remote server startup failed", &output)); + pub(crate) fn running_sessions(&self) -> io::Result<Vec<String>> { + // A fresh host must still reach the normal installation approval. + if self.candidates.is_empty() { + return Ok(Vec::new()); + } + let mut failure = String::new(); + for candidate in &self.candidates { + let output = self.ssh.shell_output( + &candidate.platform, + &candidate.executable.command(&["session", "list", "--json"]), + )?; + if output.status.code() == Some(255) { + return Err(command_failed("remote SSH connection failed", &output)); + } + if !output.status.success() { + failure = command_failed("remote session query failed", &output).to_string(); + continue; + } + match serde_json::from_slice::<RemoteSessionListJson>(&output.stdout) { + Ok(list) => { + return Ok(list + .sessions + .into_iter() + .filter(|session| { + session.running && crate::session::validate_name(&session.name).is_ok() + }) + .map(|session| session.name) + .collect()) + } + Err(error) => failure = format!("invalid remote session list: {error}"), + } + } + Err(io::Error::other(format!( + "could not discover remote Herdr sessions: {failure}; specify --remote-session to continue" + ))) } - match remote_server_status(&ssh, &prepared.remote_herdr, true)? { - RemoteServerStatus::Running { - endpoint_protocol_generation, - surface_interest, - health_check, - detached_server_daemon, - .. - } if remote_server_restart_reason( - endpoint_protocol_generation, - detached_server_daemon, + + pub(crate) fn prepare( + mut self, + session_name: &str, + ) -> io::Result<Option<crate::client::endpoint::SshMachineMetadata>> { + crate::session::validate_name(session_name) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?; + self.ssh.session_name = session_name.to_owned(); + let Self { + ssh, + remote_herdr, + candidates, + } = self; + let prepared = + prepare_discovered_remote_herdr(&ssh, remote_herdr, &candidates, false, true)?; + ensure_remote_server_ready( + &ssh, + &prepared.remote_herdr, + prepared.stop_after_install_approved, + false, true, - surface_interest, - health_check, - ) - .is_none() => - { - Ok(prepared.remote_herdr.machine_metadata().or_else(|| { - discover_remote_api_metadata(&ssh, session_name) - .inspect_err( - |error| tracing::debug!(%error, "could not capture SSH setup metadata"), - ) - .ok() - })) + )?; + + // The bridge already owns daemon startup. EOF closes only this temporary attachment, + // leaving the named server running even when no local TUI is open yet. + let command = prepared + .remote_herdr + .executable + .saved_bridge_command(session_name); + let output = ssh.shell_output(&prepared.remote_herdr.platform, &command)?; + if !output.status.success() { + return Err(command_failed("remote server startup failed", &output)); + } + match remote_server_status(&ssh, &prepared.remote_herdr, true)? { + RemoteServerStatus::Running { + endpoint_protocol_generation, + surface_interest, + health_check, + detached_server_daemon, + .. + } if remote_server_restart_reason( + endpoint_protocol_generation, + detached_server_daemon, + true, + surface_interest, + health_check, + ) + .is_none() => + { + Ok(prepared.remote_herdr.machine_metadata().or_else(|| { + discover_remote_api_metadata(&ssh, session_name) + .inspect_err( + |error| tracing::debug!(%error, "could not capture SSH setup metadata"), + ) + .ok() + })) + } + _ => Err(io::Error::other( + "remote server is not ready for saved machines", + )), } - _ => Err(io::Error::other( - "remote server is not ready for saved machines", - )), } } +#[derive(Deserialize)] +struct RemoteSessionListJson { + sessions: Vec<RemoteSessionJson>, +} + +#[derive(Deserialize)] +struct RemoteSessionJson { + name: String, + running: bool, +} + #[derive(Debug, Clone, PartialEq, Eq)] struct RemotePlatform { os: &'static str, @@ -1201,19 +1274,36 @@ pub(super) fn prepare_remote_herdr( ) -> io::Result<PreparedRemoteHerdr> { let platform = detect_remote_platform(ssh)?; let remote_herdr = RemoteHerdr::for_platform(platform); + let candidates = remote_binary_candidates(ssh, &remote_herdr)?; + prepare_discovered_remote_herdr( + ssh, + remote_herdr, + &candidates, + live_handoff_enabled, + require_surface_interest, + ) +} + +fn prepare_discovered_remote_herdr( + ssh: &RemoteSsh, + remote_herdr: RemoteHerdr, + candidates: &[RemoteHerdr], + live_handoff_enabled: bool, + require_surface_interest: bool, +) -> io::Result<PreparedRemoteHerdr> { if remote_herdr.platform.is_windows() { return prepare_windows_remote_herdr( ssh, remote_herdr, + candidates, live_handoff_enabled, require_surface_interest, ); } let override_binary = remote_binary_override_path()?; - let remote_binary_candidates = remote_binary_candidates(ssh, &remote_herdr)?; if override_binary.is_none() { - for candidate in &remote_binary_candidates { + for candidate in candidates { if remote_binary_supports_endpoint_requirement(ssh, candidate, require_surface_interest) .unwrap_or(false) { @@ -1236,7 +1326,7 @@ pub(super) fn prepare_remote_herdr( } let mut stop_after_install_approved = false; - if let Some(status_probe_herdr) = remote_binary_candidates.first().or_else(|| { + if let Some(status_probe_herdr) = candidates.first().or_else(|| { remote_binary_exists(ssh, &remote_herdr) .ok() .and_then(|exists| exists.then_some(&remote_herdr)) @@ -1299,14 +1389,14 @@ pub(super) fn find_installed_remote_herdr(ssh: &RemoteSsh) -> io::Result<RemoteH fn prepare_windows_remote_herdr( ssh: &RemoteSsh, remote_herdr: RemoteHerdr, + candidates: &[RemoteHerdr], live_handoff_enabled: bool, require_surface_interest: bool, ) -> io::Result<PreparedRemoteHerdr> { let override_package = remote_binary_override_path()?; let custom_package = override_package.is_some(); - let candidates = remote_binary_candidates(ssh, &remote_herdr)?; if !custom_package { - for candidate in &candidates { + for candidate in candidates { if remote_binary_supports_endpoint_requirement( ssh, candidate, @@ -1524,7 +1614,7 @@ fn remote_binary_candidates( fn windows_remote_binary_candidate_command() -> String { windows_powershell_script_command(&format!( - r#"function Emit-HerdrPath([string]$CandidatePath) {{ if ([string]::IsNullOrWhiteSpace($CandidatePath) -or -not (Test-Path -LiteralPath $CandidatePath -PathType Leaf)) {{ return }}; $candidateFullPath = [System.IO.Path]::GetFullPath($CandidatePath); $encodedCandidate = [System.Convert]::ToBase64String([System.Text.Encoding]::UTF8.GetBytes($candidateFullPath)); [Console]::Out.WriteLine('{WINDOWS_REMOTE_PATH_MARKER}' + $encodedCandidate) }}; $pathCommand = Get-Command herdr.exe -CommandType Application -ErrorAction SilentlyContinue | Select-Object -First 1; if ($null -ne $pathCommand) {{ Emit-HerdrPath $pathCommand.Source }}; $herdrHome = if ([string]::IsNullOrWhiteSpace($env:HERDR_HOME)) {{ Join-Path $env:USERPROFILE '.herdr' }} else {{ $env:HERDR_HOME }}; $activeJunction = Get-Item -LiteralPath (Join-Path $herdrHome 'packages\standalone\current') -Force -ErrorAction SilentlyContinue; if ($null -ne $activeJunction -and -not [string]::IsNullOrWhiteSpace([string]$activeJunction.Target)) {{ Emit-HerdrPath (Join-Path ([string]$activeJunction.Target) 'herdr.exe') }}; exit 0"# + r#"function Emit-HerdrPath([string]$CandidatePath) {{ if ([string]::IsNullOrWhiteSpace($CandidatePath) -or -not (Test-Path -LiteralPath $CandidatePath -PathType Leaf)) {{ return }}; $candidateFullPath = [System.IO.Path]::GetFullPath($CandidatePath); $encodedCandidate = [System.Convert]::ToBase64String([System.Text.Encoding]::UTF8.GetBytes($candidateFullPath)); [Console]::Out.WriteLine('{WINDOWS_REMOTE_PATH_MARKER}' + $encodedCandidate) }}; $pathCommand = Get-Command herdr.exe -CommandType Application -ErrorAction SilentlyContinue | Select-Object -First 1; if ($null -ne $pathCommand) {{ Emit-HerdrPath $pathCommand.Source }}; $userSid = [System.Security.Principal.WindowsIdentity]::GetCurrent().User.Value; Get-CimInstance Win32_Process -Filter "Name = 'herdr.exe' OR Name = 'herdr-dev.exe'" -ErrorAction SilentlyContinue | ForEach-Object {{ $owner = Invoke-CimMethod -InputObject $_ -MethodName GetOwnerSid -ErrorAction SilentlyContinue; if ($owner.ReturnValue -eq 0 -and $owner.Sid -eq $userSid) {{ Emit-HerdrPath $_.ExecutablePath }} }}; $herdrHome = if ([string]::IsNullOrWhiteSpace($env:HERDR_HOME)) {{ Join-Path $env:USERPROFILE '.herdr' }} else {{ $env:HERDR_HOME }}; $activeJunction = Get-Item -LiteralPath (Join-Path $herdrHome 'packages\standalone\current') -Force -ErrorAction SilentlyContinue; if ($null -ne $activeJunction -and -not [string]::IsNullOrWhiteSpace([string]$activeJunction.Target)) {{ Emit-HerdrPath (Join-Path ([string]$activeJunction.Target) 'herdr.exe') }}; exit 0"# )) } @@ -4614,6 +4704,55 @@ mod tests { assert!(command.contains("$activeJunction.Target")); } + #[cfg(windows)] + #[test] + fn windows_process_candidates_require_a_verified_current_user_owner() { + // Supply ownership results without needing another Windows account. + // Execute the real discovery script, including its path validation. + let fixture = r#" +function Get-Command {} +function Get-Item {} +function Get-CimInstance { + $sid = [System.Security.Principal.WindowsIdentity]::GetCurrent().User.Value + [pscustomobject]@{ ExecutablePath=$env:HERDR_TEST_EXE; OwnerSid=$sid; OwnerResult=0 } + [pscustomobject]@{ ExecutablePath="$env:SystemRoot\System32\cmd.exe"; OwnerSid='S-1-0-0'; OwnerResult=0 } + [pscustomobject]@{ ExecutablePath="$PSHOME\powershell.exe"; OwnerSid=$sid; OwnerResult=2 } +} +function Invoke-CimMethod { + param($InputObject) + [pscustomobject]@{ ReturnValue=$InputObject.OwnerResult; Sid=$InputObject.OwnerSid } +} +function Get-Process { + Get-CimInstance | ForEach-Object { [pscustomobject]@{ Path=$_.ExecutablePath } } +} +"#; + let executable = std::env::current_exe().unwrap(); + let script = format!( + "{fixture}\n{}", + decode_windows_command(&windows_remote_binary_candidate_command()) + ); + let encoded = windows_powershell_script_command(&script); + let mut command = Command::new("powershell.exe"); + command + .args(["-NoProfile", "-NonInteractive", "-EncodedCommand"]) + .arg(encoded.split_whitespace().last().unwrap()) + .env("HERDR_TEST_EXE", &executable); + crate::platform::configure_background_command(&mut command); + let output = command.output().unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let paths = String::from_utf8(output.stdout) + .unwrap() + .lines() + .filter_map(|line| line.strip_prefix(WINDOWS_REMOTE_PATH_MARKER)) + .map(|path| decode_windows_remote_path(path).unwrap()) + .collect::<Vec<_>>(); + assert_eq!(paths, [executable.to_string_lossy().into_owned()]); + } + #[test] fn remote_output_framing_discards_any_banner_and_preserves_binary() { let payload = [0, 1, 2, 0xff, b'\n']; |