fix(runtime): verify process ownership and path identity before terminating stored PIDs (#103)

backend/app/runtime/llama_server/llama_supervisor.py

@@ -95,20 +95,64 @@ def ensure_directories(self) -> None:
         self.llama_dir.mkdir(parents=True, exist_ok=True)
         self.models_dir.mkdir(parents=True, exist_ok=True)
 
-    def get_pid(self) -> Optional[int]:
-        if not self.pid_file.is_file():
-            return None
+    def _is_owned_executable(self, exe_path: Path) -> bool:
+        """Ensure the executable matches the expected llama-server binary or lives in self.llama_dir."""
         try:
-            pid = int(self.pid_file.read_text(encoding="utf-8").strip())
-            if sys.platform == "win32":
+            exe_res = exe_path.resolve()
+            expected_bin = self.get_binary_path()
+            if expected_bin and exe_res == expected_bin.resolve():
+                return True
+            llama_res = self.llama_dir.resolve()
+            if llama_res in exe_res.parents:
+                return "llama" in exe_res.name.lower()
+        except Exception:
+            pass
+        return False
+
+    def _verify_process_identity(self, pid: int) -> bool:
+        """Verify the process at pid is actually llama-server owned by this supervisor."""
+        if sys.platform == "win32":
+            try:
                 import ctypes
+                from ctypes import wintypes
                 kernel32 = ctypes.windll.kernel32
                 handle = kernel32.OpenProcess(0x1000, False, pid)
-                if handle:
+                if not handle:
+                    return False
+                try:
+                    buf = ctypes.create_unicode_buffer(1024)
+                    size = wintypes.DWORD(1024)
+                    success = kernel32.QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size))
+                    if success:
+                        exe_path = Path(buf.value)
+                        return self._is_owned_executable(exe_path)
+                    return False
+                finally:
                     kernel32.CloseHandle(handle)
-                    return pid
-            else:
-                os.kill(pid, 0)
+            except Exception:
+                return False
+        else:
+            try:
+                proc_exe = Path(f"/proc/{pid}/exe")
+                if proc_exe.exists():
+                    resolved = Path(os.readlink(proc_exe))
+                    if self._is_owned_executable(resolved):
+                        return True
+                proc_cmdline = Path(f"/proc/{pid}/cmdline")
+                if proc_cmdline.is_file():
+                    content = proc_cmdline.read_text(encoding="latin1", errors="ignore")
+                    if "llama-server" in content or str(self.llama_dir) in content:
+                        return True
+                return False
+            except (OSError, Exception):
+                return False
+
+    def get_pid(self) -> Optional[int]:
+        if not self.pid_file.is_file():
+            return None
+        try:
+            pid = int(self.pid_file.read_text(encoding="utf-8").strip())
+            if self._verify_process_identity(pid):
                 return pid
         except (ValueError, OSError):
             pass
@@ -266,11 +310,19 @@ def stop(self) -> Dict[str, Any]:
         if not pid:
             return {"success": True, "message": "llama-server is not running"}
 
+        if not self._verify_process_identity(pid):
+            self._clean_pid_file()
+            return {
+                "success": False,
+                "message": f"Process {pid} is not a verified llama-server process; refusing to terminate.",
+            }
+
         try:
             if sys.platform == "win32":
                 subprocess.run(["taskkill", "/F", "/T", "/PID", str(pid)], capture_output=True, check=False)
             else:
                 os.kill(pid, 15)
+            return {"success": True, "message": f"llama-server process {pid} terminated"}
         except Exception as e:
             return {"success": False, "message": f"Error terminating llama-server: {e}"}
         finally:

backend/app/runtime/supervisor.py

@@ -62,8 +62,28 @@ def has_isolated_env(self) -> bool:
         """Check if the sandboxed isolated virtual environment is ready."""
         return self.get_python_bin().is_file()
 
+    def _is_owned_executable(self, exe_path: Path) -> bool:
+        """Ensure the executable strictly resides in the isolated virtualenv or engine directory."""
+        try:
+            exe_res = exe_path.resolve()
+            expected_bin = self.get_python_bin().resolve()
+            if exe_res == expected_bin:
+                return True
+            runtime_res = self.runtime_dir.resolve()
+            if runtime_res in exe_res.parents:
+                return True
+            engine_res = self.engine_dir.resolve()
+            if engine_res in exe_res.parents:
+                return True
+            comfy_res = self.comfy_dir.resolve()
+            if comfy_res in exe_res.parents:
+                return True
+        except Exception:
+            pass
+        return False
+
     def _verify_process_identity(self, pid: int) -> bool:
-        """Verify the process at pid is actually python/comfy, avoiding recycled PID collision."""
+        """Verify the process at pid is actually owned by this supervisor, avoiding recycled PID collision."""
         if sys.platform == "win32":
             try:
                 import ctypes
@@ -72,25 +92,32 @@ def _verify_process_identity(self, pid: int) -> bool:
                 handle = kernel32.OpenProcess(0x1000, False, pid)  # PROCESS_QUERY_LIMITED_INFORMATION
                 if not handle:
                     return False
-                buf = ctypes.create_unicode_buffer(1024)
-                size = wintypes.DWORD(1024)
-                success = kernel32.QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size))
-                kernel32.CloseHandle(handle)
-                if success:
-                    exe_name = buf.value.lower()
-                    return "python" in exe_name
-                return False
+                try:
+                    buf = ctypes.create_unicode_buffer(1024)
+                    size = wintypes.DWORD(1024)
+                    success = kernel32.QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size))
+                    if success:
+                        exe_path = Path(buf.value)
+                        return self._is_owned_executable(exe_path)
+                    return False
+                finally:
+                    kernel32.CloseHandle(handle)
             except Exception:
                 return False
         else:
             try:
+                proc_exe = Path(f"/proc/{pid}/exe")
+                if proc_exe.exists():
+                    resolved = Path(os.readlink(proc_exe))
+                    if self._is_owned_executable(resolved):
+                        return True
                 proc_cmdline = Path(f"/proc/{pid}/cmdline")
                 if proc_cmdline.is_file():
-                    content = proc_cmdline.read_text(encoding="latin1", errors="ignore").lower()
-                    return "python" in content
-                os.kill(pid, 0)
-                return True
-            except OSError:
+                    content = proc_cmdline.read_text(encoding="latin1", errors="ignore")
+                    if str(self.comfy_dir) in content or str(self.runtime_dir) in content:
+                        return True
+                return False
+            except (OSError, Exception):
                 return False
 
     def get_pid(self) -> Optional[int]:

backend/app/runtime/webui_supervisor.py

@@ -64,8 +64,28 @@ def has_isolated_env(self) -> bool:
         """Check if the sandboxed isolated virtual environment is ready."""
         return self.get_python_bin().is_file()
 
+    def _is_owned_executable(self, exe_path: Path) -> bool:
+        """Ensure the executable strictly resides in the isolated virtualenv or engine directory."""
+        try:
+            exe_res = exe_path.resolve()
+            expected_bin = self.get_python_bin().resolve()
+            if exe_res == expected_bin:
+                return True
+            runtime_res = self.runtime_dir.resolve()
+            if runtime_res in exe_res.parents:
+                return True
+            engine_res = self.engine_dir.resolve()
+            if engine_res in exe_res.parents:
+                return True
+            webui_res = self.webui_dir.resolve()
+            if webui_res in exe_res.parents:
+                return True
+        except Exception:
+            pass
+        return False
+
     def _verify_process_identity(self, pid: int) -> bool:
-        """Verify the process at pid is actually python, avoiding recycled PID collision."""
+        """Verify the process at pid is actually owned by this supervisor, avoiding recycled PID collision."""
         if sys.platform == "win32":
             try:
                 import ctypes
@@ -74,25 +94,32 @@ def _verify_process_identity(self, pid: int) -> bool:
                 handle = kernel32.OpenProcess(0x1000, False, pid)  # PROCESS_QUERY_LIMITED_INFORMATION
                 if not handle:
                     return False
-                buf = ctypes.create_unicode_buffer(1024)
-                size = wintypes.DWORD(1024)
-                success = kernel32.QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size))
-                kernel32.CloseHandle(handle)
-                if success:
-                    exe_name = buf.value.lower()
-                    return "python" in exe_name
-                return False
+                try:
+                    buf = ctypes.create_unicode_buffer(1024)
+                    size = wintypes.DWORD(1024)
+                    success = kernel32.QueryFullProcessImageNameW(handle, 0, buf, ctypes.byref(size))
+                    if success:
+                        exe_path = Path(buf.value)
+                        return self._is_owned_executable(exe_path)
+                    return False
+                finally:
+                    kernel32.CloseHandle(handle)
             except Exception:
                 return False
         else:
             try:
+                proc_exe = Path(f"/proc/{pid}/exe")
+                if proc_exe.exists():
+                    resolved = Path(os.readlink(proc_exe))
+                    if self._is_owned_executable(resolved):
+                        return True
                 proc_cmdline = Path(f"/proc/{pid}/cmdline")
                 if proc_cmdline.is_file():
-                    content = proc_cmdline.read_text(encoding="latin1", errors="ignore").lower()
-                    return "python" in content
-                os.kill(pid, 0)
-                return True
-            except OSError:
+                    content = proc_cmdline.read_text(encoding="latin1", errors="ignore")
+                    if str(self.webui_dir) in content or str(self.runtime_dir) in content:
+                        return True
+                return False
+            except (OSError, Exception):
                 return False
 
     def get_pid(self) -> Optional[int]:

backend/tests/test_process_ownership_security.py

@@ -0,0 +1,73 @@
+import os
+from pathlib import Path
+import tempfile
+from unittest.mock import patch, MagicMock
+
+import pytest
+
+from app.runtime.supervisor import ComfySupervisor
+from app.runtime.webui_supervisor import WebUISupervisor
+from app.runtime.llama_server.llama_supervisor import LlamaServerSupervisor
+
+
+def test_comfy_rejects_unrelated_python_process():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        supervisor = ComfySupervisor(engine_dir=tmp_path / "comfyui")
+
+        # Current test runner process is Python, but outside the supervisor runtime directory
+        is_owned = supervisor._verify_process_identity(os.getpid())
+        assert is_owned is False, "ComfySupervisor must not claim ownership over an unrelated Python process"
+
+
+def test_webui_rejects_unrelated_python_process():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        supervisor = WebUISupervisor(engine_dir=tmp_path / "webui")
+
+        is_owned = supervisor._verify_process_identity(os.getpid())
+        assert is_owned is False, "WebUISupervisor must not claim ownership over an unrelated Python process"
+
+
+def test_llama_rejects_unrelated_python_process():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        supervisor = LlamaServerSupervisor(engine_dir=tmp_path / "llama")
+
+        is_owned = supervisor._verify_process_identity(os.getpid())
+        assert is_owned is False, "LlamaServerSupervisor must not claim ownership over an unrelated process"
+
+
+def test_stale_pid_cleanup_and_safe_stop():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        supervisor = ComfySupervisor(engine_dir=tmp_path / "comfyui")
+
+        # Simulate a stale PID file pointing to our own process
+        supervisor.ensure_directories()
+        supervisor.pid_file.write_text(str(os.getpid()), encoding="utf-8")
+
+        # get_pid() should detect that this process is unowned and clean up
+        assert supervisor.get_pid() is None
+        assert not supervisor.pid_file.is_file()
+
+        # stop() must NOT call taskkill
+        with patch("subprocess.run") as mock_run:
+            res = supervisor.stop()
+            assert res["success"] is True
+            assert mock_run.call_count == 0
+
+
+def test_owned_executable_matching():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        supervisor = ComfySupervisor(engine_dir=tmp_path / "comfyui")
+
+        runtime_python = supervisor.runtime_dir / "Scripts" / "python.exe"
+        assert supervisor._is_owned_executable(runtime_python) is True
+
+        comfy_script = supervisor.comfy_dir / "main.py"
+        assert supervisor._is_owned_executable(comfy_script) is True
+
+        system_python = Path(os.environ.get("WINDIR", "C:\\Windows")) / "System32" / "cmd.exe"
+        assert supervisor._is_owned_executable(system_python) is False