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