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                        // Write kernel paths for all kinds matching this arch
201                        for kind in [
202                            OpenhclKernelPackageKind::Main,
203                            OpenhclKernelPackageKind::Dev,
204                            OpenhclKernelPackageKind::Cvm,
205                            OpenhclKernelPackageKind::CvmDev,
206                        ] {
207                            if let Some(vars) = kernel_reqs.remove(&(kind, arch)) {
208                                rt.write_all(vars, kernel_path);
209                            }
210                        }
211
212                        // Write modules paths for all kinds matching this arch
213                        for kind in [
214                            OpenhclKernelPackageKind::Main,
215                            OpenhclKernelPackageKind::Dev,
216                            OpenhclKernelPackageKind::Cvm,
217                            OpenhclKernelPackageKind::CvmDev,
218                        ] {
219                            if let Some(vars) = modules_reqs.remove(&(kind, arch)) {
220                                rt.write_all(vars, modules_path);
221                            }
222                        }
223
224                        // Write package root paths (parent of kernel)
225                        if let Some(parent) = kernel_path.parent() {
226                            let parent_buf = parent.to_path_buf();
227                            for kind in [
228                                OpenhclKernelPackageKind::Main,
229                                OpenhclKernelPackageKind::Dev,
230                                OpenhclKernelPackageKind::Cvm,
231                                OpenhclKernelPackageKind::CvmDev,
232                            ] {
233                                if let Some(vars) = pkg_reqs.remove(&(kind, arch)) {
234                                    rt.write_all(vars, &parent_buf);
235                                }
236                            }
237
238                            // Write metadata paths (kernel_build_metadata.json in same dir as kernel)
239                            let metadata_path = parent_buf.join("kernel_build_metadata.json");
240                            for kind in [
241                                OpenhclKernelPackageKind::Main,
242                                OpenhclKernelPackageKind::Dev,
243                                OpenhclKernelPackageKind::Cvm,
244                                OpenhclKernelPackageKind::CvmDev,
245                            ] {
246                                if let Some(vars) = metadata_reqs.remove(&(kind, arch)) {
247                                    rt.write_all(vars, &metadata_path);
248                                }
249                            }
250                        }
251                    }
252                    Ok(())
253                }
254            });
255        }
256
257        if download_reqs.is_empty() {
258            return Ok(());
259        }
260
261        // Handle downloads
262        let extract_zip_deps = flowey_lib_common::_util::extract::extract_zip_if_new_deps(ctx);
263
264        for (kind, arch) in download_reqs {
265            let version = versions.get(&kind).expect("checked above");
266            let tag = format!(
267                "rolling-lts/hcl-{}/{}",
268                match kind {
269                    OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "main",
270                    OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => "dev",
271                },
272                version
273            );
274
275            let file_name = format!(
276                "Microsoft.OHCL.Kernel{}.{}{}-{}.tar.gz",
277                match kind {
278                    OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "",
279                    OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => ".Dev",
280                },
281                version,
282                match kind {
283                    OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Dev => "",
284                    OpenhclKernelPackageKind::Cvm | OpenhclKernelPackageKind::CvmDev => "-cvm",
285                },
286                match arch {
287                    OpenhclKernelPackageArch::X86_64 => "x64",
288                    OpenhclKernelPackageArch::Aarch64 => "arm64",
289                },
290            );
291
292            let kernel_package_tar_gz =
293                ctx.reqv(|v| flowey_lib_common::download_gh_release::Request {
294                    repo_owner: "microsoft".into(),
295                    repo_name: "OHCL-Linux-Kernel".into(),
296                    needs_auth: false,
297                    tag,
298                    file_name: file_name.clone(),
299                    path: v,
300                });
301
302            let kernel_file_name = match arch {
303                OpenhclKernelPackageArch::X86_64 => "vmlinux",
304                OpenhclKernelPackageArch::Aarch64 => "Image",
305            };
306
307            let has_kernel_req = kernel_reqs_download.contains_key(&(kind, arch));
308            let has_modules_req = modules_reqs_download.contains_key(&(kind, arch));
309            let has_pkg_req = pkg_reqs_download.contains_key(&(kind, arch));
310            let has_metadata_req = metadata_reqs_download.contains_key(&(kind, arch));
311
312            ctx.emit_rust_step("extract and resolve kernel package", |ctx| {
313                let extract_zip_deps = extract_zip_deps.clone().claim(ctx);
314                let kernel_vars = if has_kernel_req {
315                    Some(
316                        kernel_reqs_download
317                            .remove(&(kind, arch))
318                            .unwrap()
319                            .claim(ctx),
320                    )
321                } else {
322                    None
323                };
324                let modules_vars = if has_modules_req {
325                    Some(
326                        modules_reqs_download
327                            .remove(&(kind, arch))
328                            .unwrap()
329                            .claim(ctx),
330                    )
331                } else {
332                    None
333                };
334                let pkg_vars = if has_pkg_req {
335                    Some(pkg_reqs_download.remove(&(kind, arch)).unwrap().claim(ctx))
336                } else {
337                    None
338                };
339                let metadata_vars = if has_metadata_req {
340                    Some(
341                        metadata_reqs_download
342                            .remove(&(kind, arch))
343                            .unwrap()
344                            .claim(ctx),
345                    )
346                } else {
347                    None
348                };
349                let kernel_package_tar_gz = kernel_package_tar_gz.claim(ctx);
350                let file_name = file_name.clone();
351                let kernel_file_name = kernel_file_name.to_string();
352
353                move |rt| {
354                    let kernel_package_tar_gz = rt.read(kernel_package_tar_gz);
355
356                    // Extract the downloaded package
357                    let extract_dir = flowey_lib_common::_util::extract::extract_zip_if_new(
358                        rt,
359                        extract_zip_deps,
360                        &kernel_package_tar_gz,
361                        &file_name,
362                    )?;
363
364                    // The extracted directory contains: vmlinux/Image, modules/, kernel_build_metadata.json
365                    let kernel_path = extract_dir.join(&kernel_file_name);
366                    let modules_path = extract_dir.join("modules");
367                    let metadata_path = extract_dir.join("kernel_build_metadata.json");
368
369                    if let Some(vars) = kernel_vars {
370                        rt.write_all(vars, &kernel_path);
371                    }
372                    if let Some(vars) = modules_vars {
373                        rt.write_all(vars, &modules_path);
374                    }
375                    if let Some(vars) = pkg_vars {
376                        rt.write_all(vars, &extract_dir);
377                    }
378                    if let Some(vars) = metadata_vars {
379                        rt.write_all(vars, &metadata_path);
380                    }
381
382                    Ok(())
383                }
384            });
385        }
386
387        Ok(())
388    }
389}