fix: accept machine add options before the SSH target (#3897)

refs #3883

Co-authored-by: akbash-bot <300245827+akbash-bot@users.noreply.github.com>

2 files changed, 163 insertions, 24 deletions

diff --git a/src/cli/machine.rs b/src/cli/machine.rs
index 3e93d101..ab84041e 100644
--- a/src/cli/machine.rs
+++ b/src/cli/machine.rs
@@ -89,62 +89,84 @@ fn list(args: &[String]) -> std::io::Result<i32> {
     Ok(0)
 }
 
-fn add(args: &[String]) -> std::io::Result<i32> {
+#[derive(Debug, PartialEq, Eq)]
+struct AddArgs {
+    target: String,
+    label: String,
+    session: String,
+}
+
+fn parse_add_args(args: &[String]) -> Result<AddArgs, String> {
     let args = super::expand_equals_args(args, &["--label", "--remote-session"]);
-    let Some(target) = args.first().filter(|value| !value.starts_with('-')) else {
-        eprintln!(
-            "usage: herdr machine add <ssh-target> --label <label> [--remote-session <name>]"
-        );
-        return Ok(2);
-    };
+    let mut target = None;
     let mut label = None;
     let mut session = None;
-    let mut index = 1;
+    let mut index = 0;
     while index < args.len() {
         let (name, value) = match args[index].as_str() {
             "--label" | "--remote-session" => {
                 let Some(value) = args.get(index + 1) else {
-                    eprintln!("missing value for {}", args[index]);
-                    return Ok(2);
+                    return Err(format!("missing value for {}", args[index]));
                 };
                 index += 2;
                 (args[index - 2].as_str(), value.clone())
             }
+            positional if !positional.starts_with('-') && target.is_none() => {
+                target = Some(positional.to_owned());
+                index += 1;
+                continue;
+            }
             unknown => {
-                eprintln!("unknown machine add option: {unknown}");
-                return Ok(2);
+                return Err(format!("unknown machine add option: {unknown}"));
             }
         };
         match name {
             "--label" if label.is_none() => label = Some(value),
             "--remote-session" if session.is_none() => session = Some(value),
             "--remote-session" => {
-                eprintln!("--remote-session can only be specified once");
-                return Ok(2);
+                return Err("--remote-session can only be specified once".into());
             }
             "--label" => {
-                eprintln!("--label can only be specified once");
-                return Ok(2);
+                return Err("--label can only be specified once".into());
             }
             _ => unreachable!("validated machine add option"),
         }
     }
-    let Some(label) = label else {
-        eprintln!("--label is required");
-        return Ok(2);
-    };
+    let target = target.ok_or_else(|| {
+        "usage: herdr machine add <ssh-target> --label <label> [--remote-session <name>]".to_owned()
+    })?;
+    let label = label.ok_or_else(|| "--label is required".to_owned())?;
     let session = session.unwrap_or_else(|| crate::session::DEFAULT_SESSION_NAME.to_owned());
+    Ok(AddArgs {
+        target,
+        label,
+        session,
+    })
+}
+
+fn add(args: &[String]) -> std::io::Result<i32> {
+    let AddArgs {
+        target,
+        label,
+        session,
+    } = match parse_add_args(args) {
+        Ok(args) => args,
+        Err(error) => {
+            eprintln!("{error}");
+            return Ok(2);
+        }
+    };
     let mut catalog = load_catalog()?;
-    match catalog.add_ssh(label.clone(), target, session.clone()) {
+    match catalog.add_ssh(label.clone(), &target, session.clone()) {
         Ok(_) => {}
         Err(error) => {
             eprintln!("error: {error}");
             return Ok(2);
         }
     }
-    if let Err(error) = crate::remote::prepare_saved_ssh(target, &session) {
+    if let Err(error) = crate::remote::prepare_saved_ssh(&target, &session) {
         eprintln!("error: {error}; machine was not saved");
-        crate::remote::print_saved_ssh_error_hint(&error, target);
+        crate::remote::print_saved_ssh_error_hint(&error, &target);
         return Ok(1);
     }
     // Setup can wait for human approval. Do not overwrite catalog edits made meanwhile.
@@ -272,6 +294,70 @@ mod tests {
     use super::*;
 
     #[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![
+                    "--remote-session",
+                    "agents",
+                    "workstation.coder",
+                    "--label",
+                    "coder",
+                ],
+                "agents",
+            ),
+            (
+                vec![
+                    "--label=coder",
+                    "--remote-session=agents",
+                    "workstation.coder",
+                ],
+                "agents",
+            ),
+        ] {
+            let args = args.into_iter().map(str::to_owned).collect::<Vec<_>>();
+            assert_eq!(
+                parse_add_args(&args).unwrap(),
+                AddArgs {
+                    target: "workstation.coder".into(),
+                    label: "coder".into(),
+                    session: session.into(),
+                },
+                "{args:?}"
+            );
+        }
+    }
+
+    #[test]
+    fn add_parser_rejects_incomplete_duplicate_and_extra_arguments() {
+        for args in [
+            vec![],
+            vec!["--label", "coder"],
+            vec!["workstation.coder"],
+            vec!["workstation.coder", "--label"],
+            vec!["workstation.coder", "--label", "coder", "--remote-session"],
+            vec!["--label", "coder", "--label", "other", "workstation.coder"],
+            vec![
+                "workstation.coder",
+                "--label",
+                "coder",
+                "--remote-session",
+                "a",
+                "--remote-session",
+                "b",
+            ],
+            vec!["--label", "coder", "workstation.coder", "other-host"],
+            vec!["--unknown", "workstation.coder", "--label", "coder"],
+            vec!["--label", "--remote-session", "agents", "workstation.coder"],
+        ] {
+            let args = args.into_iter().map(str::to_owned).collect::<Vec<_>>();
+            assert!(parse_add_args(&args).is_err(), "{args:?}");
+        }
+    }
+
+    #[test]
     fn profile_id_parser_rejects_target_text() {
         assert!(one_profile_id(&["build.example".into()], "usage")
             .unwrap()
diff --git a/tests/machine_setup.rs b/tests/machine_setup.rs
index 5fce2865..3b577f4e 100644
--- a/tests/machine_setup.rs
+++ b/tests/machine_setup.rs
@@ -5,7 +5,7 @@ use std::io::{Read, Write};
 use std::os::unix::fs::PermissionsExt;
 use std::process::Command;
 use std::sync::mpsc;
-use std::time::Duration;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use portable_pty::{native_pty_system, CommandBuilder, PtySize};
 
@@ -206,6 +206,59 @@ fn setup_with_strict_host_key_failure(
 }
 
 #[test]
+fn machine_add_accepts_help_argument_order() {
+    let nonce = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_nanos();
+    let root = std::env::temp_dir().join(format!(
+        "herdr-machine-add-order-{}-{nonce}",
+        std::process::id()
+    ));
+    fs::create_dir(&root).unwrap();
+    let app = if cfg!(debug_assertions) {
+        "herdr-dev"
+    } else {
+        "herdr"
+    };
+    fs::create_dir_all(root.join("config").join(app)).unwrap();
+    fs::write(
+        root.join("config").join(app).join("config.toml"),
+        "onboarding = false\n[remote]\nmanage_ssh_config = false\n",
+    )
+    .unwrap();
+    let mut command = Command::new(env!("CARGO_BIN_EXE_herdr"));
+    command.args(["machine", "add", "--label", "coder", "workstation.coder"]);
+    // Reach remote preparation, but never execute SSH or start a server.
+    command.env("PATH", root.join("no-executables"));
+    command.env("HOME", &root);
+    command.env("XDG_CONFIG_HOME", root.join("config"));
+    command.env("XDG_STATE_HOME", root.join("state"));
+    command.env("XDG_RUNTIME_DIR", &root);
+    for name in [
+        "HERDR_ENV",
+        "HERDR_SESSION",
+        "HERDR_SOCKET_PATH",
+        "HERDR_CLIENT_SOCKET_PATH",
+        "HERDR_REMOTE_BINARY",
+        "HERDR_CONFIG_PATH",
+    ] {
+        command.env_remove(name);
+    }
+    let output = command.output().unwrap();
+    let saved = root
+        .join("state")
+        .join(app)
+        .join("client/endpoints.json")
+        .exists();
+    fs::remove_dir_all(root).unwrap();
+    let stderr = String::from_utf8_lossy(&output.stderr);
+    assert_eq!(output.status.code(), Some(1), "{stderr}");
+    assert!(stderr.contains("machine was not saved"), "{stderr}");
+    assert!(!saved, "failed preparation must not save a machine");
+}
+
+#[test]
 fn machine_add_reports_strict_host_key_failure() {
     let result = setup_with_strict_host_key_failure("new", "", false, true);
     assert!(!result.success, "{}", result.output);