Skip to main content

scx_cargo/
bpf_builder.rs

1// Copyright (c) Meta Platforms, Inc. and affiliates.
2//
3// This software may be used and distributed according to the terms of the
4// GNU General Public License version 2.
5
6use crate::ClangInfo;
7
8use anyhow::{Context, Result, anyhow, bail};
9use glob::glob;
10use libbpf_cargo::SkeletonBuilder;
11use libbpf_rs::Linker;
12use libbpf_rs::btf::{Btf, BtfKind, BtfType, types::DataSec};
13use std::collections::BTreeSet;
14use std::collections::HashMap;
15use std::env;
16use std::fs;
17use std::path::Path;
18use std::path::PathBuf;
19use tracing::Level;
20use tracing_subscriber::{Layer, filter, layer::SubscriberExt};
21
22#[derive(Debug)]
23/// # Build helpers for sched_ext schedulers with Rust userspace component
24///
25/// This is to be used from `build.rs` of a cargo project which implements a
26/// [sched_ext](https://github.com/sched-ext/scx) scheduler with C BPF
27/// component and Rust userspace component. `BpfBuilder` provides everything
28/// necessary to build the BPF component and generate Rust bindings.
29/// BpfBuilder provides the following.
30///
31/// 1. *`vmlinux.h` and other common BPF header files*
32///
33/// All sched_ext BPF implementations require `vmlinux.h` and many make use
34/// of common constructs such as
35/// [`user_exit_info`](https://github.com/sched-ext/scx/blob/main/scheds/include/scx/user_exit_info.h).
36/// `BpfBuilder` makes these headers available when compiling BPF source
37/// code and generating bindings for it. The included headers can be browsed
38/// at <https://github.com/sched-ext/scx/tree/main/scheds/include>.
39///
40/// These headers can be superseded using environment variables which will
41/// be discussed later.
42///
43/// 2. *Header bindings using `bindgen`*
44///
45/// If enabled with `.enable_intf()`, the input `.h` file is processed by
46/// `bindgen` to generate Rust bindings. This is useful in establishing
47/// shared constants and data types between the BPF and user components.
48///
49/// Note that the types generated with `bindgen` are different from the
50/// types used by the BPF skeleton even when they are the same types in BPF.
51/// This is a source of ugliness and we are hoping to address it by
52/// improving `libbpf-cargo` in the future.
53///
54/// 3. *BPF compilation and generation of the skeleton and its bindings*
55///
56/// If enabled with `.enable_skel()`, the input `.bpf.c` file is compiled
57/// and its skeleton and bindings are generated using `libbpf-cargo`.
58///
59/// ## An Example
60///
61/// This section shows how `BpfBuilder` can be used in an example project.
62/// For a concrete example, take a look at
63/// [`scx_rusty`](https://github.com/sched-ext/scx/tree/main/scheds/rust/scx_rusty).
64///
65/// A minimal source tree using all features would look like the following:
66///
67/// ```text
68/// scx_hello_world
69/// |-- Cargo.toml
70/// |-- build.rs
71/// \-- src
72///     |-- main.rs
73///     |-- bpf_intf.rs
74///     |-- bpf_skel.rs
75///     \-- bpf
76///         |-- intf.h
77///         \-- main.c
78/// ```
79///
80/// The following three files would contain the actual implementation:
81///
82/// - `src/main.rs`: Rust userspace component which loads the BPF blob and
83/// interacts it using the generated bindings.
84///
85/// - `src/bpf/intf.h`: C header file defining constants and structs
86/// that will be used by both the BPF and userspace components.
87///
88/// - `src/bpf/main.c`: C source code implementing the BPF component -
89/// including `struct sched_ext_ops`.
90///
91/// And then there are boilerplates to generate the bindings and make them
92/// available as modules to `main.rs`.
93///
94/// - `Cargo.toml`: Includes `scx_cargo` in the `[build-dependencies]`
95/// section.
96///
97/// - `build.rs`: Uses `scx_cargo::BpfBuilder` to build and generate
98/// bindings for the BPF component. For this project, it can look like the
99/// following.
100///
101/// ```should_panic
102/// fn main() {
103///     scx_cargo::BpfBuilder::new()
104///         .unwrap()
105///         .enable_intf("src/bpf/intf.h", "bpf_intf.rs")
106///         .enable_skel("src/bpf/main.bpf.c", "bpf")
107///         .build()
108///         .unwrap();
109/// }
110/// ```
111///
112/// - `bpf_intf.rs`: Import the bindings generated by `bindgen` into a
113/// module. Above, we told `.enable_intf()` to generate the bindings into
114/// `bpf_intf.rs`, so the file would look like the following. The `allow`
115/// directives are useful if the header is including `vmlinux.h`.
116///
117/// ```ignore
118/// #![allow(non_upper_case_globals)]
119/// #![allow(non_camel_case_types)]
120/// #![allow(non_snake_case)]
121/// #![allow(dead_code)]
122///
123/// include!(concat!(env!("OUT_DIR"), "/bpf_intf.rs"));
124/// ```
125///
126/// - `bpf_skel.rs`: Import the BPF skeleton bindings generated by
127/// `libbpf-cargo` into a module. Above, we told `.enable_skel()` to use the
128/// skeleton name `bpf`, so the file would look like the following.
129///
130/// ```ignore
131/// include!(concat!(env!("OUT_DIR"), "/bpf_skel.rs"));
132/// ```
133///
134/// ## Compiler Flags and Environment Variables
135///
136/// BPF being its own CPU architecture and independent runtime environment,
137/// build environment and steps are already rather complex. The need to
138/// interface between two different languages - C and Rust - adds further
139/// complexities. `BpfBuilder` automates most of the process. The determined
140/// build environment is recorded in the `build.rs` output and can be
141/// obtained with a command like the following:
142///
143/// ```text
144/// $ grep '^scx_utils:clang=' target/release/build/scx_rusty-*/output
145/// ```
146///
147/// While the automatic settings should work most of the time, there can be
148/// times when overriding them is necessary. The following environment
149/// variables can be used to customize the build environment.
150///
151/// - `BPF_CLANG`: The clang command to use. (Default: `clang`)
152///
153/// - `BPF_CFLAGS`: Compiler flags to use when building BPF source code. If
154///   specified, the flags from this variable are the only flags passed to
155///   the compiler. `BpfBuilder` won't generate any flags including `-I`
156///   flags for the common header files and other `CFLAGS` related variables
157///   are ignored.
158///
159/// - `BPF_BASE_CFLAGS`: Override the non-include part of cflags.
160///
161/// - `BPF_EXTRA_CFLAGS_PRE_INCL`: Add cflags before the automatic include
162///   search path options. Header files in the search paths added by this
163///   variable will supersede the automatic ones.
164///
165/// - `BPF_EXTRA_CFLAGS_POST_INCL`: Add cflags after the automatic include
166///   search path options. Header paths added by this variable will be
167///   searched only if the target header file can't be found in the
168///   automatic header paths.
169///
170/// - `RUSTFLAGS`: This is a generic `cargo` flag and can be useful for
171///   specifying extra linker flags.
172///
173/// - `SCX_ALLOW_OLD_CLANG`: Build a cid-form scheduler with a clang older
174///   than 22 anyway. clang before 22 mishandles arena memory, so such a build
175///   fails by default; setting this turns the failure into a warning.
176///
177/// A common case for using the above flags is using the latest `libbpf`
178/// from the kernel tree. Let's say the kernel tree is at `$KERNEL` and
179/// `libbpf`. The following builds `libbpf` shipped with the kernel:
180///
181/// ```test
182/// $ cd $KERNEL
183/// $ make -C tools/bpf/bpftool
184/// ```
185///
186/// To link the scheduler against the resulting `libbpf`:
187///
188/// ```test
189/// $ env BPF_EXTRA_CFLAGS_POST_INCL=$KERNEL/tools/bpf/bpftool/libbpf/include \
190///   RUSTFLAGS="-C link-args=-lelf -C link-args=-lz -C link-args=-lzstd \
191///   -L$KERNEL/tools/bpf/bpftool/libbpf" cargo build --release
192/// ```
193pub struct BpfBuilder {
194    clang: ClangInfo,
195    cflags: Vec<String>,
196    out_dir: PathBuf,
197    sources: BTreeSet<String>,
198
199    intf_input_output: Option<(String, String)>,
200    skel_input_name: Option<(String, String)>,
201}
202
203impl BpfBuilder {
204    const BPF_H_TAR: &'static [u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bpf_h.tar"));
205
206    fn install_bpf_h<P: AsRef<Path>>(dest: P) -> Result<()> {
207        let mut ar = tar::Archive::new(Self::BPF_H_TAR);
208        ar.unpack(dest)?;
209        Ok(())
210    }
211
212    fn determine_cflags<P>(clang: &ClangInfo, out_dir: P) -> Result<Vec<String>>
213    where
214        P: AsRef<Path> + std::fmt::Debug,
215    {
216        let bpf_h = out_dir
217            .as_ref()
218            .join("scx_utils-bpf_h")
219            .to_str()
220            .ok_or(anyhow!(
221                "{:?}/scx_utils-bph_h can't be converted to str",
222                out_dir
223            ))?
224            .to_string();
225        Self::install_bpf_h(&bpf_h)?;
226
227        let mut cflags = Vec::<String>::new();
228
229        cflags.append(&mut match env::var("BPF_BASE_CFLAGS") {
230            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
231            _ => clang.determine_base_cflags()?,
232        });
233
234        cflags.append(&mut match env::var("BPF_EXTRA_CFLAGS_PRE_INCL") {
235            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
236            _ => vec![],
237        });
238
239        cflags.push(format!(
240            "-I{}/arch/{}",
241            bpf_h,
242            clang.kernel_target().unwrap()
243        ));
244        cflags.push(format!("-I{}", bpf_h));
245        cflags.push(format!("-I{}/bpf-compat", bpf_h));
246
247        cflags.append(&mut match env::var("BPF_EXTRA_CFLAGS_POST_INCL") {
248            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
249            _ => vec![],
250        });
251
252        Ok(cflags)
253    }
254
255    /// Create a new `BpfBuilder` struct. Call `enable` and `set` methods to
256    /// configure and `build` method to compile and generate bindings. See
257    /// the struct documentation for details.
258    pub fn new() -> Result<Self> {
259        let out_dir = PathBuf::from(env::var("OUT_DIR")?);
260
261        let clang = ClangInfo::new()?;
262        let mut cflags = match env::var("BPF_CFLAGS") {
263            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
264            _ => Self::determine_cflags(&clang, &out_dir)?,
265        };
266
267        // Add a target architecture definition (e.g., __SCX_TARGET_ARCH_x86)
268        // for target architecture-specific optimizations in BPF code.
269        cflags.push(format!("-D__SCX_TARGET_ARCH_{}", clang.kernel_target()?));
270
271        println!("scx_utils:clang={:?} {:?}", clang, cflags);
272
273        Ok(Self {
274            clang,
275            cflags,
276            out_dir,
277
278            sources: BTreeSet::new(),
279            intf_input_output: None,
280            skel_input_name: None,
281        })
282    }
283
284    /// Enable generation of header bindings using `bindgen`. `@input` is
285    /// the `.h` file defining the constants and types to be shared between
286    /// BPF and Rust components. `@output` is the `.rs` file to be
287    /// generated.
288    pub fn enable_intf(&mut self, input: &str, output: &str) -> &mut Self {
289        self.intf_input_output = Some((input.into(), output.into()));
290        self
291    }
292
293    /// Enable compilation of BPF code and generation of the skeleton and
294    /// its Rust bindings. `@input` is the `.bpf.c` file containing the BPF
295    /// source code and `@output` is the `.rs` file to be generated.
296    pub fn enable_skel(&mut self, input: &str, name: &str) -> &mut Self {
297        self.skel_input_name = Some((input.into(), name.into()));
298        self.sources.insert(input.into());
299
300        self
301    }
302
303    fn input_insert_deps(&self, deps: &mut BTreeSet<String>) {
304        let (input, _) = match &self.intf_input_output {
305            Some(pair) => pair,
306            None => return,
307        };
308
309        // Tell cargo to invalidate the built crate whenever the wrapper changes
310        deps.insert(input.to_string());
311    }
312
313    fn bindgen_bpf_intf(&self) -> Result<()> {
314        let (input, output) = match &self.intf_input_output {
315            Some(pair) => pair,
316            None => return Ok(()),
317        };
318
319        // The bindgen::Builder is the main entry point to bindgen, and lets
320        // you build up options for the resulting bindings.
321        let bindings = bindgen::Builder::default()
322            // Should run clang with the same -I options as BPF compilation.
323            .clang_args(
324                self.cflags
325                    .iter()
326                    .chain(["-target".into(), "bpf".into()].iter()),
327            )
328            // The input header we would like to generate bindings for.
329            .header(input)
330            // Tell cargo to invalidate the built crate whenever any of the
331            // included header files changed.
332            .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
333            .wrap_unsafe_ops(true)
334            .generate()
335            .context("Unable to generate bindings")?;
336
337        bindings
338            .write_to_file(self.out_dir.join(output))
339            .context("Couldn't write bindings")
340    }
341
342    pub fn add_source(&mut self, input: &str) -> &mut Self {
343        self.sources.insert(input.into());
344        self
345    }
346
347    /// Add an include search path that takes precedence over the built-in
348    /// header directories (vmlinux, `scheds/include`, and `bpf-compat`). The
349    /// path is prepended to the compiler flags, so it is searched first and
350    /// headers found here shadow the bundled ones. Later calls take precedence
351    /// over earlier ones.
352    pub fn add_include_path(&mut self, path: &str) -> &mut Self {
353        self.cflags.insert(0, format!("-I{path}"));
354        self
355    }
356
357    /// Add a raw flag (e.g. a `-D<name>` define) to the BPF compiler flags.
358    pub fn add_cflag(&mut self, flag: &str) -> &mut Self {
359        self.cflags.push(flag.into());
360        self
361    }
362
363    pub fn compile_link_gen(&mut self) -> Result<()> {
364        let input = match &self.skel_input_name {
365            Some((name, _)) => name,
366            None => return Ok(()),
367        };
368
369        // Better to hardcode the name because it gets embedded in
370        // all the skeleton struct/member definitions.
371        let linkobj = self.out_dir.join("bpf.bpf.o");
372        let mut linker = Linker::new(&linkobj)?;
373
374        for filename in self.sources.iter() {
375            let name = Path::new(filename).file_name().unwrap().to_str().unwrap();
376            let obj = self.out_dir.join(name.replace(".bpf.c", ".bpf.o"));
377
378            with_clang_warnings(|| {
379                SkeletonBuilder::new()
380                    .source(filename)
381                    .rustfmt("disable_rustfmt")
382                    .obj(&obj)
383                    .clang(&self.clang.clang)
384                    .clang_args(&self.cflags)
385                    .build()
386            })?;
387
388            fixup_arena_externs(&obj)?;
389            linker.add_file(&obj)?;
390        }
391
392        linker.link()?;
393        fixup_arena_datasec(&linkobj)?;
394        self.check_cid_clang(&linkobj)?;
395
396        self.bindgen_bpf_intf()?;
397
398        let skel_path = self.out_dir.join("bpf_skel.rs");
399
400        SkeletonBuilder::new()
401            .obj(&linkobj)
402            .clang(&self.clang.clang)
403            .clang_args(&self.cflags)
404            .rustfmt("disable_rustfmt")
405            .reference_obj(true)
406            .generate(&skel_path)?;
407
408        let mut deps = BTreeSet::new();
409        self.add_src_deps(&mut deps, input)?;
410        for filename in self.sources.iter() {
411            deps.insert(filename.to_string());
412        }
413
414        self.gen_cargo_reruns(Some(&deps))?;
415
416        Ok(())
417    }
418
419    fn gen_bpf_skel(&self, deps: &mut BTreeSet<String>) -> Result<()> {
420        let (input, name) = match &self.skel_input_name {
421            Some(pair) => pair,
422            None => return Ok(()),
423        };
424
425        let obj = self.out_dir.join(format!("{name}.bpf.o"));
426        let skel_path = self.out_dir.join(format!("{name}_skel.rs"));
427
428        with_clang_warnings(|| {
429            SkeletonBuilder::new()
430                .source(input)
431                .obj(&obj)
432                .clang(&self.clang.clang)
433                .clang_args(&self.cflags)
434                .rustfmt("disable_rustfmt")
435                .reference_obj(true)
436                .build_and_generate(&skel_path)
437        })?;
438
439        self.add_src_deps(deps, input)?;
440
441        Ok(())
442    }
443
444    fn add_src_deps(&self, deps: &mut BTreeSet<String>, input: &str) -> Result<()> {
445        let c_path = PathBuf::from(input);
446        let dir = c_path
447            .parent()
448            .ok_or(anyhow!("Source {:?} doesn't have parent dir", c_path))?
449            .to_str()
450            .ok_or(anyhow!("Parent dir of {:?} isn't a UTF-8 string", c_path))?;
451
452        for path in glob(&format!("{dir}/*.[hc]"))?.filter_map(Result::ok) {
453            deps.insert(
454                path.to_str()
455                    .ok_or(anyhow!("Path {:?} is not a valid string", path))?
456                    .to_string(),
457            );
458        }
459
460        Ok(())
461    }
462
463    /// clang before 22 mishandles arena memory: memset() on it is lowered
464    /// without the address space cast. cid-form schedulers are arena based, so
465    /// refuse to build one with such a clang unless SCX_ALLOW_OLD_CLANG turns
466    /// the refusal into a warning.
467    fn check_cid_clang(&self, linkobj: &Path) -> Result<()> {
468        if version_compare::compare(&self.clang.ver, "22") != Ok(version_compare::Cmp::Lt) {
469            return Ok(());
470        }
471        let Some(ops) = cid_struct_ops_type(linkobj)? else {
472            return Ok(());
473        };
474        let msg = format!(
475            "{} registers cid-form struct_ops {ops} but is built with clang {}, and \
476             clang before 22 mishandles arena memory (memset() on it is lowered without \
477             the address space cast), which cid-form schedulers depend on. Use clang >= 22, \
478             or set SCX_ALLOW_OLD_CLANG=1 to build anyway.",
479            env::var("CARGO_PKG_NAME").unwrap_or_default(),
480            self.clang.ver
481        );
482        if env::var_os("SCX_ALLOW_OLD_CLANG").is_some_and(|v| !v.is_empty()) {
483            println!("cargo:warning={msg}");
484            return Ok(());
485        }
486        bail!(msg)
487    }
488
489    fn gen_cargo_reruns(&self, dependencies: Option<&BTreeSet<String>>) -> Result<()> {
490        println!("cargo:rerun-if-env-changed=BPF_CLANG");
491        println!("cargo:rerun-if-env-changed=SCX_ALLOW_OLD_CLANG");
492        println!("cargo:rerun-if-env-changed=BPF_CFLAGS");
493        println!("cargo:rerun-if-env-changed=BPF_BASE_CFLAGS");
494        println!("cargo:rerun-if-env-changed=BPF_EXTRA_CFLAGS_PRE_INCL");
495        println!("cargo:rerun-if-env-changed=BPF_EXTRA_CFLAGS_POST_INCL");
496        if let Some(deps) = dependencies {
497            for dep in deps.iter() {
498                println!("cargo:rerun-if-changed={dep}");
499            }
500        };
501
502        for source in self.sources.iter() {
503            println!("cargo:rerun-if-changed={source}");
504        }
505
506        Ok(())
507    }
508
509    /// Build and generate the enabled bindings.
510    pub fn build(&self) -> Result<()> {
511        let mut deps = BTreeSet::new();
512
513        self.input_insert_deps(&mut deps);
514
515        self.bindgen_bpf_intf()?;
516        self.gen_bpf_skel(&mut deps)?;
517        self.gen_cargo_reruns(Some(&deps))?;
518        Ok(())
519    }
520}
521
522/// The cid-form struct_ops type the linked object registers, if any. The
523/// struct_ops DATASECs list the ops variables; a cid-form one has type struct
524/// sched_ext_ops_cid or a CO-RE variant of it.
525fn cid_struct_ops_type(linkobj: &Path) -> Result<Option<String>> {
526    let btf = Btf::from_path(linkobj)
527        .with_context(|| format!("Failed to parse BTF from {}", linkobj.display()))?;
528
529    for sec in [".struct_ops", ".struct_ops.link"] {
530        let Some(datasec) = btf.type_by_name::<DataSec>(sec) else {
531            continue;
532        };
533        for vsi in datasec.iter() {
534            let Some(ty) = btf
535                .type_by_id::<BtfType>(vsi.ty)
536                .and_then(|var| var.next_type())
537                .map(|ty| ty.skip_mods_and_typedefs())
538            else {
539                continue;
540            };
541            if ty.kind() != BtfKind::Struct {
542                continue;
543            }
544            let name = ty.name().unwrap_or_default().to_string_lossy();
545            if name == "sched_ext_ops_cid" || name.starts_with("sched_ext_ops_cid___") {
546                return Ok(Some(name.into_owned()));
547            }
548        }
549    }
550    Ok(None)
551}
552
553// libbpf's static linker (libbpf 1.7 and 1.8) mishandles an arena global that
554// one compilation unit defines and another declares extern, in two ways. An
555// unresolved extern comes out of a link bound to the output's .addr_space.1
556// section as a size-zero NOTYPE symbol instead of staying undefined, so the
557// next link takes it for a definition. libbpf-cargo's build() links every
558// per-source object once, so the bound externs reach the final link here.
559// Separately, the merged .addr_space.1 BTF DATASEC keeps the extern's
560// placeholder offset when the declaring unit is linked before the defining
561// one, which the skeleton generators reject. The ELF symbols are correct in
562// both cases, so the two fixups below restore the objects from them. Both
563// edit bytes in place and never change sizes. Drop them once libbpf carries
564// the fix.
565
566const ARENA_SEC: &str = ".addr_space.1";
567const SHT_SYMTAB: u32 = 2;
568const STT_NOTYPE: u8 = 0;
569const STT_OBJECT: u8 = 1;
570const BTF_MAGIC: u16 = 0xeB9F;
571const BTF_KIND_VAR: u32 = 14;
572const BTF_KIND_DATASEC: u32 = 15;
573
574fn le16(b: &[u8], o: usize) -> u16 {
575    u16::from_le_bytes(b[o..o + 2].try_into().unwrap())
576}
577
578fn le32(b: &[u8], o: usize) -> u32 {
579    u32::from_le_bytes(b[o..o + 4].try_into().unwrap())
580}
581
582fn le64(b: &[u8], o: usize) -> u64 {
583    u64::from_le_bytes(b[o..o + 8].try_into().unwrap())
584}
585
586fn cstr(b: &[u8], o: usize) -> String {
587    let len = b[o..].iter().position(|&c| c == 0).unwrap_or(0);
588    String::from_utf8_lossy(&b[o..o + len]).into_owned()
589}
590
591/// The pieces of a little-endian ELF64 object the arena fixups touch.
592struct ArenaElf {
593    arena_shndx: u16,
594    symtab: (usize, usize),
595    strtab: usize,
596    btf: Option<(usize, usize)>,
597}
598
599impl ArenaElf {
600    /// Returns None when the object has no arena data section.
601    fn parse(elf: &[u8], path: &Path) -> Result<Option<Self>> {
602        if elf.len() < 0x40 || &elf[0..4] != b"\x7fELF" || elf[4] != 2 || elf[5] != 1 {
603            bail!("{}: not a little-endian ELF64 object", path.display());
604        }
605        let shoff = le64(elf, 0x28) as usize;
606        let shentsize = le16(elf, 0x3a) as usize;
607        let shnum = le16(elf, 0x3c) as usize;
608        let shdr = |i: usize| shoff + i * shentsize;
609        let shstr = le64(elf, shdr(le16(elf, 0x3e) as usize) + 24) as usize;
610        let (mut arena_shndx, mut symtab, mut btf) = (None, None, None);
611
612        for i in 0..shnum {
613            let h = shdr(i);
614            let name = cstr(elf, shstr + le32(elf, h) as usize);
615            let (off, size) = (le64(elf, h + 24) as usize, le64(elf, h + 32) as usize);
616
617            if name == ARENA_SEC {
618                arena_shndx = Some(i as u16);
619            } else if name == ".BTF" {
620                btf = Some((off, size));
621            } else if le32(elf, h + 4) == SHT_SYMTAB {
622                let strtab = le64(elf, shdr(le32(elf, h + 40) as usize) + 24) as usize;
623                symtab = Some(((off, size), strtab));
624            }
625        }
626        let (Some(arena_shndx), Some((symtab, strtab))) = (arena_shndx, symtab) else {
627            return Ok(None);
628        };
629        Ok(Some(Self {
630            arena_shndx,
631            symtab,
632            strtab,
633            btf,
634        }))
635    }
636
637    /// File offsets of the 24-byte symbol table entries.
638    fn syms(&self) -> impl Iterator<Item = usize> {
639        (self.symtab.0..self.symtab.0 + self.symtab.1).step_by(24)
640    }
641}
642
643/// A libbpf linker pass binds unresolved arena externs to the arena section
644/// as size-zero NOTYPE symbols, with the declaring object's section offset as
645/// their value. Mark them undefined again so the next pass treats them as
646/// externs.
647fn fixup_arena_externs(path: &Path) -> Result<()> {
648    let mut elf = fs::read(path)?;
649    let Some(parts) = ArenaElf::parse(&elf, path)? else {
650        return Ok(());
651    };
652    let mut changed = false;
653
654    for o in parts.syms() {
655        if le16(&elf, o + 6) != parts.arena_shndx
656            || elf[o + 4] & 0xf != STT_NOTYPE
657            || le64(&elf, o + 16) != 0
658        {
659            continue;
660        }
661        elf[o + 6..o + 8].copy_from_slice(&0u16.to_le_bytes());
662        elf[o + 8..o + 16].copy_from_slice(&0u64.to_le_bytes());
663        changed = true;
664    }
665    if changed {
666        fs::write(path, elf)?;
667    }
668    Ok(())
669}
670
671/// Rewrite the .addr_space.1 BTF DATASEC of a linked object from its ELF
672/// symbols and sort it by offset.
673fn fixup_arena_datasec(path: &Path) -> Result<()> {
674    let mut elf = fs::read(path)?;
675    let Some(parts) = ArenaElf::parse(&elf, path)? else {
676        return Ok(());
677    };
678    let Some((btf_off, _)) = parts.btf else {
679        return Ok(());
680    };
681    if le16(&elf, btf_off) != BTF_MAGIC {
682        bail!("{}: bad BTF magic", path.display());
683    }
684
685    // The DATASEC lists static arena globals too, so keep local symbols. Two
686    // units with a same-named static would collide and are rejected below.
687    let mut syms = HashMap::new();
688    for o in parts.syms() {
689        if le16(&elf, o + 6) != parts.arena_shndx || elf[o + 4] & 0xf != STT_OBJECT {
690            continue;
691        }
692        let name = cstr(&elf, parts.strtab + le32(&elf, o) as usize);
693        if syms.insert(name.clone(), le64(&elf, o + 8)).is_some() {
694            bail!("{}: duplicate arena symbol {name}", path.display());
695        }
696    }
697
698    let hdr_len = le32(&elf, btf_off + 4) as usize;
699    let types = btf_off + hdr_len + le32(&elf, btf_off + 8) as usize;
700    let types_end = types + le32(&elf, btf_off + 12) as usize;
701    let strs = btf_off + hdr_len + le32(&elf, btf_off + 16) as usize;
702    let mut var_names = vec![String::new()];
703    let mut datasec = None;
704    let mut o = types;
705
706    while o < types_end {
707        let name = cstr(&elf, strs + le32(&elf, o) as usize);
708        let info = le32(&elf, o + 4);
709        let (kind, vlen) = ((info >> 24) & 0x1f, (info & 0xffff) as usize);
710        let extra = match kind {
711            1 | 14 | 17 => 4,
712            3 => 12,
713            4 | 5 | 15 | 19 => vlen * 12,
714            6 | 13 => vlen * 8,
715            _ => 0,
716        };
717
718        var_names.push(if kind == BTF_KIND_VAR {
719            name.clone()
720        } else {
721            String::new()
722        });
723        if kind == BTF_KIND_DATASEC && name == ARENA_SEC {
724            datasec = Some((o + 12, vlen));
725        }
726        o += 12 + extra;
727    }
728    let Some((secinfos, vlen)) = datasec else {
729        return Ok(());
730    };
731
732    let mut entries: Vec<[u8; 12]> = (0..vlen)
733        .map(|i| {
734            elf[secinfos + i * 12..secinfos + (i + 1) * 12]
735                .try_into()
736                .unwrap()
737        })
738        .collect();
739    for e in entries.iter_mut() {
740        let name = &var_names[le32(e, 0) as usize];
741        let Some(&value) = syms.get(name) else {
742            bail!(
743                "{}: arena variable {name} has no ELF symbol",
744                path.display()
745            );
746        };
747        e[4..8].copy_from_slice(&(value as u32).to_le_bytes());
748    }
749    entries.sort_by_key(|e| le32(e, 4));
750    for (i, e) in entries.iter().enumerate() {
751        elf[secinfos + i * 12..secinfos + (i + 1) * 12].copy_from_slice(e);
752    }
753    fs::write(path, elf)?;
754    Ok(())
755}
756
757// Helper function to set up tracing and output compiler warnings
758fn with_clang_warnings<F, R>(f: F) -> Result<R>
759where
760    F: FnOnce() -> Result<R>,
761{
762    let fmt =
763        tracing_subscriber::fmt::layer().event_format(libbpf_cargo::util::CargoWarningFormatter);
764    let filter = filter::Targets::new()
765        .with_target("compiler-stdout", Level::INFO)
766        .with_target("compiler-stderr", Level::INFO);
767
768    let subscriber = tracing_subscriber::registry().with(fmt.with_filter(filter));
769
770    // Execute the closure with a tracing guard
771    {
772        let _guard = tracing::subscriber::set_default(subscriber);
773        f()
774    }
775}
776
777#[cfg(test)]
778mod tests {
779    use regex::Regex;
780    use sscanf::sscanf;
781
782    use crate::ClangInfo;
783
784    #[test]
785    fn test_bpf_builder_new() {
786        let td = tempfile::tempdir().unwrap();
787        unsafe { std::env::set_var("OUT_DIR", td.path()) };
788
789        let res = super::BpfBuilder::new();
790        assert!(res.is_ok(), "Failed to create BpfBuilder ({res:?})");
791    }
792
793    #[test]
794    fn test_vmlinux_h_ver_sha1() {
795        let clang_info = ClangInfo::new().unwrap();
796
797        let mut ar = tar::Archive::new(super::BpfBuilder::BPF_H_TAR);
798        let mut found = false;
799
800        let pattern = Regex::new(r"arch\/.*\/vmlinux-.*.h").unwrap();
801
802        for entry in ar.entries().unwrap() {
803            let entry = entry.unwrap();
804            let file_name = entry.header().path().unwrap();
805            let file_name_str = file_name.to_string_lossy().to_owned();
806            if file_name_str.contains(&clang_info.kernel_target().unwrap()) {
807                found = true;
808            }
809            if !pattern.find(&file_name_str).is_some() {
810                continue;
811            }
812
813            println!("checking {file_name_str}");
814
815            let (arch, ver, sha1) =
816                sscanf!(file_name_str, "arch/{String}/vmlinux-v{String}-g{String}.h").unwrap();
817            println!("vmlinux.h: arch={:?} ver={:?} sha1={:?}", arch, ver, sha1,);
818
819            assert!(
820                regex::Regex::new(r"^([1-9][0-9]*\.[0-9]+[a-z0-9-]*)$")
821                    .unwrap()
822                    .is_match(&ver)
823            );
824            assert!(
825                regex::Regex::new(r"^[0-9a-z]{12}$")
826                    .unwrap()
827                    .is_match(&sha1)
828            );
829        }
830
831        assert!(found);
832    }
833}