Skip to main content

flowey_lib_hvlite/
resolve_openhcl_kernel_package.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Resolve OpenHCL kernel packages - either by downloading from GitHub Release
5//! or using local paths
6
7use crate::common::CommonArch;
8use flowey::node::prelude::*;
9use std::collections::BTreeMap;
10
11#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Copy)]
12pub enum OpenhclKernelPackageKind {
13    Main,
14    Cvm,
15    Dev,
16    CvmDev,
17}
18
19impl OpenhclKernelPackageKind {
20    pub fn is_dev(self) -> bool {
21        match self {
22            Self::Main | Self::Cvm => false,
23            Self::Dev | Self::CvmDev => true,
24        }
25    }
26}
27
28flowey_config! {
29    /// Config for the resolve_openhcl_kernel_package node.
30    pub struct Config {
31        /// Version strings keyed by package kind.
32        pub versions: BTreeMap<OpenhclKernelPackageKind, String>,
33        /// Local paths keyed by architecture (kernel binary, modules directory).
34        pub local_paths: BTreeMap<CommonArch, (ConfigVar<PathBuf>, ConfigVar<PathBuf>)>,
35    }
36}
37
38flowey_request! {
39    #[expect(clippy::enum_variant_names)]
40    pub enum Request {
41        /// Get path to the kernel binary
42        GetKernel {
43            kind: OpenhclKernelPackageKind,
44            arch: CommonArch,
45            kernel: WriteVar<PathBuf>,
46        },
47        /// Get path to the kernel modules directory
48        GetModules {
49            kind: OpenhclKernelPackageKind,
50            arch: CommonArch,
51            modules: WriteVar<PathBuf>,
52        },
53        /// Get path to the package root (for metadata files, etc)
54        GetPackageRoot {
55            kind: OpenhclKernelPackageKind,
56            arch: CommonArch,
57            pkg: WriteVar<PathBuf>,
58        },
59        /// Get path to the kernel build metadata file
60        GetMetadata {
61            kind: OpenhclKernelPackageKind,
62            arch: CommonArch,
63            metadata: WriteVar<PathBuf>,
64        },
65    }
66}
67
68new_flow_node_with_config!(struct Node);
69
70impl FlowNodeWithConfig for Node {
71    type Request = Request;
72    type Config = Config;
73
74    fn imports(ctx: &mut ImportCtx<'_>) {
75        ctx.import::<flowey_lib_common::install_dist_pkg::Node>();
76        ctx.import::<flowey_lib_common::download_gh_release::Node>();
77    }
78
79    fn emit(
80        config: Config,
81        requests: Vec<Self::Request>,
82        ctx: &mut NodeCtx<'_>,
83    ) -> anyhow::Result<()> {
84        let versions = config.versions;
85        let local_paths = config.local_paths;
86        let mut kernel_reqs: BTreeMap<
87            (OpenhclKernelPackageKind, CommonArch),
88            Vec<WriteVar<PathBuf>>,
89        > = BTreeMap::new();
90        let mut modules_reqs: BTreeMap<
91            (OpenhclKernelPackageKind, CommonArch),
92            Vec<WriteVar<PathBuf>>,
93        > = BTreeMap::new();
94        let mut pkg_reqs: BTreeMap<(OpenhclKernelPackageKind, CommonArch), Vec<WriteVar<PathBuf>>> =
95            BTreeMap::new();
96        let mut metadata_reqs: BTreeMap<
97            (OpenhclKernelPackageKind, CommonArch),
98            Vec<WriteVar<PathBuf>>,
99        > = BTreeMap::new();
100
101        for req in requests {
102            match req {
103                Request::GetKernel { kind, arch, kernel } => {
104                    kernel_reqs.entry((kind, arch)).or_default().push(kernel);
105                }
106                Request::GetModules {
107                    kind,
108                    arch,
109                    modules,
110                } => {
111                    modules_reqs.entry((kind, arch)).or_default().push(modules);
112                }
113                Request::GetPackageRoot { kind, arch, pkg } => {
114                    pkg_reqs.entry((kind, arch)).or_default().push(pkg);
115                }
116                Request::GetMetadata {
117                    kind,
118                    arch,
119                    metadata,
120                } => {
121                    metadata_reqs
122                        .entry((kind, arch))
123                        .or_default()
124                        .push(metadata);
125                }
126            }
127        }
128
129        // Collect all architectures that need resolution
130        let all_reqs: std::collections::BTreeSet<(OpenhclKernelPackageKind, CommonArch)> =
131            kernel_reqs
132                .keys()
133                .chain(modules_reqs.keys())
134                .chain(pkg_reqs.keys())
135                .chain(metadata_reqs.keys())
136                .cloned()
137                .collect();
138
139        // Verify we have either local paths or versions for each requested architecture
140        for (kind, arch) in &all_reqs {
141            if !local_paths.contains_key(arch) && !versions.contains_key(kind) {
142                if kind.is_dev() {
143                    anyhow::bail!(
144                        "OpenHCL dev kernel support is disabled; provide local kernel paths for \
145                         {:?} to enable {:?}",
146                        arch,
147                        kind,
148                    );
149                }
150                anyhow::bail!(
151                    "Must provide either SetLocal for {:?} or SetVersion for {:?}",
152                    arch,
153                    kind
154                );
155            }
156        }
157
158        if all_reqs.is_empty() {
159            return Ok(());
160        }
161
162        // Partition requests into local vs download
163        let (local_reqs, download_reqs): (Vec<_>, Vec<_>) = all_reqs
164            .into_iter()
165            .partition(|(_, arch)| local_paths.contains_key(arch));
166
167        // Split the request maps into local and download portions
168        let (kernel_reqs_local, mut kernel_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) =
169            kernel_reqs
170                .into_iter()
171                .partition(|((_, arch), _)| local_paths.contains_key(arch));
172        let (modules_reqs_local, mut modules_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) =
173            modules_reqs
174                .into_iter()
175                .partition(|((_, arch), _)| local_paths.contains_key(arch));
176        let (pkg_reqs_local, mut pkg_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) = pkg_reqs
177            .into_iter()
178            .partition(|((_, arch), _)| local_paths.contains_key(arch));
179        let (metadata_reqs_local, mut metadata_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) =
180            metadata_reqs
181                .into_iter()
182                .partition(|((_, arch), _)| local_paths.contains_key(arch));
183
184        // Handle local paths
185        if !local_reqs.is_empty() {
186            ctx.emit_rust_step("use local kernel package", |ctx| {
187                let mut kernel_reqs = kernel_reqs_local.claim(ctx);
188                let mut modules_reqs = modules_reqs_local.claim(ctx);
189                let mut pkg_reqs = pkg_reqs_local.claim(ctx);
190                let mut metadata_reqs = metadata_reqs_local.claim(ctx);
191                let local_paths: BTreeMap<_, _> = local_paths
192                    .into_iter()
193                    .map(|(arch, (k, m))| (arch, (k.claim(ctx), m.claim(ctx))))
194                    .collect();
195                let local_reqs = local_reqs.clone();
196
197                move |rt| {
198                    for (_, arch) in local_reqs {
199                        let (kernel_var, modules_var) = local_paths.get(&arch).unwrap();
200                        let kernel_path = rt.read(kernel_var.clone());
201                        let modules_path = rt.read(modules_var.clone());
202
203                        log::info!(
204                            "using local kernel at {:?} and modules at {:?}",
205                            kernel_path,
206                            modules_path
207                        );
208
209                        // Write kernel paths for all kinds matching this arch
210                        for kind in [
211                            OpenhclKernelPackageKind::Main,
212                            OpenhclKernelPackageKind::Dev,
213                            OpenhclKernelPackageKind::Cvm,
214                            OpenhclKernelPackageKind::CvmDev,
215                        ] {
216                            if let Some(vars) = kernel_reqs.remove(&(kind, arch)) {
217                                rt.write_all(vars, &kernel_path);
218                            }
219                        }
220
221                        // Write modules paths for all kinds matching this arch
222                        for kind in [
223                            OpenhclKernelPackageKind::Main,
224                            OpenhclKernelPackageKind::Dev,
225                            OpenhclKernelPackageKind::Cvm,
226                            OpenhclKernelPackageKind::CvmDev,
227                        ] {
228                            if let Some(vars) = modules_reqs.remove(&(kind, arch)) {
229                                rt.write_all(vars, &modules_path);
230                            }
231                        }
232
233                        // Write package root paths (parent of kernel)
234                        if let Some(parent) = kernel_path.parent() {
235                            let parent_buf = parent.to_path_buf();
236                            for kind in [
237                                OpenhclKernelPackageKind::Main,
238                                OpenhclKernelPackageKind::Dev,
239                                OpenhclKernelPackageKind::Cvm,
240                                OpenhclKernelPackageKind::CvmDev,
241                            ] {
242                                if let Some(vars) = pkg_reqs.remove(&(kind, arch)) {
243                                    rt.write_all(vars, &parent_buf);
244                                }
245                            }
246
247                            // Write metadata paths (kernel_build_metadata.json in same dir as kernel)
248                            let metadata_path = parent_buf.join("kernel_build_metadata.json");
249                            for kind in [
250                                OpenhclKernelPackageKind::Main,
251                                OpenhclKernelPackageKind::Dev,
252                                OpenhclKernelPackageKind::Cvm,
253                                OpenhclKernelPackageKind::CvmDev,
254                            ] {
255                                if let Some(vars) = metadata_reqs.remove(&(kind, arch)) {
256                                    rt.write_all(vars, &metadata_path);
257                                }
258                            }
259                        }
260                    }
261                    Ok(())
262                }
263            });
264        }
265
266        if download_reqs.is_empty() {
267            return Ok(());
268        }
269
270        // Handle downloads
271        let extract_zip_deps = flowey_lib_common::_util::extract::extract_zip_if_new_deps(ctx);
272
273        for (kind, arch) in download_reqs {
274            let version = versions.get(&kind).expect("checked above");
275            let tag = format!(
276                "rolling-lts/hcl-{}/{}",
277                match kind {
278                    OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "main",
279                    OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => "dev",
280                },
281                version
282            );
283
284            let file_name = format!(
285                "Microsoft.OHCL.Kernel{}.{}{}-{}.tar.gz",
286                match kind {
287                    OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "",
288                    OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => ".Dev",
289                },
290                version,
291                match kind {
292                    OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Dev => "",
293                    OpenhclKernelPackageKind::Cvm | OpenhclKernelPackageKind::CvmDev => "-cvm",
294                },
295                match arch {
296                    CommonArch::X86_64 => "x64",
297                    CommonArch::Aarch64 => "arm64",
298                },
299            );
300
301            let kernel_package_tar_gz =
302                ctx.reqv(|v| flowey_lib_common::download_gh_release::Request {
303                    repo_owner: "microsoft".into(),
304                    repo_name: "OHCL-Linux-Kernel".into(),
305                    needs_auth: false,
306                    tag,
307                    file_name: file_name.clone(),
308                    path: v,
309                });
310
311            let kernel_file_name = match arch {
312                CommonArch::X86_64 => "vmlinux",
313                CommonArch::Aarch64 => "Image",
314            };
315
316            let has_kernel_req = kernel_reqs_download.contains_key(&(kind, arch));
317            let has_modules_req = modules_reqs_download.contains_key(&(kind, arch));
318            let has_pkg_req = pkg_reqs_download.contains_key(&(kind, arch));
319            let has_metadata_req = metadata_reqs_download.contains_key(&(kind, arch));
320
321            ctx.emit_rust_step("extract and resolve kernel package", |ctx| {
322                let extract_zip_deps = extract_zip_deps.clone().claim(ctx);
323                let kernel_vars = if has_kernel_req {
324                    Some(
325                        kernel_reqs_download
326                            .remove(&(kind, arch))
327                            .unwrap()
328                            .claim(ctx),
329                    )
330                } else {
331                    None
332                };
333                let modules_vars = if has_modules_req {
334                    Some(
335                        modules_reqs_download
336                            .remove(&(kind, arch))
337                            .unwrap()
338                            .claim(ctx),
339                    )
340                } else {
341                    None
342                };
343                let pkg_vars = if has_pkg_req {
344                    Some(pkg_reqs_download.remove(&(kind, arch)).unwrap().claim(ctx))
345                } else {
346                    None
347                };
348                let metadata_vars = if has_metadata_req {
349                    Some(
350                        metadata_reqs_download
351                            .remove(&(kind, arch))
352                            .unwrap()
353                            .claim(ctx),
354                    )
355                } else {
356                    None
357                };
358                let kernel_package_tar_gz = kernel_package_tar_gz.claim(ctx);
359                let file_name = file_name.clone();
360                let kernel_file_name = kernel_file_name.to_string();
361
362                move |rt| {
363                    let kernel_package_tar_gz = rt.read(kernel_package_tar_gz);
364
365                    // Extract the downloaded package
366                    let extract_dir = flowey_lib_common::_util::extract::extract_zip_if_new(
367                        rt,
368                        extract_zip_deps,
369                        &kernel_package_tar_gz,
370                        &file_name,
371                    )?;
372
373                    // The extracted directory contains: vmlinux/Image, modules/, kernel_build_metadata.json
374                    let kernel_path = extract_dir.join(&kernel_file_name);
375                    let modules_path = extract_dir.join("modules");
376                    let metadata_path = extract_dir.join("kernel_build_metadata.json");
377
378                    if let Some(vars) = kernel_vars {
379                        rt.write_all(vars, &kernel_path);
380                    }
381                    if let Some(vars) = modules_vars {
382                        rt.write_all(vars, &modules_path);
383                    }
384                    if let Some(vars) = pkg_vars {
385                        rt.write_all(vars, &extract_dir);
386                    }
387                    if let Some(vars) = metadata_vars {
388                        rt.write_all(vars, &metadata_path);
389                    }
390
391                    Ok(())
392                }
393            });
394        }
395
396        Ok(())
397    }
398}