1use petri_artifacts_common::capabilities;
7use std::collections::BTreeSet;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum ExecutionEnvironment {
12 Baremetal,
14 Nested,
16}
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum Vendor {
21 Amd,
23 Intel,
25 Arm,
27}
28
29impl Vendor {
30 pub fn host() -> Self {
32 #[cfg(target_arch = "x86_64")]
34 {
35 let result =
36 safe_intrinsics::cpuid(x86defs::cpuid::CpuidFunction::VendorAndMaxFunction.0, 0);
37 let vendor =
38 x86defs::cpuid::Vendor::from_ebx_ecx_edx(result.ebx, result.ecx, result.edx);
39 if vendor.is_amd_compatible() {
40 Vendor::Amd
41 } else {
42 assert!(vendor.is_intel_compatible());
43 Vendor::Intel
44 }
45 }
46 #[cfg(not(target_arch = "x86_64"))]
48 {
49 Vendor::Arm
50 }
51 }
52}
53
54#[derive(Clone, Copy, Debug, PartialEq)]
56pub enum IsolationType {
57 Vbs,
59 Snp,
61 Tdx,
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67pub enum VmmType {
68 OpenVmm,
70 HyperV,
72 Qemu,
74}
75
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum OpenVmmHypervisor {
79 Mshv,
81 Kvm,
83 Whp,
85 Hvf,
87}
88
89#[derive(Debug, Clone)]
91pub struct VmHostInfo {
92 pub vbs_supported: bool,
94 pub snp_status: bool,
96 pub tdx_status: bool,
98}
99
100#[derive(Debug, Clone)]
102pub struct HostContext {
103 pub vm_host_info: Option<VmHostInfo>,
105 pub vendor: Vendor,
107 pub execution_environment: ExecutionEnvironment,
109 pub vpci_supported: bool,
111 pub openvmm_hypervisor: Option<OpenVmmHypervisor>,
113}
114
115impl HostContext {
116 pub async fn new() -> Self {
118 let is_nested = {
119 #[cfg(target_arch = "x86_64")]
121 {
122 let result = safe_intrinsics::cpuid(
123 hvdef::HV_CPUID_FUNCTION_MS_HV_ENLIGHTENMENT_INFORMATION,
124 0,
125 );
126 hvdef::HvEnlightenmentInformation::from(
127 result.eax as u128
128 | (result.ebx as u128) << 32
129 | (result.ecx as u128) << 64
130 | (result.edx as u128) << 96,
131 )
132 .nested()
133 }
134 #[cfg(not(target_arch = "x86_64"))]
136 {
137 false
138 }
139 };
140
141 let vendor = Vendor::host();
142
143 let vm_host_info = {
144 #[cfg(windows)]
145 {
146 crate::vm::hyperv::powershell::run_get_vm_host()
147 .await
148 .ok()
149 .map(|info| VmHostInfo {
150 vbs_supported: info.guest_isolation_types.contains(
151 &crate::vm::hyperv::powershell::HyperVGuestStateIsolationType::Vbs,
152 ),
153 snp_status: info.snp_status,
154 tdx_status: info.tdx_status,
155 })
156 }
157 #[cfg(not(windows))]
158 {
159 None
160 }
161 };
162
163 let vpci_supported = cfg!(windows);
165
166 let openvmm_hypervisor = if cfg!(target_os = "linux") {
167 if fs_err::File::open("/dev/mshv").is_ok() {
168 Some(OpenVmmHypervisor::Mshv)
169 } else if fs_err::File::options()
170 .read(true)
171 .write(true)
172 .open("/dev/kvm")
173 .is_ok()
174 {
175 Some(OpenVmmHypervisor::Kvm)
176 } else {
177 None
178 }
179 } else if cfg!(windows) {
180 Some(OpenVmmHypervisor::Whp)
181 } else if cfg!(target_os = "macos") {
182 Some(OpenVmmHypervisor::Hvf)
183 } else {
184 None
185 };
186
187 Self {
188 vm_host_info,
189 vendor,
190 execution_environment: if is_nested {
191 ExecutionEnvironment::Nested
192 } else {
193 ExecutionEnvironment::Baremetal
194 },
195 vpci_supported,
196 openvmm_hypervisor,
197 }
198 }
199}
200
201pub enum TestRequirement {
203 ExecutionEnvironment(ExecutionEnvironment),
205 Vendor(Vendor),
207 Isolation(IsolationType),
209 RequiresCapability {
220 name: &'static str,
222 vmm: VmmType,
224 },
225 And(Box<TestRequirement>, Box<TestRequirement>),
227 Or(Box<TestRequirement>, Box<TestRequirement>),
229 Not(Box<TestRequirement>),
231 Any,
233}
234
235impl TestRequirement {
236 pub fn and(self, other: TestRequirement) -> TestRequirement {
238 TestRequirement::And(Box::new(self), Box::new(other))
239 }
240
241 pub fn or(self, other: TestRequirement) -> TestRequirement {
243 TestRequirement::Or(Box::new(self), Box::new(other))
244 }
245
246 #[expect(clippy::should_implement_trait)]
248 pub fn not(self) -> TestRequirement {
249 TestRequirement::Not(Box::new(self))
250 }
251
252 pub fn is_satisfied(&self, context: &HostContext) -> bool {
254 match self {
255 TestRequirement::ExecutionEnvironment(env) => context.execution_environment == *env,
256 TestRequirement::Vendor(vendor) => context.vendor == *vendor,
257 TestRequirement::Isolation(isolation_type) => {
258 if let Some(vm_host_info) = &context.vm_host_info {
259 match isolation_type {
260 IsolationType::Vbs => vm_host_info.vbs_supported,
261 IsolationType::Snp => vm_host_info.snp_status,
262 IsolationType::Tdx => vm_host_info.tdx_status,
263 }
264 } else {
265 false
266 }
267 }
268 TestRequirement::RequiresCapability { name, vmm } => {
269 available_capabilities(context, *vmm).contains(name)
270 }
271 TestRequirement::And(req1, req2) => {
272 req1.is_satisfied(context) && req2.is_satisfied(context)
273 }
274 TestRequirement::Or(req1, req2) => {
275 req1.is_satisfied(context) || req2.is_satisfied(context)
276 }
277 TestRequirement::Not(req) => !req.is_satisfied(context),
278 TestRequirement::Any => true,
279 }
280 }
281}
282
283pub fn known_capability(name: &str) -> Option<&'static str> {
285 capabilities::known(name)
286}
287
288pub fn is_known_capability(name: &str) -> bool {
290 known_capability(name).is_some()
291}
292
293fn available_capabilities(context: &HostContext, vmm: VmmType) -> BTreeSet<&'static str> {
294 let mut capabilities = BTreeSet::new();
295
296 if context.vpci_supported {
297 capabilities.insert(capabilities::VPCI);
298 }
299
300 if !matches!(
306 (vmm, context.openvmm_hypervisor),
307 (VmmType::OpenVmm, Some(OpenVmmHypervisor::Mshv))
308 ) {
309 capabilities.insert(capabilities::WINDOWS_PARTITION_RESET);
310 }
311
312 match std::env::var("PETRI_CAPABILITIES") {
313 Ok(env_capabilities) => {
314 for capability in env_capabilities.split(',').map(str::trim) {
315 if capability.is_empty() {
316 continue;
317 }
318 let capability = known_capability(capability)
319 .unwrap_or_else(|| panic!("unknown PETRI_CAPABILITIES entry: {capability}"));
320 capabilities.insert(capability);
321 }
322 }
323 Err(std::env::VarError::NotPresent) => {}
324 Err(std::env::VarError::NotUnicode(_)) => {
325 panic!("PETRI_CAPABILITIES is not valid UTF-8")
326 }
327 }
328
329 capabilities
330}
331
332pub struct TestCaseRequirements {
334 requirements: TestRequirement,
335}
336
337impl TestCaseRequirements {
338 pub fn new(requirements: TestRequirement) -> Self {
340 Self { requirements }
341 }
342}
343
344pub fn can_run_test_with_context(
346 config: Option<&TestCaseRequirements>,
347 context: &HostContext,
348) -> bool {
349 if let Some(config) = config {
350 config.requirements.is_satisfied(context)
351 } else {
352 true
353 }
354}
355
356#[cfg(test)]
357mod tests {
358 use super::*;
359
360 fn host_context(openvmm_hypervisor: OpenVmmHypervisor) -> HostContext {
361 HostContext {
362 vm_host_info: None,
363 vendor: Vendor::Intel,
364 execution_environment: ExecutionEnvironment::Baremetal,
365 vpci_supported: false,
366 openvmm_hypervisor: Some(openvmm_hypervisor),
367 }
368 }
369
370 #[test]
371 fn capabilities_are_evaluated_for_the_selected_vmm() {
372 let requirement = |vmm| TestRequirement::RequiresCapability {
373 name: capabilities::WINDOWS_PARTITION_RESET,
374 vmm,
375 };
376 let mshv = host_context(OpenVmmHypervisor::Mshv);
377
378 assert!(!requirement(VmmType::OpenVmm).is_satisfied(&mshv));
379 assert!(requirement(VmmType::HyperV).is_satisfied(&mshv));
380 assert!(requirement(VmmType::OpenVmm).is_satisfied(&host_context(OpenVmmHypervisor::Kvm)));
381 }
382}