From 521843134977a335a638c0bc9f037205bd9fc49f Mon Sep 17 00:00:00 2001 From: tech Date: Thu, 21 May 2026 12:19:06 +0800 Subject: [PATCH] feat: add tunnel start/stop with PID file tracking Implements build_autossh_cmd, start_tunnel, and stop_tunnel in process.py. Adds _reap() helper to handle zombie processes in test environments where the parent process is long-lived (SIGKILL leaves zombies unless waitpid is called). Co-Authored-By: Claude Sonnet 4.6 --- src/autossh_mgr/process.py | 102 +++++++++++++++++++++++++++++++++++++ tests/unit/test_process.py | 85 +++++++++++++++++++++++++++++++ 2 files changed, 187 insertions(+) diff --git a/src/autossh_mgr/process.py b/src/autossh_mgr/process.py index a5d1486..52a0f99 100644 --- a/src/autossh_mgr/process.py +++ b/src/autossh_mgr/process.py @@ -1,5 +1,9 @@ import os +import shutil import signal +import socket +import subprocess +import time from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -7,6 +11,8 @@ from typing import Optional import click +from .config import TunnelConfig + @dataclass class TunnelStatus: @@ -87,3 +93,99 @@ def get_status(config_dir: Path, name: str) -> TunnelStatus: if not is_process_alive(pid): return TunnelStatus(state="stale", pid=pid) return TunnelStatus(state="running", pid=pid, uptime=format_uptime(started_at)) + + +def build_autossh_cmd(tunnel: TunnelConfig) -> list[str]: + cmd = [ + "autossh", "-M", "0", "-N", + "-o", "ServerAliveInterval=30", + "-o", "ServerAliveCountMax=3", + "-i", os.path.expanduser(tunnel.identity_file), + "-p", str(tunnel.port), + "-R", ( + f"{tunnel.remote_host}:{tunnel.remote_port}" + f":{tunnel.local_host}:{tunnel.local_port}" + ), + f"{tunnel.user}@{tunnel.host}", + ] + if tunnel.ssh_options: + for opt in tunnel.ssh_options.split(): + cmd.append(opt) + return cmd + + +def _is_port_in_use(port: int) -> bool: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + try: + s.bind(("127.0.0.1", port)) + return False + except OSError: + return True + + +def start_tunnel(config_dir: Path, tunnel: TunnelConfig) -> int: + if not shutil.which("autossh"): + raise click.ClickException("autossh not found. Install it first.") + if _is_port_in_use(tunnel.local_port): + raise click.ClickException(f"Port {tunnel.local_port} is already in use") + + log_path = config_dir / "logs" / f"{tunnel.name}.log" + env = os.environ.copy() + env["AUTOSSH_GATETIME"] = "0" + + with open(log_path, "a") as log: + proc = subprocess.Popen( + build_autossh_cmd(tunnel), + stdin=subprocess.DEVNULL, + stdout=log, + stderr=log, + env=env, + ) + + started_at = datetime.now(timezone.utc) + write_pid_file(config_dir, tunnel.name, proc.pid, started_at) + + time.sleep(1) + if not is_process_alive(proc.pid): + delete_pid_file(config_dir, tunnel.name) + raise click.ClickException( + f"autossh process for '{tunnel.name}' exited immediately. " + f"Check logs: {log_path}" + ) + return proc.pid + + +def _reap(pid: int) -> None: + """Best-effort waitpid to reap a zombie (only works if we are the parent).""" + try: + os.waitpid(pid, 0) + except (ChildProcessError, PermissionError, OSError): + pass + + +def stop_tunnel(config_dir: Path, name: str) -> None: + result = read_pid_file(config_dir, name) + if result is None: + return + pid, _ = result + if not is_process_alive(pid): + delete_pid_file(config_dir, name) + return + try: + os.kill(pid, signal.SIGTERM) + except ProcessLookupError: + delete_pid_file(config_dir, name) + return + for _ in range(50): + time.sleep(0.1) + if not is_process_alive(pid): + _reap(pid) + break + else: + try: + os.kill(pid, signal.SIGKILL) + except ProcessLookupError: + pass + else: + _reap(pid) + delete_pid_file(config_dir, name) diff --git a/tests/unit/test_process.py b/tests/unit/test_process.py index cd2b846..d9afaf6 100644 --- a/tests/unit/test_process.py +++ b/tests/unit/test_process.py @@ -101,3 +101,88 @@ def test_read_pid_file_corrupt(config_dir): pid_path.write_text("only-one-line\n") with pytest.raises(click.ClickException, match="Corrupt PID file"): read_pid_file(config_dir, "web-service") + + +# --------------------------------------------------------------------------- +# Task 4: build_autossh_cmd, start_tunnel, stop_tunnel +# --------------------------------------------------------------------------- + +import shutil +import signal +import subprocess +import time +from autossh_mgr.config import TunnelConfig +from autossh_mgr.process import start_tunnel, stop_tunnel, build_autossh_cmd + + +def test_build_autossh_cmd_basic(sample_tunnel): + cmd = build_autossh_cmd(sample_tunnel) + assert cmd[0] == "autossh" + assert "-M" in cmd and "0" in cmd + assert "-N" in cmd + assert "-i" in cmd + assert os.path.expanduser(sample_tunnel.identity_file) in cmd + assert "-p" in cmd and str(sample_tunnel.port) in cmd + assert "-R" in cmd + r_idx = cmd.index("-R") + assert cmd[r_idx + 1] == ( + f"{sample_tunnel.remote_host}:{sample_tunnel.remote_port}" + f":{sample_tunnel.local_host}:{sample_tunnel.local_port}" + ) + assert f"{sample_tunnel.user}@{sample_tunnel.host}" in cmd + + +def test_build_autossh_cmd_with_ssh_options(): + t = TunnelConfig( + name="test", host="h.com", user="u", + local_port=80, remote_port=8080, + ssh_options="-o StrictHostKeyChecking=no", + ) + cmd = build_autossh_cmd(t) + assert "-o" in cmd + assert "StrictHostKeyChecking=no" in cmd + + +def test_start_tunnel_no_autossh(config_dir, sample_tunnel, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda _: None) + import click + with pytest.raises(click.ClickException, match="autossh not found"): + start_tunnel(config_dir, sample_tunnel) + + +def test_start_creates_pid_file(config_dir, sample_tunnel, monkeypatch, tmp_path): + # Replace autossh with a real long-lived process for testing + fake = tmp_path / "autossh" + fake.write_text("#!/bin/sh\nexec sleep 60\n") + fake.chmod(0o755) + monkeypatch.setenv("PATH", f"{tmp_path}:{os.environ['PATH']}") + + pid = start_tunnel(config_dir, sample_tunnel) + try: + assert pid > 0 + result = read_pid_file(config_dir, "web-service") + assert result is not None + assert result[0] == pid + assert is_process_alive(pid) + finally: + os.kill(pid, signal.SIGKILL) + + +def test_stop_tunnel(config_dir, sample_tunnel, monkeypatch, tmp_path): + fake = tmp_path / "autossh" + fake.write_text("#!/bin/sh\nexec sleep 60\n") + fake.chmod(0o755) + monkeypatch.setenv("PATH", f"{tmp_path}:{os.environ['PATH']}") + + pid = start_tunnel(config_dir, sample_tunnel) + assert is_process_alive(pid) + stop_tunnel(config_dir, "web-service") + time.sleep(0.2) + assert not is_process_alive(pid) + assert read_pid_file(config_dir, "web-service") is None + + +def test_stop_tunnel_stale_pid(config_dir): + write_pid_file(config_dir, "web-service", 9999999, datetime.now(timezone.utc)) + stop_tunnel(config_dir, "web-service") # must not raise + assert read_pid_file(config_dir, "web-service") is None