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

refs #3883

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

1 files changed, 109 insertions, 23 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()