nixos-rebuild-ng: cleanup SSH ControlMaster at exit

This commit is contained in:
Thiago Kenji Okada 2024-11-25 19:19:20 +00:00
parent f443299c58
commit e37e7e348d
2 changed files with 10 additions and 1 deletions

View File

@ -1,4 +1,5 @@
import argparse import argparse
import atexit
import json import json
import os import os
import sys import sys
@ -22,7 +23,7 @@ from .nix import (
switch_to_configuration, switch_to_configuration,
upgrade_channels, upgrade_channels,
) )
from .process import Remote from .process import Remote, cleanup_ssh
from .utils import flags_to_dict, info from .utils import flags_to_dict, info
VERBOSE = 0 VERBOSE = 0
@ -176,6 +177,7 @@ def execute(argv: list[str]) -> None:
# Will be cleaned up on exit automatically. # Will be cleaned up on exit automatically.
tmpdir = TemporaryDirectory(prefix="nixos-rebuild.") tmpdir = TemporaryDirectory(prefix="nixos-rebuild.")
tmpdir_path = Path(tmpdir.name) tmpdir_path = Path(tmpdir.name)
atexit.register(cleanup_ssh, tmpdir_path)
profile = Profile.from_name(args.profile_name) profile = Profile.from_name(args.profile_name)
target_host = Remote.from_arg(args.target_host, not args.no_ssh_tty, tmpdir_path) target_host = Remote.from_arg(args.target_host, not args.no_ssh_tty, tmpdir_path)

View File

@ -26,6 +26,7 @@ class Remote:
def from_arg(cls, host: str | None, tty: bool | None, tmp_dir: Path) -> Self | None: def from_arg(cls, host: str | None, tty: bool | None, tmp_dir: Path) -> Self | None:
if host: if host:
opts = os.getenv("NIX_SSHOPTS", "").split() + [ opts = os.getenv("NIX_SSHOPTS", "").split() + [
# SSH ControlMaster flags, allow for faster re-connection
"-o", "-o",
"ControlMaster=auto", "ControlMaster=auto",
"-o", "-o",
@ -38,6 +39,12 @@ class Remote:
return None return None
def cleanup_ssh(tmp_dir: Path) -> None:
"Close SSH ControlMaster connection."
for ctrl in tmp_dir.glob("ssh-*"):
subprocess.run(["ssh", "-o", f"ControlPath={ctrl}", "exit"], check=False)
def run_wrapper( def run_wrapper(
args: Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]], args: Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]],
*, *,