flowey_lib_hvlite/
install_vmm_tests_deps.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Hyper-V test pre-reqs
5
6use flowey::node::prelude::*;
7use std::collections::BTreeSet;
8
9const HYPERV_TESTS_REQUIRED_FEATURES: [&str; 3] = [
10    "Microsoft-Hyper-V",
11    "Microsoft-Hyper-V-Management-PowerShell",
12    "Microsoft-Hyper-V-Management-Clients",
13];
14
15const WHP_TESTS_REQUIRED_FEATURES: [&str; 1] = ["HypervisorPlatform"];
16
17const VIRT_REG_PATH: &str = r#"HKLM\Software\Microsoft\Windows NT\CurrentVersion\Virtualization"#;
18
19#[derive(Serialize, Deserialize, Debug, PartialEq)]
20pub enum VmmTestsDepSelections {
21    Windows {
22        hyperv: bool,
23        whp: bool,
24        hardware_isolation: bool,
25    },
26    Linux,
27}
28
29flowey_config! {
30    /// Config for the install_vmm_tests_deps node.
31    pub struct Config {
32        /// Specify the necessary dependencies
33        pub selections: Option<VmmTestsDepSelections>,
34        /// Automatically install dependencies (requires admin privileges).
35        ///
36        /// When false, assume all dependencies are already present and skip
37        /// checks that require admin privileges (e.g., DISM.exe).
38        ///
39        /// Must be set to true/false when running locally.
40        pub auto_install: Option<bool>,
41    }
42}
43
44flowey_request! {
45    pub enum Request {
46        /// Install the dependencies
47        Install(WriteVar<SideEffect>),
48        /// Generate a list of commands that would install the dependencies
49        GetCommands(WriteVar<Vec<String>>),
50    }
51}
52
53new_flow_node_with_config!(struct Node);
54
55impl FlowNodeWithConfig for Node {
56    type Request = Request;
57    type Config = Config;
58
59    fn imports(_ctx: &mut ImportCtx<'_>) {}
60
61    fn emit(
62        config: Config,
63        requests: Vec<Self::Request>,
64        ctx: &mut NodeCtx<'_>,
65    ) -> anyhow::Result<()> {
66        let mut installed = Vec::new();
67        let mut write_commands = Vec::new();
68        for req in requests {
69            match req {
70                Request::Install(v) => installed.push(v),
71                Request::GetCommands(v) => write_commands.push(v),
72            }
73        }
74
75        let installed = installed;
76        let write_commands = write_commands;
77
78        // Return if no requests specified
79        if installed.is_empty() && write_commands.is_empty() {
80            return Ok(());
81        }
82
83        let selections = config
84            .selections
85            .ok_or(anyhow::anyhow!("missing config: selections"))?;
86        let auto_install = config.auto_install;
87        let installing = !installed.is_empty();
88
89        match selections {
90            VmmTestsDepSelections::Windows {
91                hyperv,
92                whp,
93                hardware_isolation,
94            } => {
95                ctx.emit_rust_step("install vmm tests deps (windows)", move |ctx| {
96                    installed.claim(ctx);
97                    let write_commands = write_commands.claim(ctx);
98
99                    move |rt| {
100                        let mut commands = Vec::new();
101
102                        if !matches!(rt.platform(), FlowPlatform::Windows)
103                            && !flowey_lib_common::_util::running_in_wsl(rt)
104                        {
105                            anyhow::bail!("Must be on Windows or WSL2 to install Windows deps.")
106                        }
107
108                        // Resolve auto_install for local backend
109                        let auto_install = match rt.backend() {
110                            FlowBackend::Local => auto_install.ok_or_else(|| {
111                                anyhow::anyhow!("Missing essential request: AutoInstall")
112                            })?,
113                            // CI backends always auto-install
114                            FlowBackend::Ado | FlowBackend::Github => true,
115                        };
116
117                        // TODO: add these features and reg keys to the initial CI image
118
119                        // Select required features
120                        let mut features_to_enable = BTreeSet::new();
121                        if hyperv {
122                            features_to_enable.append(&mut HYPERV_TESTS_REQUIRED_FEATURES.into());
123                        }
124                        if whp {
125                            features_to_enable.append(&mut WHP_TESTS_REQUIRED_FEATURES.into());
126                        }
127
128                        // Check if features are already enabled (requires admin, so skip if not auto_install)
129                        if installing && auto_install && !features_to_enable.is_empty() {
130                            let features = flowey::shell_cmd!(rt, "DISM.exe /Online /Get-Features").output()?;
131                            assert!(features.status.success());
132                            let features = String::from_utf8_lossy(&features.stdout).to_string();
133                            let mut feature = None;
134                            for line in features.lines() {
135                                if let Some((k, v)) = line.split_once(":") {
136                                    if let Some(f) = feature {
137                                        assert_eq!(k.trim(), "State");
138                                        match v.trim() {
139                                            "Enabled" => {
140                                                assert!(features_to_enable.remove(f));
141                                            }
142                                            "Disabled" => {}
143                                            _ => anyhow::bail!("Unknown feature enablement state"),
144                                        }
145                                        feature = None;
146                                    } else if k.trim() == "Feature Name" {
147                                        let new_feature = v.trim();
148                                        feature = features_to_enable.contains(new_feature).then_some(new_feature);
149                                    }
150                                }
151                            }
152                        } else if installing && !auto_install && !features_to_enable.is_empty() {
153                            // Not auto-installing, assume features are already present
154                            log::info!("Skipping Windows feature check (requires admin). Assuming features are already enabled.");
155                            features_to_enable.clear();
156                        }
157
158                        // Prompt before enabling when running locally
159                        if installing && auto_install && !features_to_enable.is_empty() && matches!(rt.backend(), FlowBackend::Local) {
160                            let mut features_to_install_string = String::new();
161                            for feature in features_to_enable.iter() {
162                                features_to_install_string.push_str(feature);
163                                features_to_install_string.push('\n');
164                            }
165
166                            log::warn!(
167                                r#"
168================================================================================
169To run the VMM tests, the following features need to be enabled:
170{features_to_install_string}
171
172You may need to restart your system for the changes to take effect.
173
174If you're OK with installing these features, please press <enter>.
175Otherwise, press `ctrl-c` to cancel the run.
176================================================================================
177"#
178                            );
179                            let _ = std::io::stdin().read_line(&mut String::new());
180                        }
181
182                        // Install the features
183                        for feature in features_to_enable {
184                            if installing && auto_install {
185                                flowey::shell_cmd!(rt, "DISM.exe /Online /NoRestart /Enable-Feature /All /FeatureName:{feature}").run()?;
186                            }
187                            commands.push(format!("DISM.exe /Online /NoRestart /Enable-Feature /All /FeatureName:{feature}"));
188                        }
189
190                        // Select required reg keys
191                        let mut reg_keys_to_set = BTreeSet::new();
192                        if hyperv {
193                            // Allow loading IGVM from file (to run custom OpenHCL firmware)
194                            reg_keys_to_set.insert("AllowFirmwareLoadFromFile");
195                            // Enable COM3 and COM4 for Hyper-V VMs so we can get the OpenHCL KMSG logs over serial
196                            reg_keys_to_set.insert("EnableAdditionalComPorts");
197
198                            if hardware_isolation {
199                                reg_keys_to_set.insert("EnableHardwareIsolation");
200                            }
201                        }
202
203                        // Check if reg keys are set (skip if not auto_install, assume already set)
204                        if installing && auto_install && !reg_keys_to_set.is_empty() {
205                            let output = flowey::shell_cmd!(rt, "reg.exe query {VIRT_REG_PATH}").output()?;
206                            if output.status.success() {
207                                let output = String::from_utf8_lossy(&output.stdout).to_string();
208                                for line in output.lines() {
209                                    let components = line.split_whitespace().collect::<Vec<_>>();
210                                    if components.len() == 3
211                                        && reg_keys_to_set.contains(components[0])
212                                        && components[1] == "REG_DWORD"
213                                        && components[2] == "0x1"
214                                    {
215                                        assert!(reg_keys_to_set.remove(components[0]));
216                                    }
217                                }
218                            }
219                        } else if installing && !auto_install && !reg_keys_to_set.is_empty() {
220                            // Not auto-installing, assume reg keys are already set
221                            log::info!("Skipping registry key check. Assuming keys are already set.");
222                            reg_keys_to_set.clear();
223                        }
224
225                        // Prompt before changing registry when running locally
226                        if installing && auto_install && !reg_keys_to_set.is_empty() && matches!(rt.backend(), FlowBackend::Local) {
227                            let mut reg_keys_to_set_string = String::new();
228                            for feature in reg_keys_to_set.iter() {
229                                reg_keys_to_set_string.push_str(feature);
230                                reg_keys_to_set_string.push('\n');
231                            }
232
233                            log::warn!(
234                                r#"
235================================================================================
236To run the VMM tests, the following registry keys need to be set to 1:
237{reg_keys_to_set_string}
238
239If you're OK with changing the registry, please press <enter>.
240Otherwise, press `ctrl-c` to cancel the run.
241================================================================================
242"#
243                            );
244                            let _ = std::io::stdin().read_line(&mut String::new());
245                        }
246
247                        // Modify the registry
248                        for v in reg_keys_to_set {
249                            // TODO: figure out why reg.exe is not found if I
250                            // render the command as a string first and share
251                            if installing && auto_install {
252                                flowey::shell_cmd!(rt, "reg.exe add {VIRT_REG_PATH} /v {v} /t REG_DWORD /d 1 /f").run()?;
253                            }
254                            commands.push(format!("reg.exe add \"{VIRT_REG_PATH}\" /v {v} /t REG_DWORD /d 1 /f"));
255                        }
256
257                        for write_cmds in write_commands {
258                            rt.write(write_cmds, &commands);
259                        }
260
261                        Ok(())
262                    }
263                });
264            }
265            VmmTestsDepSelections::Linux => {
266                ctx.emit_rust_step("install vmm tests deps (linux)", |ctx| {
267                    installed.claim(ctx);
268                    let write_commands = write_commands.claim(ctx);
269
270                    |rt| {
271                        for write_cmds in write_commands {
272                            rt.write(write_cmds, &Vec::new());
273                        }
274
275                        Ok(())
276                    }
277                });
278            }
279        }
280
281        Ok(())
282    }
283}
flowey_lib_hvlite/install_vmm_tests_deps.rs

flowey_lib_hvlite/
install_vmm_tests_deps.rs
