src/pqc_hypervisor_attestation/backends/memory.py
2.1 KB · 52 lines · python Raw
1 """In-memory backend — reference implementation for tests and demos."""
2
3 from __future__ import annotations
4
5 from pqc_hypervisor_attestation.backends.base import AttestationBackend
6 from pqc_hypervisor_attestation.errors import InvalidRegionError
7 from pqc_hypervisor_attestation.region import MemoryRegion, RegionSnapshot
8
9
10 class InMemoryBackend(AttestationBackend):
11 """Deterministic in-memory backend.
12
13 Holds region bytes in a dict keyed by ``region_id``. Suitable for
14 tests, tutorials, and CI — it has no dependency on any TEE hardware.
15 """
16
17 name = "in-memory"
18 platform = "in-memory"
19
20 def __init__(self) -> None:
21 self._regions: dict[str, MemoryRegion] = {}
22 self._contents: dict[str, bytes] = {}
23 self._workload_map: dict[str, list[str]] = {} # workload_id -> [region_id]
24
25 def register(self, workload_id: str, region: MemoryRegion, content: bytes) -> None:
26 """Register a region with initial content for the given workload."""
27 self._regions[region.region_id] = region
28 self._contents[region.region_id] = content
29 ids = self._workload_map.setdefault(workload_id, [])
30 if region.region_id not in ids:
31 ids.append(region.region_id)
32
33 def update(self, region_id: str, new_content: bytes) -> None:
34 """Overwrite a region's bytes to simulate a legitimate mutation
35 (or an attacker tampering with memory)."""
36 if region_id not in self._regions:
37 raise InvalidRegionError(f"no region with id {region_id}")
38 self._contents[region_id] = new_content
39
40 def list_regions(self, workload_id: str) -> list[MemoryRegion]:
41 ids = self._workload_map.get(workload_id, [])
42 return [self._regions[rid] for rid in ids]
43
44 def snapshot(self, region: MemoryRegion) -> RegionSnapshot:
45 if region.region_id not in self._contents:
46 raise InvalidRegionError(
47 f"region {region.region_id} not registered"
48 )
49 return RegionSnapshot.create(
50 region.region_id, self._contents[region.region_id]
51 )
52