From 752c092c47a8ed8d10dcfe4f1e24ab963cc59ae8 Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Sun, 1 Dec 2024 12:47:20 +0000
Subject: [PATCH] nixos-rebuild-ng: use shlex.quote instead of join in
 run_wrapper

---
 .../src/nixos_rebuild/process.py              |  9 +++++----
 .../nixos-rebuild-ng/src/tests/test_main.py   | 20 ++++++++++++++++---
 .../src/tests/test_process.py                 | 13 ++++++++++--
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py
index a9e8d7dff295..1232275d23c3 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py
@@ -105,10 +105,11 @@ def run_wrapper(
             *SSH_DEFAULT_OPTS,
             remote.host,
             "--",
-            # sadly SSH just join all remaining parameters, expanding glob and
-            # ignoring quotes
-            # so we need to use shlex.join() to safely join the arguments
-            shlex.join(str(a) for a in args),
+            # SSH will join the parameters here and pass it to the shell, so we
+            # need to quote it to avoid issues.
+            # We can't use `shlex.join`, otherwise we will hit MAX_ARG_STRLEN
+            # limits when the command becomes too big.
+            *[shlex.quote(str(a)) for a in args],
         ]
     else:
         if extra_env:
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py
index 97ae0072a829..9337fda44dc0 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py
@@ -278,7 +278,12 @@ def test_execute_nix_switch_flake_target_host(
                     *nr.process.SSH_DEFAULT_OPTS,
                     "user@localhost",
                     "--",
-                    f"sudo nix-env -p /nix/var/nix/profiles/system --set {config_path}",
+                    "sudo",
+                    "nix-env",
+                    "-p",
+                    "/nix/var/nix/profiles/system",
+                    "--set",
+                    str(config_path),
                 ],
                 check=True,
                 **DEFAULT_RUN_KWARGS,
@@ -289,7 +294,11 @@ def test_execute_nix_switch_flake_target_host(
                     *nr.process.SSH_DEFAULT_OPTS,
                     "user@localhost",
                     "--",
-                    f"sudo env NIXOS_INSTALL_BOOTLOADER=0 {config_path / 'bin/switch-to-configuration'} switch",
+                    "sudo",
+                    "env",
+                    "NIXOS_INSTALL_BOOTLOADER=0",
+                    f"{config_path / 'bin/switch-to-configuration'}",
+                    "switch",
                 ],
                 check=True,
                 **DEFAULT_RUN_KWARGS,
@@ -360,7 +369,12 @@ def test_execute_nix_switch_flake_build_host(
                     *nr.process.SSH_DEFAULT_OPTS,
                     "user@localhost",
                     "--",
-                    f"nix --extra-experimental-features 'nix-command flakes' build '{config_path}^*' --print-out-paths",
+                    "nix",
+                    "--extra-experimental-features",
+                    "'nix-command flakes'",
+                    "build",
+                    f"'{config_path}^*'",
+                    "--print-out-paths",
                 ],
                 check=True,
                 stdout=PIPE,
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_process.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_process.py
index 53ed6824e39f..6c2ca2289245 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_process.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_process.py
@@ -51,7 +51,9 @@ def test_run(mock_run: Any) -> None:
             *p.SSH_DEFAULT_OPTS,
             "user@localhost",
             "--",
-            "test --with 'some flags'",
+            "test",
+            "--with",
+            "'some flags'",
         ],
         check=True,
         text=True,
@@ -75,7 +77,14 @@ def test_run(mock_run: Any) -> None:
             *p.SSH_DEFAULT_OPTS,
             "user@localhost",
             "--",
-            "sudo --prompt= --stdin env FOO=bar test --with flags",
+            "sudo",
+            "--prompt=",
+            "--stdin",
+            "env",
+            "FOO=bar",
+            "test",
+            "--with",
+            "flags",
         ],
         check=True,
         env=None,