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}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub enum OpenVmmHypervisor {
77 Mshv,
79 Kvm,
81 Whp,
83 Hvf,
85}
86
87#[derive(Debug, Clone)]
89pub struct VmHostInfo {
90 pub vbs_supported: bool,
92 pub snp_status: bool,
94 pub tdx_status: bool,
96}
97
98#[derive(Debug, Clone)]
100pub struct HostContext {
101 pub vm_host_info: Option<VmHostInfo>,
103 pub vendor: Vendor,
105 pub execution_environment: ExecutionEnvironment,
107 pub vpci_supported: bool,
109 pub openvmm_hypervisor: Option<OpenVmmHypervisor>,
111}
112
113impl HostContext {
114 pub async fn new() -> Self {
116 let is_nested = {
117 #[cfg(target_arch = "x86_64")]
119 {
120 let result = safe_intrinsics::cpuid(
121 hvdef::HV_CPUID_FUNCTION_MS_HV_ENLIGHTENMENT_INFORMATION,
122 0,
123 );
124 hvdef::HvEnlightenmentInformation::from(
125 result.eax as u128
126 | (result.ebx as u128) << 32
127 | (result.ecx as u128) << 64
128 | (result.edx as u128) << 96,
129 )
130 .nested()
131 }
132 #[cfg(not(target_arch = "x86_64"))]
134 {
135 false
136 }
137 };
138
139 let vendor = Vendor::host();
140
141 let vm_host_info = {
142 #[cfg(windows)]
143 {
144 crate::vm::hyperv::powershell::run_get_vm_host()
145 .await
146 .ok()
147 .map(|info| VmHostInfo {
148 vbs_supported: info.guest_isolation_types.contains(
149 &crate::vm::hyperv::powershell::HyperVGuestStateIsolationType::Vbs,
150 ),
151 snp_status: info.snp_status,
152 tdx_status: info.tdx_status,
153 })
154 }
155 #[cfg(not(windows))]
156 {
157 None
158 }
159 };
160
161 let vpci_supported = cfg!(windows);
163
164 let openvmm_hypervisor = if cfg!(target_os = "linux") {
165 if fs_err::File::open("/dev/mshv").is_ok() {
166 Some(OpenVmmHypervisor::Mshv)
167 } else if fs_err::File::options()
168 .read(true)
169 .write(true)
170 .open("/dev/kvm")
171 .is_ok()
172 {
173 Some(OpenVmmHypervisor::Kvm)
174 } else {
175 None
176 }
177 } else if cfg!(windows) {
178 Some(OpenVmmHypervisor::Whp)
179 } else if cfg!(target_os = "macos") {
180 Some(OpenVmmHypervisor::Hvf)
181 } else {
182 None
183 };
184
185 Self {
186 vm_host_info,
187 vendor,
188 execution_environment: if is_nested {
189 ExecutionEnvironment::Nested
190 } else {
191 ExecutionEnvironment::Baremetal
192 },
193 vpci_supported,
194 openvmm_hypervisor,
195 }
196 }
197}
198
199pub enum TestRequirement {
201 ExecutionEnvironment(ExecutionEnvironment),
203 Vendor(Vendor),
205 Isolation(IsolationType),
207 RequiresCapability {
218 name: &'static str,
220 vmm: VmmType,
222 },
223 And(Box<TestRequirement>, Box<TestRequirement>),
225 Or(Box<TestRequirement>, Box<TestRequirement>),
227 Not(Box<TestRequirement>),
229 Any,
231}
232
233impl TestRequirement {
234 pub fn and(self, other: TestRequirement) -> TestRequirement {
236 TestRequirement::And(Box::new(self), Box::new(other))
237 }
238
239 pub fn or(self, other: TestRequirement) -> TestRequirement {
241 TestRequirement::Or(Box::new(self), Box::new(other))
242 }
243
244 #[expect(clippy::should_implement_trait)]
246 pub fn not(self) -> TestRequirement {
247 TestRequirement::Not(Box::new(self))
248 }
249
250 pub fn is_satisfied(&self, context: &HostContext) -> bool {
252 match self {
253 TestRequirement::ExecutionEnvironment(env) => context.execution_environment == *env,
254 TestRequirement::Vendor(vendor) => context.vendor == *vendor,
255 TestRequirement::Isolation(isolation_type) => {
256 if let Some(vm_host_info) = &context.vm_host_info {
257 match isolation_type {
258 IsolationType::Vbs => vm_host_info.vbs_supported,
259 IsolationType::Snp => vm_host_info.snp_status,
260 IsolationType::Tdx => vm_host_info.tdx_status,
261 }
262 } else {
263 false
264 }
265 }
266 TestRequirement::RequiresCapability { name, vmm } => {
267 available_capabilities(context, *vmm).contains(name)
268 }
269 TestRequirement::And(req1, req2) => {
270 req1.is_satisfied(context) && req2.is_satisfied(context)
271 }
272 TestRequirement::Or(req1, req2) => {
273 req1.is_satisfied(context) || req2.is_satisfied(context)
274 }
275 TestRequirement::Not(req) => !req.is_satisfied(context),
276 TestRequirement::Any => true,
277 }
278 }
279}
280
281pub fn known_capability(name: &str) -> Option<&'static str> {
283 capabilities::known(name)
284}
285
286pub fn is_known_capability(name: &str) -> bool {
288 known_capability(name).is_some()
289}
290
291fn available_capabilities(context: &HostContext, vmm: VmmType) -> BTreeSet<&'static str> {
292 let mut capabilities = BTreeSet::new();
293
294 if context.vpci_supported {
295 capabilities.insert(capabilities::VPCI);
296 }
297
298 if !matches!(
304 (vmm, context.openvmm_hypervisor),
305 (VmmType::OpenVmm, Some(OpenVmmHypervisor::Mshv))
306 ) {
307 capabilities.insert(capabilities::WINDOWS_PARTITION_RESET);
308 }
309
310 match std::env::var("PETRI_CAPABILITIES") {
311 Ok(env_capabilities) => {
312 for capability in env_capabilities.split(',').map(str::trim) {
313 if capability.is_empty() {
314 continue;
315 }
316 let capability = known_capability(capability)
317 .unwrap_or_else(|| panic!("unknown PETRI_CAPABILITIES entry: {capability}"));
318 capabilities.insert(capability);
319 }
320 }
321 Err(std::env::VarError::NotPresent) => {}
322 Err(std::env::VarError::NotUnicode(_)) => {
323 panic!("PETRI_CAPABILITIES is not valid UTF-8")
324 }
325 }
326
327 capabilities
328}
329
330#[derive(Debug, Clone)]
332pub struct TestEvaluationResult {
333 pub test_name: String,
335 pub can_run: bool,
337}
338
339impl TestEvaluationResult {
340 pub fn new(test_name: &str) -> Self {
342 Self {
343 test_name: test_name.to_string(),
344 can_run: true,
345 }
346 }
347}
348
349pub struct TestCaseRequirements {
351 requirements: TestRequirement,
352}
353
354impl TestCaseRequirements {
355 pub fn new(requirements: TestRequirement) -> Self {
357 Self { requirements }
358 }
359}
360
361pub fn can_run_test_with_context(
363 config: Option<&TestCaseRequirements>,
364 context: &HostContext,
365) -> bool {
366 if let Some(config) = config {
367 config.requirements.is_satisfied(context)
368 } else {
369 true
370 }
371}
372
373#[cfg(test)]
374mod tests {
375 use super::*;
376
377 fn host_context(openvmm_hypervisor: OpenVmmHypervisor) -> HostContext {
378 HostContext {
379 vm_host_info: None,
380 vendor: Vendor::Intel,
381 execution_environment: ExecutionEnvironment::Baremetal,
382 vpci_supported: false,
383 openvmm_hypervisor: Some(openvmm_hypervisor),
384 }
385 }
386
387 #[test]
388 fn capabilities_are_evaluated_for_the_selected_vmm() {
389 let requirement = |vmm| TestRequirement::RequiresCapability {
390 name: capabilities::WINDOWS_PARTITION_RESET,
391 vmm,
392 };
393 let mshv = host_context(OpenVmmHypervisor::Mshv);
394
395 assert!(!requirement(VmmType::OpenVmm).is_satisfied(&mshv));
396 assert!(requirement(VmmType::HyperV).is_satisfied(&mshv));
397 assert!(requirement(VmmType::OpenVmm).is_satisfied(&host_context(OpenVmmHypervisor::Kvm)));
398 }
399}