feat: manage multiple ssh machines from one client (#3670)

1 files changed, 126 insertions, 0 deletions

diff --git a/src/remote/args.rs b/src/remote/args.rs
new file mode 100644
index 00000000..476acf70
--- /dev/null
+++ b/src/remote/args.rs
@@ -0,0 +1,126 @@
+pub(crate) const REATTACH_COMMAND_ENV_VAR: &str = "HERDR_REATTACH_COMMAND";
+pub(crate) const REMOTE_KEYBINDINGS_ENV_VAR: &str = "HERDR_REMOTE_KEYBINDINGS";
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum RemoteKeybindings {
+    Local,
+    Server,
+}
+
+impl RemoteKeybindings {
+    pub(super) fn parse(value: &str) -> Result<Self, String> {
+        match value {
+            "local" => Ok(Self::Local),
+            "server" => Ok(Self::Server),
+            _ => Err("--remote-keybindings must be 'local' or 'server'".to_string()),
+        }
+    }
+
+    pub(super) fn as_str(self) -> &'static str {
+        match self {
+            Self::Local => "local",
+            Self::Server => "server",
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct RemoteLaunch {
+    pub(crate) target: String,
+    pub(crate) keybindings: RemoteKeybindings,
+    pub(crate) live_handoff: bool,
+}
+
+pub(crate) fn extract_remote_args(
+    args: &[String],
+) -> Result<(Vec<String>, Option<RemoteLaunch>), String> {
+    let mut cleaned = Vec::with_capacity(args.len());
+    if let Some(program) = args.first() {
+        cleaned.push(program.clone());
+    }
+
+    let mut remote_target = None;
+    let mut keybindings = RemoteKeybindings::Local;
+    let mut keybindings_seen = false;
+    let mut live_handoff = false;
+    let mut index = 1;
+    while index < args.len() {
+        let arg = &args[index];
+        if arg == "--" {
+            cleaned.extend_from_slice(&args[index..]);
+            break;
+        }
+        if arg == "--handoff" {
+            live_handoff = true;
+            index += 1;
+            continue;
+        }
+        if arg == "--remote" {
+            if remote_target.is_some() {
+                return Err("--remote can only be specified once".to_string());
+            }
+            let Some(value) = args.get(index + 1) else {
+                return Err("missing value for --remote".to_string());
+            };
+            remote_target = Some(validate_remote_target(value)?.to_owned());
+            index += 2;
+            continue;
+        }
+        if let Some(value) = arg.strip_prefix("--remote=") {
+            if remote_target.is_some() {
+                return Err("--remote can only be specified once".to_string());
+            }
+            remote_target = Some(validate_remote_target(value)?.to_owned());
+            index += 1;
+            continue;
+        }
+        if arg == "--remote-keybindings" {
+            if keybindings_seen {
+                return Err("--remote-keybindings can only be specified once".to_string());
+            }
+            let Some(value) = args.get(index + 1) else {
+                return Err("missing value for --remote-keybindings".to_string());
+            };
+            keybindings = RemoteKeybindings::parse(value)?;
+            keybindings_seen = true;
+            index += 2;
+            continue;
+        }
+        if let Some(value) = arg.strip_prefix("--remote-keybindings=") {
+            if keybindings_seen {
+                return Err("--remote-keybindings can only be specified once".to_string());
+            }
+            keybindings = RemoteKeybindings::parse(value)?;
+            keybindings_seen = true;
+            index += 1;
+            continue;
+        }
+
+        cleaned.push(arg.clone());
+        index += 1;
+    }
+
+    let remote = remote_target.map(|target| RemoteLaunch {
+        target,
+        keybindings,
+        live_handoff,
+    });
+    if remote.is_none() && keybindings_seen {
+        return Err("--remote-keybindings requires --remote".to_string());
+    }
+    if remote.is_none() && live_handoff {
+        cleaned.push("--handoff".to_string());
+    }
+
+    Ok((cleaned, remote))
+}
+
+pub(crate) fn validate_remote_target(target: &str) -> Result<&str, String> {
+    if target.is_empty() {
+        return Err("missing value for --remote".to_string());
+    }
+    if target.starts_with('-') {
+        return Err("--remote target must not start with '-'".to_string());
+    }
+    Ok(target)
+}