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, PartialEq)]
20pub enum VmmTestsDepSelections {
21    Windows {
22        hyperv: bool,
23        whp: bool,
24        hardware_isolation: bool,
25    },
26    Linux,
27}
28
29flowey_request! {
30    pub enum Request {
31        /// Specify the necessary dependencies
32        Select(VmmTestsDepSelections),
33        /// Install the dependencies
34        Install(WriteVar<SideEffect>),
35        /// Generate a list of commands that would install the dependencies
36        GetCommands(WriteVar<Vec<String>>),
37    }
38}
39
40new_flow_node!(struct Node);
41
42impl FlowNode for Node {
43    type Request = Request;
44
45    fn imports(_ctx: &mut ImportCtx<'_>) {}
46
47    fn emit(requests: Vec<Self::Request>, ctx: &mut NodeCtx<'_>) -> anyhow::Result<()> {
48        let mut selections = None;
49        let mut installed = Vec::new();
50        let mut write_commands = Vec::new();
51        for req in requests {
52            match req {
53                Request::Select(v) => same_across_all_reqs("Select", &mut selections, v)?,
54                Request::Install(v) => installed.push(v),
55                Request::GetCommands(v) => write_commands.push(v),
56            }
57        }
58        let selections = selections.ok_or(anyhow::anyhow!("Missing essential request: Select"))?;
59        let installed = installed;
60        let write_commands = write_commands;
61        if installed.is_empty() && write_commands.is_empty() {
62            return Ok(());
63        }
64        let installing = !installed.is_empty();
65
66        match selections {
67            VmmTestsDepSelections::Windows {
68                hyperv,
69                whp,
70                hardware_isolation,
71            } => {
72                ctx.emit_rust_step("install vmm tests deps (windows)", move |ctx| {
73                    installed.claim(ctx);
74                    let write_commands = write_commands.claim(ctx);
75
76                    move |rt| {
77                        let sh = xshell::Shell::new()?;
78                        let mut commands = Vec::new();
79
80                        if !matches!(rt.platform(), FlowPlatform::Windows)
81                            && !flowey_lib_common::_util::running_in_wsl(rt)
82                        {
83                            anyhow::bail!("Must be on Windows or WSL2 to install Windows deps.")
84                        }
85
86                        // TODO: add these features and reg keys to the initial CI image
87
88                        // Select required features
89                        let mut features_to_enable = BTreeSet::new();
90                        if hyperv {
91                            features_to_enable.append(&mut HYPERV_TESTS_REQUIRED_FEATURES.into());
92                        }
93                        if whp {
94                            features_to_enable.append(&mut WHP_TESTS_REQUIRED_FEATURES.into());
95                        }
96
97                        // Check if features are already enabled
98                        if installing && !features_to_enable.is_empty() {
99                            let features = xshell::cmd!(sh, "DISM.exe /Online /Get-Features").output()?;
100                            assert!(features.status.success());
101                            let features = String::from_utf8_lossy(&features.stdout).to_string();
102                            let mut feature = None;
103                            for line in features.lines() {
104                                if let Some((k, v)) = line.split_once(":") {
105                                    if let Some(f) = feature {
106                                        assert_eq!(k.trim(), "State");
107                                        match v.trim() {
108                                            "Enabled" => {
109                                                assert!(features_to_enable.remove(f));
110                                            }
111                                            "Disabled" => {}
112                                            _ => anyhow::bail!("Unknown feature enablement state"),
113                                        }
114                                        feature = None;
115                                    } else if k.trim() == "Feature Name" {
116                                        let new_feature = v.trim();
117                                        feature = features_to_enable.contains(new_feature).then_some(new_feature);
118                                    }
119                                }
120                            }
121                        }
122
123                        // Prompt before enabling when running locally
124                        if installing && !features_to_enable.is_empty() && matches!(rt.backend(), FlowBackend::Local) {
125                            let mut features_to_install_string = String::new();
126                            for feature in features_to_enable.iter() {
127                                features_to_install_string.push_str(feature);
128                                features_to_install_string.push('\n');
129                            }
130
131                            log::warn!(
132                                r#"
133================================================================================
134To run the VMM tests, the following features need to be enabled:
135{features_to_install_string}
136
137You may need to restart your system for the changes to take effect.
138
139If you're OK with installing these features, please press <enter>.
140Otherwise, press `ctrl-c` to cancel the run.
141================================================================================
142"#
143                            );
144                            let _ = std::io::stdin().read_line(&mut String::new());
145                        }
146
147                        // Install the features
148                        for feature in features_to_enable {
149                            if installing {
150                                xshell::cmd!(sh, "DISM.exe /Online /NoRestart /Enable-Feature /All /FeatureName:{feature}").run()?;
151                            }
152                            commands.push(format!("DISM.exe /Online /NoRestart /Enable-Feature /All /FeatureName:{feature}"));
153                        }
154
155                        // Select required reg keys
156                        let mut reg_keys_to_set = BTreeSet::new();
157                        if hyperv {
158                            // Allow loading IGVM from file (to run custom OpenHCL firmware)
159                            reg_keys_to_set.insert("AllowFirmwareLoadFromFile");
160                            // Enable COM3 and COM4 for Hyper-V VMs so we can get the OpenHCL KMSG logs over serial
161                            reg_keys_to_set.insert("EnableAdditionalComPorts");
162
163                            if hardware_isolation {
164                                reg_keys_to_set.insert("EnableHardwareIsolation");
165                            }
166                        }
167
168                        // Check if reg keys are set
169                        if installing && !reg_keys_to_set.is_empty() {
170                            let output = xshell::cmd!(sh, "reg.exe query {VIRT_REG_PATH}").output()?;
171                            if output.status.success() {
172                                let output = String::from_utf8_lossy(&output.stdout).to_string();
173                                for line in output.lines() {
174                                    let components = line.split_whitespace().collect::<Vec<_>>();
175                                    if components.len() == 3
176                                        && reg_keys_to_set.contains(components[0])
177                                        && components[1] == "REG_DWORD"
178                                        && components[2] == "0x1"
179                                    {
180                                        assert!(reg_keys_to_set.remove(components[0]));
181                                    }
182                                }
183                            }
184                        }
185
186                        // Prompt before changing registry when running locally
187                        if installing && !reg_keys_to_set.is_empty() && matches!(rt.backend(), FlowBackend::Local) {
188                            let mut reg_keys_to_set_string = String::new();
189                            for feature in reg_keys_to_set.iter() {
190                                reg_keys_to_set_string.push_str(feature);
191                                reg_keys_to_set_string.push('\n');
192                            }
193
194                            log::warn!(
195                                r#"
196================================================================================
197To run the VMM tests, the following registry keys need to be set to 1:
198{reg_keys_to_set_string}
199
200If you're OK with changing the registry, please press <enter>.
201Otherwise, press `ctrl-c` to cancel the run.
202================================================================================
203"#
204                            );
205                            let _ = std::io::stdin().read_line(&mut String::new());
206                        }
207
208                        // Modify the registry
209                        for v in reg_keys_to_set {
210                            // TODO: figure out why reg.exe is not found if I
211                            // render the command as a string first and share
212                            if installing {
213                                xshell::cmd!(sh, "reg.exe add {VIRT_REG_PATH} /v {v} /t REG_DWORD /d 1 /f").run()?;
214                            }
215                            commands.push(format!("reg.exe add \"{VIRT_REG_PATH}\" /v {v} /t REG_DWORD /d 1 /f"));
216                        }
217
218                        for write_cmds in write_commands {
219                            rt.write(write_cmds, &commands);
220                        }
221
222                        Ok(())
223                    }
224                });
225            }
226            VmmTestsDepSelections::Linux => {
227                ctx.emit_rust_step("install vmm tests deps (linux)", |ctx| {
228                    installed.claim(ctx);
229                    let write_commands = write_commands.claim(ctx);
230
231                    |rt| {
232                        for write_cmds in write_commands {
233                            rt.write(write_cmds, &Vec::new());
234                        }
235
236                        Ok(())
237                    }
238                });
239            }
240        }
241
242        Ok(())
243    }
244}
flowey_lib_hvlite/install_vmm_tests_deps.rs

flowey_lib_hvlite/
install_vmm_tests_deps.rs
