feat: add tests for known hosts management and enhance validation in known_hosts.py

tests/test_config.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import os
 import unittest
+from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import patch
 
@@ -68,3 +69,22 @@ class ArchitectureSettingsTests(unittest.TestCase):
             self.assertRaisesRegex(ConfigurationError, "unsupported host architecture: riscv64"),
         ):
             Settings.from_environment()
+
+
+class KnownHostsSettingsTests(unittest.TestCase):
+    def test_known_hosts_files_are_loaded_as_explicit_absolute_paths(self) -> None:
+        with patch.dict(os.environ, {
+            "UVM_KNOWN_HOSTS_FILES": "/root/.ssh/known_hosts:/home/operator/.ssh/known_hosts",
+        }, clear=True):
+            settings = Settings.from_environment()
+        self.assertEqual(settings.known_hosts_files, (
+            Path("/root/.ssh/known_hosts"), Path("/home/operator/.ssh/known_hosts"),
+        ))
+        with patch.dict(os.environ, {}, clear=True):
+            self.assertIsNone(Settings.from_environment().known_hosts_files)
+
+    def test_relative_or_empty_path_entries_are_rejected(self) -> None:
+        for value in ("known_hosts", "/tmp/known_hosts:", ":/tmp/known_hosts"):
+            with self.subTest(value=value), patch.dict(os.environ, {"UVM_KNOWN_HOSTS_FILES": value}, clear=True):
+                with self.assertRaisesRegex(ConfigurationError, "UVM_KNOWN_HOSTS_FILES"):
+                    Settings.from_environment()

tests/test_known_hosts.py

@@ -0,0 +1,277 @@
+from __future__ import annotations
+
+import os
+import errno
+import shutil
+import tempfile
+import unittest
+from dataclasses import replace
+from ipaddress import IPv4Address
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from tests.test_cli import make_vm
+from tests.test_lifecycle import LifecycleTestCase
+from uvm.config import Settings
+from uvm.domain import PortRule
+from uvm.errors import FirecrackerError, UvmError, ValidationError
+from uvm.known_hosts import KnownHostsManager, _without_hosts
+from uvm.system import CommandRunner
+
+
+class KnownHostsTests(unittest.TestCase):
+    def setUp(self) -> None:
+        temporary = tempfile.TemporaryDirectory()
+        self.addCleanup(temporary.cleanup)
+        self.directory = Path(temporary.name)
+        self.path = self.directory / "known_hosts"
+        self.vm = replace(make_vm(), ports=(PortRule(22), PortRule(2222)))
+        self.manager = KnownHostsManager(Settings(known_hosts_files=(self.path,)))
+
+    def test_removes_all_vm_key_types_and_aliases_but_preserves_other_hosts(self) -> None:
+        self.path.write_bytes(
+            b"# Known hosts\r\n"
+            b"10.42.0.2 ssh-ed25519 AAAAed first key\n"
+            b"10.42.0.2 ecdsa-sha2-nistp256 AAAAecdsa second key\n"
+            b"10.42.0.2 ssh-rsa AAAArsa third key\n"
+            b"uvm-vm-test ssh-ed25519 AAAAalias\n"
+            b"[10.42.0.2]:22 ssh-ed25519 AAAAport22\n"
+            b"[10.42.0.2]:2222 ssh-ed25519 AAAAport2222\n"
+            b"[10.42.0.2]:8022 ssh-ed25519 AAAAold-port\n"
+            b"[uvm-vm-test]:2222 ssh-ed25519 AAAAportalias\n"
+            b"10.42.0.20 ssh-ed25519 AAAAother same prefix\n"
+            b"some-host,10.42.0.2,another-host  ssh-ed25519 AAAAshared\r\n"
+            b"@cert-authority 10.42.0.2,example.com ssh-ed25519 AAAAca\n"
+            b"@revoked 10.42.0.2 ssh-ed25519 AAAArevoked\n"
+            b"10.42.* ssh-ed25519 AAAAwildcard\n"
+            b"!10.42.0.2,* ssh-ed25519 AAAAnegated\n"
+            b"|1|not-base64|invalid ssh-ed25519 AAAAmalformed\n"
+            b"# non-UTF8 comment: \xff"
+        )
+        self.manager.remove(self.vm)
+        self.assertEqual(self.path.read_bytes(),
+            b"# Known hosts\r\n"
+            b"10.42.0.20 ssh-ed25519 AAAAother same prefix\n"
+            b"some-host,another-host  ssh-ed25519 AAAAshared\r\n"
+            b"@cert-authority example.com ssh-ed25519 AAAAca\n"
+            b"10.42.* ssh-ed25519 AAAAwildcard\n"
+            b"!10.42.0.2,* ssh-ed25519 AAAAnegated\n"
+            b"|1|not-base64|invalid ssh-ed25519 AAAAmalformed\n"
+            b"# non-UTF8 comment: \xff"
+        )
+
+    def test_missing_files_and_unmatched_hosts_are_noops(self) -> None:
+        missing = self.directory / "missing" / ".ssh" / "known_hosts"
+        KnownHostsManager(Settings(known_hosts_files=(missing, self.path))).remove(self.vm)
+        self.assertEqual(list(self.directory.iterdir()), [])
+        self.path.write_bytes(b"10.42.0.3 ssh-ed25519 AAAAother\n")
+        before = self.path.stat()
+        with patch("uvm.known_hosts.os.replace") as publish:
+            self.manager.remove(self.vm)
+        publish.assert_not_called()
+        self.assertEqual(self.path.stat().st_ino, before.st_ino)
+        self.assertEqual(self.path.stat().st_mtime_ns, before.st_mtime_ns)
+
+    def test_rewrite_preserves_ownership_and_file_mode_and_is_idempotent(self) -> None:
+        self.path.write_bytes(b"10.42.0.2 ssh-ed25519 AAAAold\n10.42.0.3 ssh-ed25519 AAAAkeep\n")
+        self.path.chmod(0o640)
+        before = self.path.stat()
+        self.manager.remove(self.vm)
+        after = self.path.stat()
+        self.assertEqual((after.st_uid, after.st_gid, after.st_mode), (before.st_uid, before.st_gid, before.st_mode))
+        self.assertEqual(self.path.read_bytes(), b"10.42.0.3 ssh-ed25519 AAAAkeep\n")
+        self.manager.remove(self.vm)
+        self.assertEqual(self.path.stat().st_ino, after.st_ino)
+        self.assertEqual(list(self.directory.iterdir()), [self.path])
+
+    def test_root_rewrite_restores_the_file_owners_uid_and_gid(self) -> None:
+        self.path.write_bytes(b"10.42.0.2 ssh-ed25519 AAAAold\n")
+        inode = self.path.stat().st_ino
+        real_fstat, real_stat = os.fstat, os.stat
+
+        def ownership(metadata):
+            if metadata.st_ino != inode:
+                return metadata
+            values = list(metadata)
+            values[4:6] = [12345, 23456]
+            return os.stat_result(values)
+
+        with (
+            patch("uvm.known_hosts.os.fstat", side_effect=lambda fd: ownership(real_fstat(fd))),
+            patch("uvm.known_hosts.os.stat", side_effect=lambda *args, **kwargs: ownership(real_stat(*args, **kwargs))),
+            patch("uvm.known_hosts.os.fchown") as chown,
+        ):
+            self.manager.remove(self.vm)
+        self.assertEqual(chown.call_count, 1)
+        self.assertEqual(chown.call_args.args[1:], (12345, 23456))
+        self.assertEqual(self.path.read_bytes(), b"")
+
+    def test_extended_file_attributes_survive_replacement(self) -> None:
+        self.path.write_bytes(b"10.42.0.2 ssh-ed25519 AAAAold\n")
+        try:
+            os.setxattr(self.path, "user.uvm-fixture", b"preserved")
+        except OSError as error:
+            if error.errno in (errno.ENOTSUP, errno.EPERM):
+                self.skipTest("extended file attributes are unavailable")
+            raise
+        self.manager.remove(self.vm)
+        self.assertEqual(os.getxattr(self.path, "user.uvm-fixture"), b"preserved")
+
+    def test_default_files_include_root_and_sudo_user_from_passwd_not_home(self) -> None:
+        root = SimpleNamespace(pw_uid=0, pw_dir="/root")
+        user = SimpleNamespace(pw_uid=1000, pw_dir="/home/operator")
+        with (
+            patch("uvm.known_hosts.os.geteuid", return_value=0),
+            patch("uvm.known_hosts.pwd.getpwuid", side_effect=lambda uid: {0: root, 1000: user}[uid]),
+            patch.dict(os.environ, {"HOME": "/wrong-home", "SUDO_UID": "1000"}, clear=True),
+        ):
+            paths = KnownHostsManager(Settings())._files()
+        self.assertEqual(paths, {
+            Path("/root/.ssh/known_hosts"): 0, Path("/root/.ssh/known_hosts2"): 0,
+            Path("/home/operator/.ssh/known_hosts"): 1000, Path("/home/operator/.ssh/known_hosts2"): 1000,
+        })
+
+    def test_sudo_user_fallback_and_nonroot_context_are_scoped_to_the_launcher(self) -> None:
+        user = SimpleNamespace(pw_uid=1000, pw_dir="/home/operator")
+        with (
+            patch("uvm.known_hosts.os.geteuid", return_value=0),
+            patch("uvm.known_hosts.pwd.getpwuid", return_value=SimpleNamespace(pw_uid=0, pw_dir="/root")),
+            patch("uvm.known_hosts.pwd.getpwnam", return_value=user) as lookup,
+            patch.dict(os.environ, {"SUDO_USER": "operator"}, clear=True),
+        ):
+            self.assertIn(Path("/home/operator/.ssh/known_hosts"), KnownHostsManager(Settings())._files())
+        lookup.assert_called_once_with("operator")
+        with (
+            patch("uvm.known_hosts.os.geteuid", return_value=1000),
+            patch("uvm.known_hosts.pwd.getpwuid", return_value=user),
+            patch.dict(os.environ, {"SUDO_UID": "9999"}, clear=True),
+        ):
+            self.assertEqual(set(KnownHostsManager(Settings())._files()), {
+                Path("/home/operator/.ssh/known_hosts"), Path("/home/operator/.ssh/known_hosts2"),
+            })
+
+    def test_publication_failure_preserves_original_and_removes_temporary_file(self) -> None:
+        contents = b"10.42.0.2 ssh-ed25519 AAAAold\n"
+        self.path.write_bytes(contents)
+        with patch("uvm.known_hosts.os.replace", side_effect=OSError("publish failed")):
+            with self.assertRaisesRegex(UvmError, "publish failed"):
+                self.manager.remove(self.vm)
+        self.assertEqual(self.path.read_bytes(), contents)
+        self.assertEqual(list(self.directory.iterdir()), [self.path])
+
+    def test_concurrent_append_is_not_lost_and_cleanup_can_be_retried(self) -> None:
+        self.path.write_bytes(b"10.42.0.2 ssh-ed25519 AAAAold\n")
+        addition = b"10.42.0.4 ssh-ed25519 AAAAconcurrent\n"
+
+        def append(contents, hosts, candidates):
+            with self.path.open("ab") as output:
+                output.write(addition)
+            return _without_hosts(contents, hosts, candidates)
+
+        with patch("uvm.known_hosts._without_hosts", side_effect=append):
+            with self.assertRaisesRegex(UvmError, "changed during cleanup"):
+                self.manager.remove(self.vm)
+        self.assertTrue(self.path.read_bytes().endswith(addition))
+        self.manager.remove(self.vm)
+        self.assertEqual(self.path.read_bytes(), addition)
+
+    def test_temporary_name_collision_never_removes_a_foreign_file(self) -> None:
+        self.path.write_bytes(b"10.42.0.2 ssh-ed25519 AAAAold\n")
+        foreign = self.directory / ".known_hosts.uvm-collision"
+        foreign.write_bytes(b"keep this")
+        with patch("uvm.known_hosts.uuid.uuid4", return_value=SimpleNamespace(hex="collision")):
+            with self.assertRaises(UvmError):
+                self.manager.remove(self.vm)
+        self.assertEqual(foreign.read_bytes(), b"keep this")
+        self.assertIn(b"10.42.0.2", self.path.read_bytes())
+
+    def test_symlink_and_hardlinked_files_are_not_rewritten(self) -> None:
+        target = self.directory / "target"
+        target.write_bytes(b"10.42.0.2 ssh-ed25519 AAAAold\n")
+        for kind in ("symlink", "hardlink"):
+            with self.subTest(kind=kind):
+                self.path.unlink(missing_ok=True)
+                if kind == "symlink":
+                    self.path.symlink_to(target)
+                else:
+                    self.path.hardlink_to(target)
+                with self.assertRaises(UvmError):
+                    self.manager.remove(self.vm)
+                self.assertEqual(target.read_bytes(), b"10.42.0.2 ssh-ed25519 AAAAold\n")
+
+    @unittest.skipUnless(shutil.which("ssh-keygen"), "OpenSSH is required for hashed-host interoperability")
+    def test_removes_real_openssh_hashed_hosts_for_all_guest_keys_and_keeps_other_hosts(self) -> None:
+        runner = CommandRunner(emit=None)
+        lines = []
+        for key_type in ("ed25519", "ecdsa"):
+            key = self.directory / key_type
+            runner.run(("ssh-keygen", "-q", "-t", key_type, "-N", "", "-f", key), capture=True)
+            public = key.with_suffix(".pub").read_text().strip()
+            for host in ("10.42.0.2", "uvm-vm-test", "[10.42.0.2]:2222", "[uvm-vm-test]:2222", "10.42.0.3"):
+                lines.append(f"{host} {public}\n")
+        self.path.write_text("".join(lines))
+        runner.run(("ssh-keygen", "-H", "-f", self.path), capture=True)
+        self.assertNotIn(b"10.42.0.2", self.path.read_bytes())
+        for host in ("10.42.0.2", "uvm-vm-test", "[10.42.0.2]:2222", "[uvm-vm-test]:2222"):
+            runner.run(("ssh-keygen", "-F", host, "-f", self.path), capture=True)
+        self.manager.remove(self.vm)
+        for host in ("10.42.0.2", "uvm-vm-test", "[10.42.0.2]:2222", "[uvm-vm-test]:2222"):
+            result = runner.run(("ssh-keygen", "-F", host, "-f", self.path), capture=True, check=False)
+            self.assertEqual(result.returncode, 1)
+        runner.run(("ssh-keygen", "-F", "10.42.0.3", "-f", self.path), capture=True)
+
+
+class DeleteKnownHostsTests(LifecycleTestCase):
+    def test_deletion_cleans_keys_after_teardown_before_releasing_the_ip(self) -> None:
+        vm = self.create_vm()
+        self.host.events.clear()
+        self.service.destroy(vm.id)
+        self.assert_call_order(vm, "process.gone", "network.delete_vm_tap", "firewall.remove", "known_hosts.remove")
+        self.assertEqual(self.host.records["known_hosts.remove", vm.id].operation, "delete")
+        self.assertNotIn(vm.id, self.store.load().vms)
+
+    def test_stop_restart_and_resize_retain_host_keys(self) -> None:
+        vm = self.create_vm()
+        self.service.stop(vm.id)
+        self.service.start(vm.id)
+        self.service.restart(vm.id)
+        self.service.resize(vm.id, ram_mib=1024)
+        self.assertNotIn("known_hosts.remove", self.host.actions(vm.id))
+
+    def test_failed_stop_does_not_remove_host_keys(self) -> None:
+        vm = self.create_vm()
+        self.host.failures["process.terminate", vm.id] = FirecrackerError("still alive")
+        with self.assertRaises(UvmError):
+            self.service.destroy(vm.id)
+        self.assertNotIn("known_hosts.remove", self.host.actions(vm.id))
+
+    def test_cleanup_error_retains_the_ip_reservation_and_delete_retries(self) -> None:
+        vm = self.create_vm()
+        self.host.failures["known_hosts.remove", vm.id] = UvmError("known_hosts is not writable")
+        with self.assertRaisesRegex(UvmError, "known_hosts is not writable"):
+            self.service.destroy(vm.id)
+        self.assertEqual(self.store.get(vm.id).operation, "delete")
+        self.assertFalse(Path(vm.disk).exists())
+        with self.assertRaises(ValidationError):
+            self.create_vm(guest_ip=IPv4Address(vm.guest_ip))
+        self.service.destroy(vm.id)
+        created = self.create_vm(guest_ip=IPv4Address(vm.guest_ip))
+        self.assertEqual(created.guest_ip, vm.guest_ip)
+
+    def test_delete_removes_real_cache_entries_and_allows_reuse_without_stale_keys(self) -> None:
+        vm = self.create_vm()
+        cache = self.settings.base / "known_hosts"
+        cache.write_text(f"{vm.guest_ip} ssh-ed25519 AAAAold\n10.42.0.99 ssh-ed25519 AAAAother\n")
+        self.service._known_hosts = KnownHostsManager(replace(self.settings, known_hosts_files=(cache,)))
+        self.service.destroy(vm.id)
+        self.assertEqual(cache.read_text(), "10.42.0.99 ssh-ed25519 AAAAother\n")
+        replacement = self.create_vm(guest_ip=IPv4Address(vm.guest_ip))
+        self.assertEqual(replacement.guest_ip, vm.guest_ip)
+
+    def test_bulk_delete_uses_the_same_known_hosts_cleanup(self) -> None:
+        vms = [self.create_vm(), self.create_vm()]
+        results = self.service.run_many("delete", [vm.id for vm in vms])
+        self.assertTrue(all(result.outcome == "succeeded" for result in results))
+        for vm in vms:
+            self.assertIn("known_hosts.remove", self.host.actions(vm.id))

uvm/known_hosts.py

@@ -84,12 +84,14 @@ class KnownHostsManager:
                 return
             with os.fdopen(descriptor, "rb") as source:
                 metadata = os.fstat(source.fileno())
-                if not stat.S_ISREG(metadata.st_mode) or metadata.st_nlink != 1:
+                if not stat.S_ISREG(metadata.st_mode):
+                    raise UvmError(f"known_hosts must be a regular, single-link file: {path}")
+                fcntl.flock(source.fileno(), fcntl.LOCK_EX)
+                metadata = os.fstat(source.fileno())
+                if metadata.st_nlink != 1:
                     raise UvmError(f"known_hosts must be a regular, single-link file: {path}")
                 if owner is not None and metadata.st_uid != owner:
                     raise UvmError(f"known_hosts has an unexpected owner: {path}")
-                fcntl.flock(source.fileno(), fcntl.LOCK_EX)
-                metadata = os.fstat(source.fileno())
                 contents = source.read()
                 updated = _without_hosts(contents, hosts, candidates)
                 if updated == contents: