fix: cache saved machine connection metadata (#4256)

refs #4252

2 files changed, 217 insertions, 14 deletions

diff --git a/src/remote/attach.rs b/src/remote/attach.rs
index 6a1ee6f9..53023b83 100644
--- a/src/remote/attach.rs
+++ b/src/remote/attach.rs
@@ -86,7 +86,10 @@ pub(crate) fn run_remote(remote: RemoteLaunch) -> io::Result<()> {
     run_client_process(&local_socket, &reattach_command, remote.keybindings)
 }
 
-pub(crate) fn prepare_saved_ssh(target: &str, session_name: &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)
@@ -135,7 +138,13 @@ pub(crate) fn prepare_saved_ssh(target: &str, session_name: &str) -> io::Result<
         )
         .is_none() =>
         {
-            Ok(())
+            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",
@@ -318,11 +327,21 @@ impl RemoteExecutable {
 pub(super) struct RemoteHerdr {
     install_suffix: String,
     executable: RemoteExecutable,
+    resolved_executable: Option<String>,
     platform: RemotePlatform,
     bridge_idle_timeout: bool,
 }
 
 impl RemoteHerdr {
+    pub(super) fn machine_metadata(&self) -> Option<crate::client::endpoint::SshMachineMetadata> {
+        let executable = self.resolved_executable.clone()?;
+        let metadata = crate::client::endpoint::SshMachineMetadata {
+            os: self.platform.os.to_owned(),
+            executable,
+        };
+        metadata.is_valid().then_some(metadata)
+    }
+
     fn for_platform(platform: RemotePlatform) -> Self {
         let (install_suffix, executable) = if platform.is_windows() {
             (
@@ -337,17 +356,20 @@ impl RemoteHerdr {
         Self {
             install_suffix,
             executable,
+            resolved_executable: None,
             platform,
             bridge_idle_timeout: false,
         }
     }
 
-    fn with_shell_path(mut self, shell_path: String) -> Self {
-        self.executable = RemoteExecutable::PosixShellPath(shell_path);
+    fn with_posix_path(mut self, path: &str) -> Self {
+        self.executable = RemoteExecutable::PosixShellPath(shell_quote(path));
+        self.resolved_executable = Some(path.to_owned());
         self
     }
 
     fn with_windows_path(mut self, path: String) -> Self {
+        self.resolved_executable = Some(path.clone());
         self.executable = RemoteExecutable::WindowsPath(path);
         self
     }
@@ -364,6 +386,10 @@ fn windows_powershell_application_script(path: &str, args: &[&str]) -> String {
 }
 
 fn windows_powershell_streaming_application_command(path: &str, args: &[&str]) -> String {
+    windows_powershell_script_command(&windows_powershell_streaming_application_script(path, args))
+}
+
+fn windows_powershell_streaming_application_script(path: &str, args: &[&str]) -> String {
     let command_line = args
         .iter()
         .map(|arg| crate::platform::quote_windows_command_line_arg(arg))
@@ -371,11 +397,11 @@ fn windows_powershell_streaming_application_command(path: &str, args: &[&str]) -
         .join(" ");
     // Start-Process -Wait waits for descendants, including a cold-started server.
     // Retain the handle so Windows PowerShell 5.1 keeps the application's exit code.
-    windows_powershell_script_command(&format!(
+    format!(
         "$process = Start-Process -FilePath {} -ArgumentList {} -NoNewWindow -PassThru -ErrorAction Stop; $null = $process.Handle; $process.WaitForExit(); exit $process.ExitCode",
         crate::platform::quote_powershell_arg(path),
         crate::platform::quote_powershell_arg(&command_line),
-    ))
+    )
 }
 
 fn posix_remote_output_command(command: &str) -> String {
@@ -1248,11 +1274,29 @@ fn prepare_windows_remote_herdr(
     })
 }
 
-pub(super) fn find_installed_remote_api_herdr(
+pub(super) fn discover_remote_api_metadata(
     ssh: &RemoteSsh,
     session: &str,
-) -> io::Result<RemoteHerdr> {
+) -> io::Result<crate::client::endpoint::SshMachineMetadata> {
     let platform = detect_remote_platform(ssh)?;
+    if !platform.is_windows() {
+        let output =
+            ssh.framed_user_shell_output(&posix_remote_api_discovery_command(&platform, session))?;
+        if !output.status.success() {
+            return Err(command_failed("remote binary discovery failed", &output));
+        }
+        let metadata = crate::client::endpoint::SshMachineMetadata {
+            os: platform.os.to_owned(),
+            executable: String::from_utf8_lossy(&output.stdout).trim().to_owned(),
+        };
+        if !metadata.is_valid() {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "invalid remote Herdr executable path",
+            ));
+        }
+        return Ok(metadata);
+    }
     let remote_herdr = RemoteHerdr::for_platform(platform);
     let candidates = remote_binary_candidates(ssh, &remote_herdr)?;
     for candidate in candidates {
@@ -1264,7 +1308,10 @@ pub(super) fn find_installed_remote_api_herdr(
         if probe.status.success()
             && String::from_utf8_lossy(&probe.stdout).trim() == "herdr-api-bridge-v1"
         {
-            return Ok(candidate);
+            return Ok(crate::client::endpoint::SshMachineMetadata {
+                os: "windows".into(),
+                executable: candidate.executable.display().to_owned(),
+            });
         }
     }
     Err(io::Error::new(
@@ -1511,7 +1558,7 @@ fn remote_herdr_from_path(remote_herdr: &RemoteHerdr, path: &str) -> Option<Remo
     if is_mise_shim_path(path) {
         return None;
     }
-    Some(remote_herdr.clone().with_shell_path(shell_quote(path)))
+    Some(remote_herdr.clone().with_posix_path(path))
 }
 
 fn is_mise_shim_path(path: &str) -> bool {
@@ -2383,6 +2430,67 @@ fn confirm_remote_install(
     Ok(())
 }
 
+fn posix_remote_api_discovery_command(platform: &RemotePlatform, session: &str) -> String {
+    let script = format!(
+        r#"set -f
+candidates=$(
+command -v herdr
+{discovery}
+)
+IFS='
+'
+for candidate in $candidates; do
+    case "$candidate" in
+        */mise/shims/herdr) continue ;;
+        /*) ;;
+        *) continue ;;
+    esac
+    [ -x "$candidate" ] || continue
+    if capability=$("$candidate" --session {session} remote-api-bridge --check </dev/null 2>/dev/null) && [ "$capability" = herdr-api-bridge-v1 ]; then
+        printf '%s\n' "$candidate"
+        exit 0
+    fi
+done
+printf '%s\n' 'remote Herdr does not support machine API forwarding; update Herdr on this machine' >&2
+exit 2"#,
+        discovery = known_remote_binary_candidate_script(platform),
+        session = shell_quote(session),
+    );
+    format!(
+        "/bin/sh -c {}",
+        shell_quote(&posix_remote_output_command(&script))
+    )
+}
+
+pub(super) const STALE_API_METADATA: &str = "herdr-machine-metadata-stale-v1";
+
+pub(super) fn cached_remote_api_command(
+    metadata: &crate::client::endpoint::SshMachineMetadata,
+    session: &str,
+) -> String {
+    if metadata.os == "windows" {
+        let path = crate::platform::quote_powershell_arg(&metadata.executable);
+        let session_arg = crate::platform::quote_powershell_arg(session);
+        let probe = format!(
+            "$capability = & {path} --session {session_arg} remote-api-bridge --check 2>$null; if ($LASTEXITCODE -ne 0 -or $capability -ne 'herdr-api-bridge-v1') {{ [Console]::Error.WriteLine('{STALE_API_METADATA}'); exit 78 }}; "
+        );
+        return windows_powershell_script_command(&format!(
+            "{probe}{}",
+            windows_powershell_streaming_application_script(
+                &metadata.executable,
+                &["--session", session, "remote-api-bridge"]
+            ),
+        ));
+    }
+    let path = shell_quote(&metadata.executable);
+    let session = shell_quote(session);
+    let script = format!(
+        "if capability=$({path} --session {session} remote-api-bridge --check </dev/null 2>/dev/null) && [ \"$capability\" = herdr-api-bridge-v1 ]; then\n{}\nelse\n    printf '%s\\n' '{STALE_API_METADATA}' >&2\n    exit 78\nfi",
+        posix_remote_output_command(&format!("exec {path} --session {session} remote-api-bridge")),
+    );
+    format!("/bin/sh -c {}", shell_quote(&script))
+}
+
 pub(super) fn remote_api_bridge_command(
     remote_herdr: &RemoteHerdr,
     session_name: &str,
@@ -4070,6 +4178,63 @@ mod tests {
     }
 
     #[test]
+    fn machine_metadata_keeps_raw_resolved_paths_not_shell_expressions() {
+        let remote = RemoteHerdr::for_platform(RemotePlatform {
+            os: "linux",
+            arch: "x86_64",
+        });
+        assert!(remote.machine_metadata().is_none());
+        let path = "/home/user's files/$literal/herdr";
+        let resolved = remote.with_posix_path(path);
+        assert_eq!(resolved.machine_metadata().unwrap().executable, path);
+        assert_eq!(resolved.executable.display(), shell_quote(path));
+        let remote = RemoteHerdr::for_platform(RemotePlatform {
+            os: "windows",
+            arch: "x86_64",
+        });
+        assert!(remote.machine_metadata().is_none());
+        let path = r"C:\Users\A B\herdr.exe";
+        assert_eq!(
+            remote
+                .with_windows_path(path.into())
+                .machine_metadata()
+                .unwrap()
+                .executable,
+            path
+        );
+    }
+
+    #[test]
+    fn cached_windows_api_command_checks_before_starting_the_stream() {
+        let path = r"C:\Users\A'B\herdr.exe";
+        let command = cached_remote_api_command(
+            &crate::client::endpoint::SshMachineMetadata {
+                os: "windows".into(),
+                executable: path.into(),
+            },
+            "fleet",
+        );
+        let encoded = command.split_whitespace().last().unwrap();
+        let bytes = base64::engine::general_purpose::STANDARD
+            .decode(encoded)
+            .unwrap();
+        let words = bytes
+            .chunks_exact(2)
+            .map(|chunk| u16::from_le_bytes([chunk[0], chunk[1]]))
+            .collect::<Vec<_>>();
+        let script = String::from_utf16(&words).unwrap();
+        assert!(script.contains(&crate::platform::quote_powershell_arg(path)));
+        assert!(script.contains(STALE_API_METADATA));
+        assert!(
+            script.find("remote-api-bridge --check").unwrap()
+                < script.find("Start-Process").unwrap()
+        );
+        assert!(script.contains("$LASTEXITCODE -ne 0"));
+        assert!(script.contains("-NoNewWindow -PassThru"));
+        assert!(script.contains("--session fleet remote-api-bridge"));
+    }
+
+    #[test]
     fn windows_remote_commands_use_one_encoded_powershell_grammar() {
         let executable = RemoteExecutable::WindowsPath("herdr.exe".to_string());
         let commands = [
diff --git a/src/remote/saved.rs b/src/remote/saved.rs
index 561ee1a0..4d8bd74b 100644
--- a/src/remote/saved.rs
+++ b/src/remote/saved.rs
@@ -19,6 +19,7 @@ pub(crate) fn connect_saved_ssh(
 ) -> io::Result<SavedSshStream> {
     let ssh = validated_saved_ssh(profile_id, target, session)?;
     let remote_herdr = find_installed_remote_herdr(&ssh)?;
+    let metadata = remote_herdr.machine_metadata();
     let path = saved_bridge_path(profile_id);
     let bridge = SshStdioBridge::start(
         target.to_owned(),
@@ -29,6 +30,10 @@ pub(crate) fn connect_saved_ssh(
         true,
     )?;
     let stream = crate::ipc::connect_local_stream(&path)?;
+    if let Some(metadata) = metadata {
+        crate::client::endpoint::SshMetadataCache::new(profile_id, target, session)?
+            .store(&metadata);
+    }
     Ok(SavedSshStream {
         stream,
         bridge: SavedSshBridge { _bridge: bridge },
@@ -38,13 +43,31 @@ pub(crate) fn connect_saved_ssh(
 pub(crate) struct SavedSshApiBridge {
     path: PathBuf,
     bridge: SshStdioBridge,
+    metadata_cache: crate::client::endpoint::SshMetadataCache,
+    pub(crate) used_cached_metadata: bool,
 }
 
 impl SavedSshApiBridge {
-    pub(crate) fn start(profile_id: &str, target: &str, session: &str) -> io::Result<Self> {
+    pub(crate) fn start(
+        profile_id: &str,
+        target: &str,
+        session: &str,
+        use_cached_metadata: bool,
+    ) -> io::Result<Self> {
         let ssh = validated_saved_ssh(profile_id, target, session)?;
-        let remote_herdr = super::attach::find_installed_remote_api_herdr(&ssh, session)?;
-        let command = super::attach::remote_api_bridge_command(&remote_herdr, session, false);
+        let metadata_cache =
+            crate::client::endpoint::SshMetadataCache::new(profile_id, target, session)?;
+        let cached = use_cached_metadata.then(|| metadata_cache.load()).flatten();
+        let used_cached_metadata = cached.is_some();
+        let metadata = match cached {
+            Some(metadata) => metadata,
+            None => {
+                let metadata = super::attach::discover_remote_api_metadata(&ssh, session)?;
+                metadata_cache.store(&metadata);
+                metadata
+            }
+        };
+        let command = super::attach::cached_remote_api_command(&metadata, session);
         let path = crate::platform::remote_bridge_endpoint_path(
             &format!("herdr-api-ssh-{}-{profile_id}.sock", std::process::id()),
             &format!(
@@ -60,7 +83,12 @@ impl SavedSshApiBridge {
             ssh.options(),
             true,
         )?;
-        Ok(Self { path, bridge })
+        Ok(Self {
+            path,
+            bridge,
+            metadata_cache,
+            used_cached_metadata,
+        })
     }
 
     pub(crate) fn socket_path(&self) -> &std::path::Path {
@@ -70,6 +98,16 @@ impl SavedSshApiBridge {
     pub(crate) fn reported_failure(&self) -> Option<io::Error> {
         self.bridge.reported_failure()
     }
+
+    pub(crate) fn invalidate_metadata(&self) {
+        self.metadata_cache.invalidate();
+    }
+
+    pub(crate) fn stale_metadata_failure(error: &io::Error) -> bool {
+        error
+            .to_string()
+            .contains(super::attach::STALE_API_METADATA)
+    }
 }
 
 pub(crate) fn saved_ssh_bootstrap_command(target: &str, session: &str) -> String {