feat: enhance service discovery and proxy management

backend/app/services/mitmproxy/manager.py

@@ -20,6 +20,7 @@ class MitmdumpManager:
         self.work_dir = Path(work_dir)
         self.command = command
         self.processes: dict[int, asyncio.subprocess.Process] = {}
+        self.attached_pids: dict[int, int] = {}
         self.locks: dict[int, asyncio.Lock] = {}
 
     async def start(self, proxy: ProxyDocument) -> dict[str, object]:
@@ -67,6 +68,7 @@ class MitmdumpManager:
             raise MitmdumpError(reason)
 
         self.processes[proxy.proxy_port] = process
+        self.attached_pids.pop(proxy.proxy_port, None)
         return {
             "status": "running",
             "pid": process.pid,
@@ -79,6 +81,7 @@ class MitmdumpManager:
 
     async def stop_unlocked(self, proxy: ProxyDocument) -> dict[str, object]:
         process = self.processes.pop(proxy.proxy_port, None)
+        self.attached_pids.pop(proxy.proxy_port, None)
         if process is not None and process.returncode is None:
             await self.stop_process(process)
         elif proxy.pid is not None:
@@ -98,19 +101,70 @@ class MitmdumpManager:
             if process.returncode is None:
                 return "running"
             self.processes.pop(proxy.proxy_port, None)
+            self.attached_pids.pop(proxy.proxy_port, None)
             return "stopped"
 
-        if proxy.pid is None:
+        pid = self.attached_pids.get(proxy.proxy_port, proxy.pid)
+        if pid is None:
             return "stopped" if proxy.status == "running" else proxy.status
 
-        try:
-            os.kill(proxy.pid, 0)
-        except ProcessLookupError:
+        if not self.is_process_alive(pid):
+            self.attached_pids.pop(proxy.proxy_port, None)
+            return "stopped"
+        if not self.matches_proxy_process(proxy, pid):
+            self.attached_pids.pop(proxy.proxy_port, None)
             return "stopped"
-        except PermissionError:
-            return "running"
         return "running"
 
+    async def reconcile(self, proxy: ProxyDocument) -> dict[str, object] | None:
+        if not self.needs_recovery(proxy):
+            return None
+
+        if proxy.pid is not None and self.is_process_alive(proxy.pid):
+            if self.matches_proxy_process(proxy, proxy.pid):
+                self.reattach(proxy)
+                return {
+                    "status": "running",
+                    "pid": proxy.pid,
+                    "script_path": str(self.script_path_for(proxy)),
+                    "config_path": str(self.config_path_for(proxy)),
+                    "log_path": str(self.log_path_for(proxy)),
+                    "last_error": None,
+                }
+
+            self.attached_pids.pop(proxy.proxy_port, None)
+            return {
+                "status": "error",
+                "pid": None,
+                "last_error": (
+                    f"Found a live pid {proxy.pid}, but its command does not match "
+                    "the expected mitmdump runtime for this proxy."
+                ),
+            }
+
+        if self.needs_upstream_restore(proxy):
+            restored = await self.restore_upstream(proxy, proxy.service.upstream_port)
+            if not restored:
+                return {
+                    "status": "error",
+                    "pid": None,
+                    "last_error": (
+                        "Proxy process is gone and the upstream service port "
+                        "could not be restored automatically."
+                    ),
+                }
+            return {
+                "status": "stopped",
+                "pid": None,
+                "service.upstream_port": None,
+                "last_error": None,
+            }
+
+        if proxy.status == "running" or proxy.pid is not None:
+            return {"status": "stopped", "pid": None, "last_error": None}
+
+        return None
+
     async def prepare_upstream(self, proxy: ProxyDocument) -> int | None:
         if proxy.mode_spec is not None or proxy.mode != "reverse":
             return None
@@ -126,6 +180,8 @@ class MitmdumpManager:
             Path(proxy.service.folder),
             proxy.service.port,
             upstream_port,
+            host_to="127.0.0.1",
+            rewrite_host=True,
         )
         if not changed:
             raise MitmdumpError("Unable to move Docker service port")
@@ -135,19 +191,21 @@ class MitmdumpManager:
         self,
         proxy: ProxyDocument,
         upstream_port: int | None,
-    ) -> None:
+    ) -> bool:
         if upstream_port is None:
-            return
+            return True
         if proxy.service.name is None or proxy.service.folder is None:
-            return
+            return True
         if upstream_port == proxy.service.port:
-            return
+            return True
 
-        await docker.change_port(
+        return await docker.change_port(
             proxy.service.name,
             Path(proxy.service.folder),
             upstream_port,
             proxy.service.port,
+            host_to=proxy.service.bind_host,
+            rewrite_host=True,
         )
 
     def build_command(
@@ -200,20 +258,20 @@ class MitmdumpManager:
         return proxy.mode in {"regular", "socks5", "dns", "transparent"}
 
     def write_script(self, proxy: ProxyDocument) -> Path:
-        path = self.work_dir / "addons" / f"{proxy.proxy_port}.py"
+        path = self.script_path_for(proxy)
         path.parent.mkdir(parents=True, exist_ok=True)
         path.write_text(proxy.code, encoding="utf-8")
         return path
 
     def write_config(self, proxy: ProxyDocument) -> Path:
-        path = self.work_dir / "conf" / str(proxy.proxy_port) / "config.yaml"
+        path = self.config_path_for(proxy)
         path.parent.mkdir(parents=True, exist_ok=True)
         options = {"block_global": False, **proxy.mitm_options}
         path.write_text(yaml.safe_dump(options, sort_keys=True), encoding="utf-8")
         return path
 
     def log_path(self, proxy: ProxyDocument) -> Path:
-        path = self.work_dir / "logs" / f"{proxy.proxy_port}.log"
+        path = self.log_path_for(proxy)
         path.parent.mkdir(parents=True, exist_ok=True)
         return path
 
@@ -227,19 +285,7 @@ class MitmdumpManager:
             )
 
     async def stop_process(self, process) -> None:
-        try:
-            os.killpg(process.pid, signal.SIGTERM)
-        except ProcessLookupError:
-            return
-
-        try:
-            await asyncio.wait_for(process.wait(), timeout=5)
-        except TimeoutError:
-            try:
-                os.killpg(process.pid, signal.SIGKILL)
-            except ProcessLookupError:
-                return
-            await process.wait()
+        await self.stop_pid(process.pid)
 
     async def exited_early(self, process) -> bool:
         try:
@@ -250,21 +296,94 @@ class MitmdumpManager:
 
     async def stop_pid(self, pid: int) -> None:
         try:
-            os.kill(pid, signal.SIGTERM)
+            pgid = os.getpgid(pid)
         except ProcessLookupError:
             return
+        except PermissionError:
+            pgid = pid
 
-        for _ in range(10):
-            await asyncio.sleep(0.2)
+        try:
+            os.killpg(pgid, signal.SIGTERM)
+        except ProcessLookupError:
+            return
+        except PermissionError:
             try:
-                os.kill(pid, 0)
+                os.kill(pid, signal.SIGTERM)
             except ProcessLookupError:
                 return
 
+        for _ in range(10):
+            await asyncio.sleep(0.2)
+            if not self.is_process_alive(pid):
+                return
+
         try:
-            os.kill(pid, signal.SIGKILL)
+            os.killpg(pgid, signal.SIGKILL)
         except ProcessLookupError:
-            return
+            pass
+        except PermissionError:
+            try:
+                os.kill(pid, signal.SIGKILL)
+            except ProcessLookupError:
+                return
+
+    def reattach(self, proxy: ProxyDocument) -> None:
+        if proxy.pid is not None:
+            self.attached_pids[proxy.proxy_port] = proxy.pid
+
+    def needs_recovery(self, proxy: ProxyDocument) -> bool:
+        return (
+            proxy.status == "running"
+            or proxy.pid is not None
+            or self.needs_upstream_restore(proxy)
+        )
+
+    def needs_upstream_restore(self, proxy: ProxyDocument) -> bool:
+        upstream_port = proxy.service.upstream_port
+        return upstream_port is not None and upstream_port != proxy.service.port
+
+    def is_process_alive(self, pid: int) -> bool:
+        try:
+            os.kill(pid, 0)
+        except ProcessLookupError:
+            return False
+        except PermissionError:
+            return True
+        return True
+
+    def matches_proxy_process(self, proxy: ProxyDocument, pid: int) -> bool:
+        cmdline = self.read_cmdline(pid)
+        if not cmdline:
+            return False
+
+        expected_command = Path(self.command).name
+        if not any(Path(part).name == expected_command for part in cmdline):
+            return False
+
+        expected_script = str(self.script_path_for(proxy))
+        expected_confdir = f"confdir={self.conf_dir_for(proxy)}"
+        return expected_script in cmdline and expected_confdir in cmdline
+
+    def read_cmdline(self, pid: int) -> list[str]:
+        path = Path(f"/proc/{pid}/cmdline")
+        try:
+            raw = path.read_bytes()
+        except OSError:
+            return []
+        decoded = raw.decode("utf-8", errors="replace")
+        return [part for part in decoded.split("\x00") if part]
+
+    def script_path_for(self, proxy: ProxyDocument) -> Path:
+        return self.work_dir / "addons" / f"{proxy.proxy_port}.py"
+
+    def conf_dir_for(self, proxy: ProxyDocument) -> Path:
+        return self.work_dir / "conf" / str(proxy.proxy_port)
+
+    def config_path_for(self, proxy: ProxyDocument) -> Path:
+        return self.conf_dir_for(proxy) / "config.yaml"
+
+    def log_path_for(self, proxy: ProxyDocument) -> Path:
+        return self.work_dir / "logs" / f"{proxy.proxy_port}.log"
 
     def read_log_tail(self, path: Path, limit: int = 4000) -> str:
         if not path.exists():