345 lines
17 KiB
Python
345 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import stat
|
|
import tempfile
|
|
import unittest
|
|
from dataclasses import replace
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from uvm.config import Settings
|
|
from uvm.domain import STATE_SCHEMA_VERSION, PortRule, State, VmRecord, VmSpec, new_vm_id
|
|
from uvm.errors import ConflictError, NotFoundError, StateError
|
|
from uvm.state import StateStore
|
|
|
|
|
|
def make_vm(vm_id: str = "vm-test") -> VmRecord:
|
|
return VmRecord(
|
|
id=vm_id,
|
|
cpu=1,
|
|
ram_mib=512,
|
|
guest_ip="10.42.0.2",
|
|
gateway="10.42.0.1",
|
|
tap="uvm-test",
|
|
mac="02:fc:00:00:00:01",
|
|
socket="/tmp/firecracker.sock",
|
|
config="/tmp/config.json",
|
|
log="/tmp/firecracker.log",
|
|
disk="/tmp/rootfs.ext4",
|
|
username="admin",
|
|
password="stored-secret",
|
|
)
|
|
|
|
|
|
class StateStoreTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self._temporary_directory = tempfile.TemporaryDirectory()
|
|
self.base = Path(self._temporary_directory.name) / "uvm"
|
|
self.settings = replace(Settings(), base=self.base)
|
|
self.store = StateStore(self.settings)
|
|
|
|
def tearDown(self) -> None:
|
|
self._temporary_directory.cleanup()
|
|
|
|
def test_transaction_persists_a_vm_atomically(self) -> None:
|
|
self.store.initialize()
|
|
with self.store.transaction() as state:
|
|
state.vms["vm-test"] = make_vm()
|
|
state.next_mac_index = 2
|
|
|
|
loaded = self.store.load()
|
|
self.assertEqual(loaded.vms["vm-test"].disk, "/tmp/rootfs.ext4")
|
|
self.assertEqual(loaded.vms["vm-test"].username, "admin")
|
|
self.assertEqual(loaded.vms["vm-test"].password, "stored-secret")
|
|
self.assertNotIn("stored-secret", repr(loaded.vms["vm-test"]))
|
|
self.assertEqual(loaded.next_mac_index, 2)
|
|
self.assertEqual(stat.S_IMODE(self.settings.state_path.stat().st_mode), 0o600)
|
|
persisted = json.loads(self.settings.state_path.read_text(encoding="utf-8"))
|
|
self.assertEqual(persisted["vms"]["vm-test"]["username"], "admin")
|
|
self.assertEqual(persisted["vms"]["vm-test"]["password"], "stored-secret")
|
|
|
|
def test_legacy_state_is_loaded_and_migrated_on_next_write(self) -> None:
|
|
self.base.mkdir(parents=True)
|
|
for version in (None, 1):
|
|
with self.subTest(version=version):
|
|
raw_vm = make_vm().to_dict()
|
|
for field in (
|
|
"disk", "username", "password", "updated_at", "name", "ports",
|
|
"disk_ready", "runtime_owned", "tap_owned", "operation", "boot_id",
|
|
"spawn_pending", "network", "bridge",
|
|
):
|
|
raw_vm.pop(field)
|
|
legacy = {"vms": {"vm-test": raw_vm}}
|
|
if version is not None:
|
|
legacy["schema_version"] = version
|
|
self.settings.state_path.write_text(json.dumps(legacy), encoding="utf-8")
|
|
|
|
loaded = self.store.load()
|
|
vm = loaded.vms["vm-test"]
|
|
self.assertEqual(loaded.schema_version, 2)
|
|
self.assertEqual(vm.disk, "")
|
|
self.assertEqual(vm.username, "root")
|
|
self.assertIsNone(vm.password)
|
|
self.assertIsNone(vm.name)
|
|
self.assertIsNone(vm.ports)
|
|
self.assertIsNone(vm.network)
|
|
self.assertIsNone(vm.bridge)
|
|
self.assertFalse(vm.disk_ready)
|
|
self.assertTrue(vm.runtime_owned)
|
|
self.assertFalse(vm.tap_owned)
|
|
self.assertFalse(vm.spawn_pending)
|
|
self.assertEqual(vm.updated_at, vm.created_at)
|
|
self.assertEqual(loaded.next_mac_index, 2)
|
|
self.assertEqual(
|
|
json.loads(self.settings.state_path.read_text(encoding="utf-8")), legacy
|
|
)
|
|
|
|
with self.store.transaction() as state:
|
|
state.vms["vm-test"].status = "stopped"
|
|
|
|
persisted = json.loads(self.settings.state_path.read_text(encoding="utf-8"))
|
|
self.assertEqual(persisted["schema_version"], 2)
|
|
self.assertEqual(persisted["next_mac_index"], 2)
|
|
self.assertEqual(persisted["vms"]["vm-test"]["status"], "stopped")
|
|
self.assertIsNone(persisted["vms"]["vm-test"]["ports"])
|
|
self.assertIs(persisted["vms"]["vm-test"]["runtime_owned"], True)
|
|
self.assertIs(persisted["vms"]["vm-test"]["disk_ready"], False)
|
|
|
|
def test_round_trip_preserves_policy_ownership_and_recovery_fields(self) -> None:
|
|
for ports in (None, (), (PortRule(22), PortRule(8080, "http"))):
|
|
with self.subTest(ports=ports):
|
|
vm = replace(
|
|
make_vm(), name="app-server", ports=ports, disk_ready=True,
|
|
runtime_owned=True, tap_owned=True, operation="configuring",
|
|
boot_id="host-boot-id", spawn_pending=True, pid=123,
|
|
process_start_time="456789", network="10.42.0.0/24", bridge="uvm0",
|
|
last_error="interrupted boot",
|
|
)
|
|
with self.store.transaction() as state:
|
|
state.vms[vm.id] = vm
|
|
|
|
loaded = self.store.load()
|
|
self.assertEqual(loaded.schema_version, 2)
|
|
self.assertEqual(loaded.vms[vm.id], vm)
|
|
raw = json.loads(self.settings.state_path.read_text(encoding="utf-8"))
|
|
self.assertEqual(raw["schema_version"], 2)
|
|
self.assertEqual(
|
|
raw["vms"][vm.id]["ports"],
|
|
None if ports is None else [port.to_dict() for port in ports],
|
|
)
|
|
for field in ("disk_ready", "runtime_owned", "tap_owned", "spawn_pending"):
|
|
self.assertIs(raw["vms"][vm.id][field], True)
|
|
|
|
def test_lookup_requires_full_case_sensitive_ids_names_or_ips(self) -> None:
|
|
vm = replace(make_vm("vm-0123456789abcdef0123456789abcdef"), name="app-server")
|
|
with self.store.transaction() as state:
|
|
state.vms[vm.id] = vm
|
|
|
|
for identifier in (vm.id, vm.name, vm.guest_ip):
|
|
with self.subTest(identifier=identifier):
|
|
self.assertEqual(self.store.find_by_id_or_ip(identifier), vm)
|
|
self.assertEqual(self.store.get(vm.id), vm)
|
|
for identifier in (
|
|
"", vm.id[:15], vm.id.upper(), "app", "APP-SERVER", " app-server ", "10.42.0",
|
|
):
|
|
with self.subTest(identifier=identifier), self.assertRaises(NotFoundError):
|
|
self.store.find_by_id_or_ip(identifier)
|
|
for identifier in (vm.name, vm.guest_ip):
|
|
with self.subTest(get_identifier=identifier), self.assertRaises(NotFoundError):
|
|
self.store.get(identifier)
|
|
|
|
def test_load_rejects_a_record_whose_id_does_not_match_its_key(self) -> None:
|
|
self.base.mkdir(parents=True)
|
|
for version in (0, 1, 2):
|
|
with self.subTest(version=version):
|
|
raw = {
|
|
"schema_version": version,
|
|
"vms": {"vm-wrong-key": make_vm("vm-record-id").to_dict()},
|
|
}
|
|
self.settings.state_path.write_text(json.dumps(raw), encoding="utf-8")
|
|
with self.assertRaisesRegex(StateError, "ID.*registry key"):
|
|
self.store.load()
|
|
|
|
def test_initialize_protects_registry_credentials(self) -> None:
|
|
self.base.mkdir(parents=True)
|
|
self.settings.state_path.write_text('{"vms": {}}', encoding="utf-8")
|
|
self.settings.state_path.chmod(0o644)
|
|
|
|
self.store.initialize()
|
|
|
|
self.assertEqual(stat.S_IMODE(self.settings.state_path.stat().st_mode), 0o600)
|
|
|
|
def test_load_protects_registry_credentials_without_initialization(self) -> None:
|
|
self.base.mkdir(parents=True)
|
|
self.settings.state_path.write_text('{"vms": {}}', encoding="utf-8")
|
|
self.settings.state_path.chmod(0o644)
|
|
|
|
self.store.load()
|
|
|
|
self.assertEqual(stat.S_IMODE(self.settings.state_path.stat().st_mode), 0o600)
|
|
|
|
def test_load_explains_when_registry_permissions_require_root(self) -> None:
|
|
self.base.mkdir(parents=True)
|
|
self.settings.state_path.write_text('{"vms": {}}', encoding="utf-8")
|
|
self.settings.state_path.chmod(0o644)
|
|
|
|
with (
|
|
patch("uvm.state.os.chmod", side_effect=PermissionError),
|
|
self.assertRaisesRegex(StateError, "Run this command with sudo"),
|
|
):
|
|
self.store.load()
|
|
|
|
|
|
class DomainStateTests(unittest.TestCase):
|
|
def test_new_specs_default_to_ssh_without_overriding_an_explicit_empty_policy(self) -> None:
|
|
self.assertEqual(VmSpec(cpu=1, ram_mib=512).ports, (PortRule(22),))
|
|
self.assertEqual(VmSpec(cpu=1, ram_mib=512, ports=()).ports, ())
|
|
self.assertIsNone(make_vm().ports)
|
|
|
|
def test_port_rule_serialization_normalizes_labels_and_defaults_to_tcp(self) -> None:
|
|
for port, protocol in ((1, "tcp"), (8080, "hTtP"), (65535, "https")):
|
|
with self.subTest(port=port, protocol=protocol):
|
|
rule = PortRule.from_dict({"port": port, "protocol": protocol})
|
|
self.assertEqual(rule.to_dict(), {"port": port, "protocol": protocol.upper()})
|
|
self.assertEqual(PortRule.from_dict(rule.to_dict()), rule)
|
|
self.assertEqual(PortRule.from_dict({"port": 22}), PortRule(22, "TCP"))
|
|
|
|
def test_port_rule_deserialization_rejects_invalid_types_fields_and_values(self) -> None:
|
|
invalid = (
|
|
None, True, 22, "22", [], (), {}, {"protocol": "TCP"},
|
|
{"port": True}, {"port": False}, {"port": None}, {"port": "22"},
|
|
{"port": 22.0}, {"port": 0}, {"port": 65536},
|
|
{"port": 22, "protocol": None}, {"port": 22, "protocol": True},
|
|
{"port": 22, "protocol": "UDP"}, {"port": 22, "protocol": " TCP"},
|
|
{"port": 22, "type": "TCP"}, {"port": 22, "protocol": "TCP", "extra": 1},
|
|
)
|
|
for value in invalid:
|
|
with self.subTest(value=value), self.assertRaises(StateError):
|
|
PortRule.from_dict(value)
|
|
|
|
def test_persisted_ports_must_be_a_json_list_or_legacy_null(self) -> None:
|
|
for ports in (True, False, 22, "22", {}, (), (PortRule(22),)):
|
|
with self.subTest(ports=ports), self.assertRaisesRegex(StateError, "ports"):
|
|
VmRecord.from_dict({**make_vm().to_dict(), "ports": ports})
|
|
for ports in ([22], [None], ["22:TCP"], [{"port": True}]):
|
|
with self.subTest(ports=ports), self.assertRaises(StateError):
|
|
VmRecord.from_dict({**make_vm().to_dict(), "ports": ports})
|
|
|
|
def test_persisted_duplicate_and_conflicting_tcp_rules_are_rejected(self) -> None:
|
|
for labels in (("TCP", "TCP"), ("tcp", "TCP"), ("TCP", "HTTP"), ("HTTP", "HTTPS")):
|
|
with self.subTest(labels=labels):
|
|
raw = make_vm().to_dict()
|
|
raw["ports"] = [{"port": 80, "protocol": label} for label in labels]
|
|
with self.assertRaisesRegex(StateError, "duplicate TCP port"):
|
|
VmRecord.from_dict(raw)
|
|
|
|
def test_ownership_and_spawn_pending_fields_require_real_booleans(self) -> None:
|
|
for field in ("disk_ready", "runtime_owned", "tap_owned", "spawn_pending"):
|
|
for value in (None, 0, 1, 0.0, 1.0, "false", "true", [], {}):
|
|
with self.subTest(field=field, value=value):
|
|
raw = {**make_vm().to_dict(), field: value}
|
|
with self.assertRaisesRegex(StateError, f"{field} must be a boolean"):
|
|
VmRecord.from_dict(raw)
|
|
|
|
def test_explicit_recovery_flags_override_legacy_inference(self) -> None:
|
|
fields = ("disk_ready", "runtime_owned", "tap_owned", "spawn_pending")
|
|
for flag in (False, True):
|
|
for operation in (None, "spawning", "configuring"):
|
|
with self.subTest(flag=flag, operation=operation):
|
|
raw = replace(make_vm(), status="running", pid=123, operation=operation).to_dict()
|
|
raw.update(dict.fromkeys(fields, flag))
|
|
vm = VmRecord.from_dict(raw)
|
|
for field in fields:
|
|
self.assertIs(getattr(vm, field), flag)
|
|
self.assertIs(vm.to_dict()[field], flag)
|
|
|
|
def test_legacy_ownership_is_inferred_from_disk_status_and_pid(self) -> None:
|
|
cases = (
|
|
("running", "/tmp/rootfs.ext4", 123, (True, True, True)),
|
|
("stopped", "/tmp/rootfs.ext4", None, (True, True, False)),
|
|
("error", "/tmp/rootfs.ext4", 123, (False, True, True)),
|
|
("starting", "/tmp/rootfs.ext4", None, (False, True, False)),
|
|
("running", "", 123, (False, True, True)),
|
|
("unknown", "", None, (False, True, False)),
|
|
)
|
|
for status, disk, pid, expected in cases:
|
|
with self.subTest(status=status, disk=disk, pid=pid):
|
|
raw = replace(make_vm(), status=status, disk=disk, pid=pid).to_dict()
|
|
for field in ("disk_ready", "runtime_owned", "tap_owned", "ports"):
|
|
raw.pop(field)
|
|
vm = VmRecord.from_dict(raw)
|
|
self.assertEqual((vm.disk_ready, vm.runtime_owned, vm.tap_owned), expected)
|
|
self.assertIsNone(vm.ports)
|
|
|
|
def test_empty_disk_legacy_runtime_is_owned_but_explicit_false_is_respected(self) -> None:
|
|
for disk_fields in ({}, {"disk": ""}):
|
|
with self.subTest(disk_fields=disk_fields):
|
|
raw = replace(make_vm(), status="stopped").to_dict()
|
|
for field in ("disk", "ports", "disk_ready", "runtime_owned"):
|
|
raw.pop(field)
|
|
raw.update(disk_fields)
|
|
|
|
legacy = VmRecord.from_dict(raw)
|
|
self.assertEqual(legacy.disk, "")
|
|
self.assertIs(legacy.runtime_owned, True)
|
|
self.assertIs(legacy.disk_ready, False)
|
|
self.assertIsNone(legacy.ports)
|
|
self.assertEqual(VmRecord.from_dict(legacy.to_dict()), legacy)
|
|
|
|
new = VmRecord.from_dict({**raw, "ports": [], "runtime_owned": False})
|
|
self.assertIs(new.runtime_owned, False)
|
|
self.assertIs(new.disk_ready, False)
|
|
self.assertEqual(new.ports, ())
|
|
self.assertEqual(VmRecord.from_dict(new.to_dict()), new)
|
|
|
|
def test_missing_spawn_pending_recovers_only_a_legacy_spawning_operation(self) -> None:
|
|
for operation in (None, "spawning", "configuring", "stopping", "policy"):
|
|
with self.subTest(operation=operation):
|
|
raw = replace(make_vm(), operation=operation).to_dict()
|
|
raw.pop("spawn_pending")
|
|
vm = VmRecord.from_dict(raw)
|
|
self.assertIs(vm.spawn_pending, operation == "spawning")
|
|
self.assertEqual(VmRecord.from_dict(vm.to_dict()), vm)
|
|
|
|
def test_new_state_uses_schema_two_and_rejects_future_versions(self) -> None:
|
|
self.assertEqual(State().schema_version, 2)
|
|
with self.assertRaisesRegex(StateError, "newer than"):
|
|
State.from_dict({"schema_version": STATE_SCHEMA_VERSION + 1, "vms": {}})
|
|
|
|
def test_generated_ids_have_the_reserved_uuid_format(self) -> None:
|
|
first, second = new_vm_id(), new_vm_id()
|
|
self.assertNotEqual(first, second)
|
|
for vm_id in (first, second):
|
|
self.assertRegex(vm_id, r"\Avm-[0-9a-f]{32}\Z")
|
|
|
|
def test_exact_generated_id_wins_over_a_legacy_name_collision(self) -> None:
|
|
vm = make_vm("vm-" + "a" * 32)
|
|
other = replace(make_vm("vm-" + "b" * 32), name=vm.id, guest_ip="10.42.0.3")
|
|
state = State(vms={other.id: other, vm.id: vm})
|
|
self.assertIs(StateStore.resolve(state, vm.id), vm)
|
|
del state.vms[vm.id]
|
|
with self.assertRaises(NotFoundError):
|
|
StateStore.resolve(state, vm.id)
|
|
|
|
def test_ambiguous_names_ips_and_legacy_ids_require_an_exact_id(self) -> None:
|
|
cases = (
|
|
("vm-" + "a" * 32, "app", "app", "10.42.0.3", "app"),
|
|
("vm-" + "a" * 32, "app", "other", "10.42.0.2", "10.42.0.2"),
|
|
("vm-" + "a" * 32, "app", "10.42.0.2", "10.42.0.3", "10.42.0.2"),
|
|
("vm-legacy", "app", "vm-legacy", "10.42.0.3", "vm-legacy"),
|
|
)
|
|
for vm_id, name, other_name, other_ip, identifier in cases:
|
|
with self.subTest(identifier=identifier, other_name=other_name):
|
|
vm = replace(make_vm(vm_id), name=name)
|
|
other = replace(make_vm("vm-" + "b" * 32), name=other_name, guest_ip=other_ip)
|
|
for records in ((vm, other), (other, vm)):
|
|
state = State(vms={record.id: record for record in records})
|
|
with self.assertRaisesRegex(ConflictError, "ambiguous.*exact VM ID"):
|
|
StateStore.resolve(state, identifier)
|
|
self.assertIs(StateStore.resolve(state, other.id), other)
|
|
|
|
def test_multiple_matching_fields_on_one_vm_are_not_ambiguous(self) -> None:
|
|
vm = replace(make_vm("vm-legacy"), name="vm-legacy")
|
|
self.assertIs(StateStore.resolve(State(vms={vm.id: vm}), "vm-legacy"), vm)
|