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