Fix Crosvm USB device tracking

nix/nixos-module.nix

@@ -94,6 +94,9 @@ in
         );
         Restart = "on-failure";
         RestartSec = "5s";
+        StateDirectory = "vhotplug";
+        StateDirectoryMode = "0700";
+        UMask = "0077";
       };
     };
   };

nix/nixos-test.nix

@@ -48,6 +48,8 @@ pkgs.testers.runNixOSTest {
     machine.start()
     machine.wait_for_unit("vhotplug.service")
     machine.succeed("systemctl status vhotplug.service")
+    machine.succeed("test $(stat -c %a /var/lib/vhotplug) = 700")
+    machine.succeed("test $(systemctl show -P UMask vhotplug.service) = 0077")
 
     # Verify vhotplug is waiting for devices
     machine.wait_until_succeeds("journalctl -u vhotplug.service | grep -q 'Waiting for new devices'")

nix/package.nix

@@ -22,7 +22,9 @@ python3Packages.buildPythonApplication {
     qemu-qmp
   ];
 
-  doCheck = false;
+  nativeCheckInputs = with python3Packages; [
+    pytestCheckHook
+  ];
 
   meta = with lib; {
     description = "Hot-plugging USB and PCI devices to virtual machines";

tests/test_crosvmlink.py

@@ -0,0 +1,138 @@
+import asyncio
+from typing import Any
+
+import pytest
+
+from vhotplug.crosvmlink import CrosvmLink, CrosvmVMUnavailableError
+from vhotplug.usb import USBInfo
+
+
+class FakeProcess:
+    def __init__(self, stdout: str, returncode: int = 0) -> None:
+        self.stdout = stdout
+        self.returncode = returncode
+
+    async def communicate(self) -> tuple[bytes, bytes]:
+        return self.stdout.encode(), b""
+
+
+def fake_crosvm(
+    monkeypatch: pytest.MonkeyPatch,
+    list_output: str | list[str] = "devices",
+    attach_output: str = "ok 2",
+    detach_output: str | None = None,
+    list_returncode: int = 0,
+) -> list[tuple[str, ...]]:
+    calls: list[tuple[str, ...]] = []
+    list_outputs = [list_output] if isinstance(list_output, str) else list(list_output)
+
+    async def execute(*args: str, **_kwargs: Any) -> FakeProcess:
+        calls.append(args)
+        if args[1:3] == ("usb", "list"):
+            output = list_outputs.pop(0) if len(list_outputs) > 1 else list_outputs[0]
+            return FakeProcess(output, list_returncode)
+        if args[1:3] == ("usb", "attach"):
+            return FakeProcess(attach_output)
+        if args[1:3] == ("usb", "detach"):
+            return FakeProcess(detach_output if detach_output is not None else f"ok {args[3]}")
+        raise AssertionError(f"Unexpected Crosvm command: {args}")
+
+    monkeypatch.setattr(asyncio, "create_subprocess_exec", execute)
+    return calls
+
+
+def usb() -> USBInfo:
+    return USBInfo(device_node="/dev/bus/usb/001/002", vid="046d", pid="c52b")
+
+
+def link(monkeypatch: pytest.MonkeyPatch) -> CrosvmLink:
+    result = CrosvmLink("/run/app-vm.sock", "/nix/store/crosvm/bin/crosvm")
+    monkeypatch.setattr(result, "_wait_for_boot", lambda: True)
+    return result
+
+
+def test_adds_identical_device_when_port_is_not_known(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = fake_crosvm(monkeypatch, "devices 1 046d c52b")
+
+    port = asyncio.run(link(monkeypatch).add_usb_device(usb()))
+
+    assert port == 2
+    assert any(command[1:3] == ("usb", "attach") for command in calls)
+
+
+def test_reuses_persisted_port(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = fake_crosvm(monkeypatch, "devices 2 046d c52b")
+
+    port = asyncio.run(link(monkeypatch).add_usb_device(usb(), known_port=2))
+
+    assert port == 2
+    assert not any(command[1:3] == ("usb", "attach") for command in calls)
+
+
+def test_removes_exact_port_among_identical_devices(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = fake_crosvm(monkeypatch, "devices 1 046d c52b 2 046d c52b")
+
+    asyncio.run(link(monkeypatch).remove_usb_device(usb(), known_port=2))
+
+    detach_commands = [command for command in calls if command[1:3] == ("usb", "detach")]
+    assert detach_commands == [("/nix/store/crosvm/bin/crosvm", "usb", "detach", "2", "/run/app-vm.sock")]
+
+
+def test_refuses_to_remove_different_device_from_known_port(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = fake_crosvm(monkeypatch, "devices 2 1050 0407")
+
+    with pytest.raises(RuntimeError, match="different device"):
+        asyncio.run(link(monkeypatch).remove_usb_device(usb(), known_port=2))
+
+    assert not any(command[1:3] == ("usb", "detach") for command in calls)
+
+
+def test_refuses_ambiguous_legacy_removal(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = fake_crosvm(monkeypatch, "devices 1 046d c52b 2 046d c52b")
+
+    with pytest.raises(RuntimeError, match="cannot be identified safely"):
+        asyncio.run(link(monkeypatch).remove_usb_device(usb()))
+
+    assert not any(command[1:3] == ("usb", "detach") for command in calls)
+
+
+def test_malformed_ok_adopts_new_port_without_retry(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = fake_crosvm(monkeypatch, ["devices 1 1234 5678", "devices 1 1234 5678 2 046d c52b"], "ok")
+
+    port = asyncio.run(link(monkeypatch).add_usb_device(usb()))
+
+    assert port == 2
+    assert len([command for command in calls if command[1:3] == ("usb", "attach")]) == 1
+
+
+def test_malformed_list_port_raises_runtime_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    fake_crosvm(monkeypatch, "devices invalid 046d c52b")
+
+    with pytest.raises(RuntimeError, match="invalid literal"):
+        asyncio.run(link(monkeypatch).remove_usb_device(usb(), known_port=2))
+
+
+def test_empty_detach_response_raises_runtime_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    fake_crosvm(monkeypatch, "devices 2 046d c52b", detach_output="")
+
+    with pytest.raises(RuntimeError, match="empty response"):
+        asyncio.run(link(monkeypatch).remove_usb_device(usb(), known_port=2))
+
+
+def test_unavailable_vm_has_distinct_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    fake_crosvm(monkeypatch, list_returncode=1)
+    monkeypatch.setattr("vhotplug.crosvmlink.is_unix_socket_alive", lambda *_args: False)
+
+    with pytest.raises(CrosvmVMUnavailableError):
+        asyncio.run(link(monkeypatch).remove_usb_device(usb(), known_port=2))
+
+
+def test_no_available_port_does_not_detach_other_devices(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = fake_crosvm(monkeypatch, "devices 1 1234 5678", "no_available_port")
+    crosvm = link(monkeypatch)
+    crosvm.vm_retry_count = 0
+
+    with pytest.raises(RuntimeError, match="timed out"):
+        asyncio.run(crosvm.add_usb_device(usb()))
+
+    assert not any(command[1:3] == ("usb", "detach") for command in calls)

tests/test_device.py

@@ -0,0 +1,135 @@
+import asyncio
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from vhotplug.crosvmlink import CrosvmVMUnavailableError
+from vhotplug.device import _attach_device_to_vm, _remove_device_from_vm
+from vhotplug.usb import USBInfo
+
+
+class FakeState:
+    def __init__(self, current_vm: str | None = "app-vm") -> None:
+        self.current_vm = current_vm
+        self.removed = False
+        self.attached_port: int | None = None
+
+    def get_vm_for_device(self, _dev_info: USBInfo) -> str | None:
+        return self.current_vm
+
+    def get_crosvm_usb_port(self, _dev_info: USBInfo, _vm_name: str, _socket_path: str) -> int:
+        return 7
+
+    def remove_vm_for_device(self, _dev_info: USBInfo) -> None:
+        self.removed = True
+
+    def set_vm_for_device(self, _dev_info: USBInfo, vm_name: str) -> None:
+        self.current_vm = vm_name
+
+    def set_crosvm_usb_port(
+        self,
+        _dev_info: USBInfo,
+        _vm_name: str,
+        _socket_path: str,
+        port: int | None,
+    ) -> None:
+        self.attached_port = port
+
+    def clear_disconnected(self, _dev_info: USBInfo) -> bool:
+        return False
+
+
+def usb() -> USBInfo:
+    return USBInfo(device_node="/dev/bus/usb/001/002", sys_name="1-2.1")
+
+
+def app_context(state: FakeState) -> Any:
+    config = SimpleNamespace(usb_authorization_enabled=lambda: False)
+    return SimpleNamespace(dev_state=state, config=config, api_server=None)
+
+
+def test_vm_unavailable_still_clears_local_state(monkeypatch: pytest.MonkeyPatch) -> None:
+    state = FakeState()
+
+    async def unavailable(*_args: Any) -> None:
+        raise CrosvmVMUnavailableError("VM unavailable")
+
+    monkeypatch.setattr("vhotplug.device.vmm_remove_device", unavailable)
+
+    asyncio.run(
+        _remove_device_from_vm(
+            app_context(state),
+            usb(),
+            {"name": "app-vm", "type": "crosvm", "socket": "/run/app-vm.sock"},
+        )
+    )
+
+    assert state.removed
+
+
+def test_detach_refusal_preserves_local_state(monkeypatch: pytest.MonkeyPatch) -> None:
+    state = FakeState()
+
+    async def refuse(*_args: Any) -> None:
+        raise RuntimeError("refusing to detach")
+
+    monkeypatch.setattr("vhotplug.device.vmm_remove_device", refuse)
+
+    with pytest.raises(RuntimeError, match="refusing"):
+        asyncio.run(
+            _remove_device_from_vm(
+                app_context(state),
+                usb(),
+                {"name": "app-vm", "type": "crosvm", "socket": "/run/app-vm.sock"},
+            )
+        )
+
+    assert not state.removed
+
+
+def test_move_does_not_attach_when_old_vm_removal_fails(monkeypatch: pytest.MonkeyPatch) -> None:
+    state = FakeState(current_vm="old-vm")
+    attached = False
+
+    async def remove(*_args: Any) -> None:
+        raise RuntimeError("old VM detach failed")
+
+    async def attach(*_args: Any) -> None:
+        nonlocal attached
+        attached = True
+
+    monkeypatch.setattr("vhotplug.device.remove_device", remove)
+    monkeypatch.setattr("vhotplug.device.vmm_add_device", attach)
+
+    with pytest.raises(RuntimeError, match="old VM detach failed"):
+        asyncio.run(
+            _attach_device_to_vm(
+                app_context(state),
+                usb(),
+                {"name": "new-vm", "type": "crosvm", "socket": "/run/new-vm.sock"},
+            )
+        )
+
+    assert not attached
+
+
+def test_crosvm_attach_passes_and_records_guest_port(monkeypatch: pytest.MonkeyPatch) -> None:
+    state = FakeState(current_vm=None)
+
+    async def attach(*_args: Any) -> int:
+        assert _args[-1] == 7
+        return 9
+
+    monkeypatch.setattr("vhotplug.device.vmm_add_device", attach)
+
+    asyncio.run(
+        _attach_device_to_vm(
+            app_context(state),
+            usb(),
+            {"name": "app-vm", "type": "crosvm", "socket": "/run/app-vm.sock"},
+        )
+    )
+
+    assert state.current_vm == "app-vm"
+    assert state.attached_port == 9