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