use std::env;
use std::path::PathBuf;
use std::process::{Command, Stdio};
use libbpf_cargo::SkeletonBuilder;
include!("src/kernel_path.rs");
fn main() {
let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
println!("cargo:rerun-if-env-changed=KTSTR_KERNEL");
println!("cargo:rerun-if-changed=src/kernel_path.rs");
println!("cargo:rerun-if-changed=src/bpf/vmlinux_gen.c");
let ktstr_kernel = env::var("KTSTR_KERNEL").ok();
let vmlinux_h = out_dir.join("vmlinux.h");
let hash_path = out_dir.join("vmlinux.btf.hash");
let current_btf = resolve_btf(ktstr_kernel.as_deref());
let current_hash: Option<String> = current_btf.as_ref().and_then(|p| match std::fs::read(p) {
Ok(bytes) => Some(format!("{:016x}", siphash_13(&bytes))),
Err(e) => {
println!(
"cargo:warning=BTF source {} present but unreadable \
({e}); skipping hash check, reusing existing vmlinux.h",
p.display(),
);
None
}
});
let stored_hash: Option<String> = std::fs::read_to_string(&hash_path)
.ok()
.map(|s| s.trim().to_string());
let should_regen =
!vmlinux_h.exists() || (current_hash.is_some() && current_hash != stored_hash);
if should_regen {
let btf_source = current_btf.unwrap_or_else(|| {
panic!(
"no BTF source found. Set KTSTR_KERNEL to a kernel build \
directory, or ensure /sys/kernel/btf/vmlinux exists."
);
});
println!("generating vmlinux.h from {}", btf_source.display());
let libbpf_include =
PathBuf::from(env::var("DEP_BPF_INCLUDE").expect("DEP_BPF_INCLUDE not set"));
let vmlinux_gen_bin = out_dir.join("vmlinux_gen");
let driver_src = out_dir.join("vmlinux_gen_main.c");
std::fs::write(
&driver_src,
format!(
r#"
extern int generate_vmlinux_h(const char *, const char *);
int main(void) {{
return generate_vmlinux_h("{btf}", "{out}") == 0 ? 0 : 1;
}}
"#,
btf = btf_source.display(),
out = vmlinux_h.display(),
),
)
.expect("write driver source");
let libbpf_lib_dir = libbpf_include.parent().unwrap();
let compiler = cc::Build::new().get_compiler();
let status = Command::new(compiler.path())
.args([
"src/bpf/vmlinux_gen.c",
driver_src.to_str().unwrap(),
"-o",
vmlinux_gen_bin.to_str().unwrap(),
&format!("-I{}", libbpf_include.display()),
&format!("-L{}", libbpf_lib_dir.display()),
"-lbpf",
"-lelf",
"-lz",
])
.status()
.expect("compile vmlinux_gen");
assert!(status.success(), "failed to compile vmlinux_gen");
let status = Command::new(&vmlinux_gen_bin)
.status()
.expect("run vmlinux_gen");
assert!(
status.success(),
"vmlinux_gen failed — check BTF source: {}",
btf_source.display()
);
let hash_opt: Option<String> = match current_hash.as_deref() {
Some(h) => Some(h.to_string()),
None => match std::fs::read(&btf_source) {
Ok(bytes) => Some(format!("{:016x}", siphash_13(&bytes))),
Err(e) => {
println!(
"cargo:warning=post-regen BTF re-read failed ({e}); \
skipping hash sidecar — next build.rs run will \
regenerate conservatively"
);
None
}
},
};
if let Some(hash) = hash_opt {
std::fs::write(&hash_path, format!("{hash}\n"))
.unwrap_or_else(|e| panic!("write BTF hash sidecar {}: {e}", hash_path.display()));
}
}
if cfg!(target_arch = "aarch64") {
let content = std::fs::read_to_string(&vmlinux_h).expect("read vmlinux.h");
if !content.contains("struct user_pt_regs {") {
use std::io::Write;
let mut f = std::fs::OpenOptions::new()
.append(true)
.open(&vmlinux_h)
.expect("open vmlinux.h for append");
writeln!(
f,
"\n/* Added by build.rs: arm64 UAPI type needed by bpf_tracing.h */\n\
struct user_pt_regs {{\n\
\t__u64 regs[31];\n\
\t__u64 sp;\n\
\t__u64 pc;\n\
\t__u64 pstate;\n\
}};\n"
)
.expect("append user_pt_regs to vmlinux.h");
}
}
let clang_args = [
format!("-I{}", out_dir.display()),
format!("-I{}", "src/bpf"),
];
let skel_path = out_dir.join("probe_skel.rs");
SkeletonBuilder::new()
.source("src/bpf/probe.bpf.c")
.obj(out_dir.join("probe.o"))
.clang_args(clang_args.clone())
.reference_obj(true)
.build_and_generate(&skel_path)
.expect("build probe BPF skeleton");
let fentry_skel_path = out_dir.join("fentry_probe_skel.rs");
SkeletonBuilder::new()
.source("src/bpf/fentry_probe.bpf.c")
.obj(out_dir.join("fentry_probe.o"))
.clang_args(clang_args)
.reference_obj(true)
.build_and_generate(&fentry_skel_path)
.expect("build fentry probe BPF skeleton");
println!("cargo::rerun-if-changed=src/bpf/probe.bpf.c");
println!("cargo::rerun-if-changed=src/bpf/fentry_probe.bpf.c");
println!("cargo::rerun-if-changed=src/bpf/intf.h");
generate_shift_registry(&out_dir);
let busybox_bin = out_dir.join("busybox");
if !busybox_bin.exists() {
println!("cargo:warning=compiling busybox (first build only)...");
if Command::new("make").arg("--version").output().is_err() {
panic!(
"busybox build requires 'make' — install build-essential \
(Debian/Ubuntu) or base-devel (Fedora/Arch)"
);
}
if Command::new("gcc").arg("--version").output().is_err() {
panic!(
"busybox build requires 'gcc' — install build-essential \
(Debian/Ubuntu) or base-devel (Fedora/Arch)"
);
}
let busybox_src = out_dir.join("busybox-src");
if busybox_src.exists() && !busybox_src.join("Makefile").exists() {
std::fs::remove_dir_all(&busybox_src).expect("remove incomplete busybox-src");
}
if !busybox_src.join("Makefile").exists() {
let tarball_url = "https://github.com/mirror/busybox/archive/refs/tags/1_36_1.tar.gz";
let github_token = std::env::var("GITHUB_TOKEN").ok();
let attempt = |attempt_idx: u32| -> Result<(), String> {
let extract_dir = out_dir.join("busybox-extract");
if extract_dir.exists() {
let _ = std::fs::remove_dir_all(&extract_dir);
}
let client = reqwest::blocking::Client::builder()
.timeout(std::time::Duration::from_secs(120))
.connect_timeout(std::time::Duration::from_secs(30))
.user_agent(concat!("ktstr-build/", env!("CARGO_PKG_VERSION")))
.build()
.map_err(|e| format!("http client: {e}"))?;
let mut req = client.get(tarball_url);
if let Some(ref token) = github_token {
req = req.bearer_auth(token);
}
let resp = req
.send()
.and_then(|r| r.error_for_status())
.map_err(|e| format!("attempt {attempt_idx} request: {e}"))?;
let body = resp
.bytes()
.map_err(|e| format!("attempt {attempt_idx} body: {e}"))?;
let gz = flate2::read::GzDecoder::new(std::io::Cursor::new(body));
let mut archive = tar::Archive::new(gz);
archive
.unpack(&extract_dir)
.map_err(|e| format!("extract: {e}"))?;
let inner = extract_dir.join("busybox-1_36_1");
std::fs::rename(&inner, &busybox_src).map_err(|e| {
format!(
"expected extracted directory {} — tarball layout may have changed: {e}",
inner.display()
)
})?;
std::fs::remove_dir_all(&extract_dir).ok();
Ok(())
};
const MAX_TARBALL_ATTEMPTS: u32 = 4;
let mut tarball_err: Option<String> = None;
for i in 1..=MAX_TARBALL_ATTEMPTS {
println!(
"cargo:warning=downloading busybox source tarball (attempt {i}/{MAX_TARBALL_ATTEMPTS}) from {tarball_url}"
);
match attempt(i) {
Ok(()) => {
tarball_err = None;
break;
}
Err(e) => {
println!("cargo:warning=busybox tarball attempt {i} failed: {e}");
tarball_err = Some(e);
if i < MAX_TARBALL_ATTEMPTS {
let backoff = 1u64 << i;
std::thread::sleep(std::time::Duration::from_secs(backoff));
}
}
}
}
if !busybox_src.join("Makefile").exists() {
let tarball_err = tarball_err.unwrap_or_else(|| "unknown".to_string());
let git_url = "https://github.com/mirror/busybox.git";
println!(
"cargo:warning=busybox tarball failed ({tarball_err}), \
cloning {git_url} (requires network)"
);
if busybox_src.exists() {
std::fs::remove_dir_all(&busybox_src).expect("remove partial busybox-src");
}
let extract_dir = out_dir.join("busybox-extract");
if extract_dir.exists() {
std::fs::remove_dir_all(&extract_dir).ok();
}
let interrupt = std::sync::atomic::AtomicBool::new(false);
let clone_err = (|| -> Result<(), Box<dyn std::error::Error>> {
let mut prep = gix::prepare_clone(git_url, &busybox_src)?
.with_shallow(gix::remote::fetch::Shallow::DepthAtRemote(
1.try_into().expect("non-zero"),
))
.with_ref_name(Some("1_36_1"))?;
let (mut checkout, _) =
prep.fetch_then_checkout(gix::progress::Discard, &interrupt)?;
let (_repo, _) = checkout.main_worktree(gix::progress::Discard, &interrupt)?;
println!("cargo:warning=busybox source cloned via git");
Ok(())
})()
.err();
if !busybox_src.join("Makefile").exists() {
let clone_err = clone_err
.map(|e| e.to_string())
.unwrap_or_else(|| "checkout missing Makefile".to_string());
panic!(
"failed to obtain busybox source.\n\
tarball ({tarball_url}): {tarball_err}\n\
git clone ({git_url}): {clone_err}\n\
Check network connectivity. First build requires internet access."
);
}
}
}
let status = Command::new("make")
.arg("defconfig")
.current_dir(&busybox_src)
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
.status()
.expect("make defconfig");
assert!(status.success(), "busybox make defconfig failed");
let config_path = busybox_src.join(".config");
let config = std::fs::read_to_string(&config_path).expect("read busybox .config");
let config = config
.replace("# CONFIG_STATIC is not set", "CONFIG_STATIC=y")
.replace("CONFIG_TC=y", "# CONFIG_TC is not set");
std::fs::write(&config_path, config).expect("write patched busybox .config");
let status = Command::new("make")
.arg("oldconfig")
.current_dir(&busybox_src)
.stdin(Stdio::null())
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
.status()
.expect("make oldconfig");
assert!(status.success(), "busybox make oldconfig failed");
let nproc = std::thread::available_parallelism()
.map(|n| n.get().to_string())
.unwrap_or_else(|_| "1".to_string());
let status = Command::new("make")
.arg(format!("-j{nproc}"))
.current_dir(&busybox_src)
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
.status()
.expect("busybox make");
assert!(status.success(), "busybox build failed");
std::fs::copy(busybox_src.join("busybox"), &busybox_bin)
.expect("copy busybox binary to OUT_DIR");
}
}
fn generate_shift_registry(out_dir: &std::path::Path) {
use std::fmt::Write;
println!("cargo::rerun-if-changed=src/budget.rs");
let budget_rs = std::fs::read_to_string("src/budget.rs")
.expect("read src/budget.rs for shift-registry scan");
let mut shifts: Vec<(u32, String)> = Vec::new();
for line in budget_rs.lines() {
let line = line.trim();
let Some(rest) = line.strip_prefix("const ") else {
continue;
};
let Some((name_part, val_part)) = rest.split_once(": u32 = ") else {
continue;
};
let name = name_part.trim();
if !name.ends_with("_SHIFT") {
continue;
}
let val_str = val_part.trim_end_matches(';').trim();
let val: u32 = val_str.parse().unwrap_or_else(|e| {
panic!("shift-registry scan: parse `{val_str}` as u32 for {name}: {e}")
});
shifts.push((val, name.to_string()));
}
shifts.sort_by_key(|(v, _)| *v);
let mut out = String::from(
"// Generated by build.rs. Lists every `const *_SHIFT: u32 = N;`\n\
// declaration in src/budget.rs, sorted by shift value. The\n\
// budget tests assert their hand-classified one-bit and\n\
// multi-bit enumerations cover every entry so a new SHIFT\n\
// cannot land without being classified into the right test.\n\
pub(crate) const ALL_SHIFTS: &[(u32, &str)] = &[\n",
);
for (v, name) in &shifts {
writeln!(out, " ({v}, \"{name}\"),").expect("write shift entry");
}
out.push_str("];\n");
std::fs::write(out_dir.join("shift_registry.rs"), out).expect("write shift_registry.rs");
}
fn siphash_13(bytes: &[u8]) -> u64 {
use siphasher::sip::SipHasher13;
use std::hash::Hasher;
let mut h = SipHasher13::new_with_keys(0, 0);
h.write(bytes);
h.finish()
}