diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py index 71e630262608..f0bbabdb66ab 100644 --- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py +++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py @@ -6,7 +6,6 @@ import os import sys from pathlib import Path from subprocess import run -from tempfile import TemporaryDirectory from typing import assert_never from . import nix @@ -164,16 +163,11 @@ def execute(argv: list[str]) -> None: flake_build_flags = common_build_flags | vars(args_groups["flake_build_flags"]) copy_flags = common_flags | vars(args_groups["copy_flags"]) - # Will be cleaned up on exit automatically. - tmpdir = TemporaryDirectory(prefix="nixos-rebuild.") - tmpdir_path = Path(tmpdir.name) - atexit.register(cleanup_ssh, tmpdir_path) + atexit.register(cleanup_ssh) profile = Profile.from_arg(args.profile_name) - build_host = Remote.from_arg( - args.build_host, False, tmpdir_path, validate_opts=False - ) - target_host = Remote.from_arg(args.target_host, args.ask_sudo_password, tmpdir_path) + build_host = Remote.from_arg(args.build_host, False, validate_opts=False) + target_host = Remote.from_arg(args.target_host, args.ask_sudo_password) build_attr = BuildAttr.from_arg(args.attr, args.file) flake = Flake.from_arg(args.flake, target_host) action = Action(args.action) @@ -200,8 +194,7 @@ def execute(argv: list[str]) -> None: argv[0], new, ) - cleanup_ssh(tmpdir_path) - tmpdir.cleanup() + cleanup_ssh() os.execve(new, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"}) if args.upgrade or args.upgrade_all: 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 1c118d48b850..54200f14d85a 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 @@ -5,10 +5,22 @@ import subprocess from dataclasses import dataclass from getpass import getpass from pathlib import Path +from tempfile import TemporaryDirectory from typing import Self, Sequence, TypedDict, Unpack logger = logging.getLogger(__name__) +TMPDIR = TemporaryDirectory(prefix="nixos-rebuild.") +TMPDIR_PATH = Path(TMPDIR.name) +SSH_DEFAULT_OPTS = [ + "-o", + "ControlMaster=auto", + "-o", + f"ControlPath={TMPDIR_PATH / "ssh-%n"}", + "-o", + "ControlPersist=60", +] + @dataclass(frozen=True) class Remote: @@ -21,7 +33,6 @@ class Remote: cls, host: str | None, ask_sudo_password: bool | None, - tmp_dir: Path, validate_opts: bool = True, ) -> Self | None: if not host: @@ -30,15 +41,6 @@ class Remote: opts = os.getenv("NIX_SSHOPTS", "").split() if validate_opts: cls._validate_opts(opts, ask_sudo_password) - opts += [ - # SSH ControlMaster flags, allow for faster re-connection - "-o", - "ControlMaster=auto", - "-o", - f"ControlPath={tmp_dir / "ssh-%n"}", - "-o", - "ControlPersist=60", - ] sudo_password = None if ask_sudo_password: sudo_password = getpass(f"[sudo] password for {host}: ") @@ -66,12 +68,13 @@ class RunKwargs(TypedDict, total=False): stdout: int | None -def cleanup_ssh(tmp_dir: Path) -> None: +def cleanup_ssh() -> None: "Close SSH ControlMaster connection." - for ctrl in tmp_dir.glob("ssh-*"): + for ctrl in TMPDIR_PATH.glob("ssh-*"): subprocess.run( ["ssh", "-o", f"ControlPath={ctrl}", "-O", "exit", "dummyhost"], check=False ) + TMPDIR.cleanup() def run_wrapper( @@ -99,6 +102,7 @@ def run_wrapper( args = [ "ssh", *remote.opts, + *SSH_DEFAULT_OPTS, remote.host, "--", # sadly SSH just join all remaining parameters, expanding glob and 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 d647f78460b3..97ae0072a829 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 @@ -218,9 +218,9 @@ def test_execute_nix_switch_flake(mock_run: Any, tmp_path: Path) -> None: @patch.dict(nr.process.os.environ, {}, clear=True) @patch(get_qualified_name(nr.process.subprocess.run), autospec=True) -@patch(get_qualified_name(nr.TemporaryDirectory, nr)) # can't autospec +@patch(get_qualified_name(nr.cleanup_ssh, nr), autospec=True) def test_execute_nix_switch_flake_target_host( - mock_tmpdir: Any, + mock_cleanup_ssh: Any, mock_run: Any, tmp_path: Path, ) -> None: @@ -236,7 +236,6 @@ def test_execute_nix_switch_flake_target_host( # switch_to_configuration CompletedProcess([], 0), ] - mock_tmpdir.return_value.name = "/tmp/test" nr.execute( [ @@ -276,12 +275,7 @@ def test_execute_nix_switch_flake_target_host( call( [ "ssh", - "-o", - "ControlMaster=auto", - "-o", - "ControlPath=/tmp/test/ssh-%n", - "-o", - "ControlPersist=60", + *nr.process.SSH_DEFAULT_OPTS, "user@localhost", "--", f"sudo nix-env -p /nix/var/nix/profiles/system --set {config_path}", @@ -292,12 +286,7 @@ def test_execute_nix_switch_flake_target_host( call( [ "ssh", - "-o", - "ControlMaster=auto", - "-o", - "ControlPath=/tmp/test/ssh-%n", - "-o", - "ControlPersist=60", + *nr.process.SSH_DEFAULT_OPTS, "user@localhost", "--", f"sudo env NIXOS_INSTALL_BOOTLOADER=0 {config_path / 'bin/switch-to-configuration'} switch", @@ -311,9 +300,9 @@ def test_execute_nix_switch_flake_target_host( @patch.dict(nr.process.os.environ, {}, clear=True) @patch(get_qualified_name(nr.process.subprocess.run), autospec=True) -@patch(get_qualified_name(nr.TemporaryDirectory, nr)) # can't autospec +@patch(get_qualified_name(nr.cleanup_ssh, nr), autospec=True) def test_execute_nix_switch_flake_build_host( - mock_tmpdir: Any, + mock_cleanup_ssh: Any, mock_run: Any, tmp_path: Path, ) -> None: @@ -331,7 +320,6 @@ def test_execute_nix_switch_flake_build_host( # switch_to_configuration CompletedProcess([], 0), ] - mock_tmpdir.return_value.name = "/tmp/test" nr.execute( [ @@ -369,12 +357,7 @@ def test_execute_nix_switch_flake_build_host( call( [ "ssh", - "-o", - "ControlMaster=auto", - "-o", - "ControlPath=/tmp/test/ssh-%n", - "-o", - "ControlPersist=60", + *nr.process.SSH_DEFAULT_OPTS, "user@localhost", "--", f"nix --extra-experimental-features 'nix-command flakes' build '{config_path}^*' --print-out-paths", 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 998ad5dedc83..53ed6824e39f 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 @@ -1,6 +1,5 @@ -from pathlib import Path from typing import Any -from unittest.mock import call, patch +from unittest.mock import patch import nixos_rebuild.models as m import nixos_rebuild.process as p @@ -8,29 +7,7 @@ import nixos_rebuild.process as p from .helpers import get_qualified_name -@patch(get_qualified_name(p.subprocess.run)) -def test_cleanup_ssh(mock_run: Any, tmp_path: Path) -> None: - (tmp_path / "ssh-conn").touch() - - p.cleanup_ssh(tmp_path) - mock_run.assert_has_calls( - [ - call( - [ - "ssh", - "-o", - f"ControlPath={tmp_path}/ssh-conn", - "-O", - "exit", - "dummyhost", - ], - check=False, - ) - ] - ) - - -@patch(get_qualified_name(p.subprocess.run)) +@patch(get_qualified_name(p.subprocess.run), autospec=True) def test_run(mock_run: Any) -> None: p.run_wrapper(["test", "--with", "flags"], check=True) mock_run.assert_called_with( @@ -67,7 +44,15 @@ def test_run(mock_run: Any) -> None: remote=m.Remote("user@localhost", ["--ssh", "opt"], "password"), ) mock_run.assert_called_with( - ["ssh", "--ssh", "opt", "user@localhost", "--", "test --with 'some flags'"], + [ + "ssh", + "--ssh", + "opt", + *p.SSH_DEFAULT_OPTS, + "user@localhost", + "--", + "test --with 'some flags'", + ], check=True, text=True, errors="surrogateescape", @@ -87,6 +72,7 @@ def test_run(mock_run: Any) -> None: "ssh", "--ssh", "opt", + *p.SSH_DEFAULT_OPTS, "user@localhost", "--", "sudo --prompt= --stdin env FOO=bar test --with flags", @@ -99,38 +85,20 @@ def test_run(mock_run: Any) -> None: ) -def test_remote_from_name(monkeypatch: Any, tmpdir: Path) -> None: +def test_remote_from_name(monkeypatch: Any) -> None: monkeypatch.setenv("NIX_SSHOPTS", "") - assert m.Remote.from_arg("user@localhost", None, tmpdir) == m.Remote( + assert m.Remote.from_arg("user@localhost", None, False) == m.Remote( "user@localhost", - opts=[ - "-o", - "ControlMaster=auto", - "-o", - f"ControlPath={tmpdir / "ssh-%n"}", - "-o", - "ControlPersist=60", - ], + opts=[], sudo_password=None, ) # get_qualified_name doesn't work because getpass is aliased to another # function with patch(f"{p.__name__}.getpass", return_value="password"): - monkeypatch.setenv("NIX_SSHOPTS", "-f foo -b bar") - assert m.Remote.from_arg("user@localhost", True, tmpdir) == m.Remote( + monkeypatch.setenv("NIX_SSHOPTS", "-f foo -b bar -t") + assert m.Remote.from_arg("user@localhost", True, True) == m.Remote( "user@localhost", - opts=[ - "-f", - "foo", - "-b", - "bar", - "-o", - "ControlMaster=auto", - "-o", - f"ControlPath={tmpdir / "ssh-%n"}", - "-o", - "ControlPersist=60", - ], + opts=["-f", "foo", "-b", "bar", "-t"], sudo_password="password", )