lockc 0.1.2

extern crate bindgen;

use std::{
    env, fs,
    io::{self, BufRead, Write},
    path, process,
    string::String,
};

use anyhow::Result;
use libbpf_cargo::SkeletonBuilder;
use tempfile::tempdir;

static CLANG_DEFAULT: &str = "/usr/bin/clang";
static HEADER_MAP_STRUCTS: &str = "src/bpf/map_structs.h";
static VMLINUX_URL: &str =
    "https://raw.githubusercontent.com/libbpf/libbpf-bootstrap/master/vmlinux/vmlinux_508.h";

/// Downloads vmlinux.h from github if it can't be generated.
fn download_btf(mut f: fs::File) -> Result<()> {
    let mut res = reqwest::blocking::get(VMLINUX_URL)?;
    io::copy(&mut res, &mut f)?;

    Ok(())
}

fn generate_btf<P: AsRef<path::Path>>(out_path: P) -> Result<()> {
    let vmlinux_path = out_path.as_ref().join("vmlinux.h");
    let mut f = fs::File::create(vmlinux_path)?;
    match process::Command::new("bpftool")
        .arg("btf")
        .arg("dump")
        .arg("file")
        .arg("/sys/kernel/btf/vmlinux")
        .arg("format")
        .arg("c")
        .output()
    {
        Ok(output) => {
            if output.status.success() {
                f.write_all(&output.stdout)?;
            } else {
                download_btf(f)?;
            }
        }
        Err(_) => download_btf(f)?,
    };

    Ok(())
}

fn generate_bpf_skel<P: AsRef<path::Path>>(out_path: P) -> Result<()> {
    let bpf_dir = path::Path::new("src").join("bpf");
    let src = bpf_dir.join("lockc.bpf.c");
    let tmp_dir = tempdir()?;
    let skel_unfiltered = tmp_dir.path().join("lockc.skel.rs");
    let clang = match env::var("CLANG") {
        Ok(val) => val,
        Err(_) => String::from(CLANG_DEFAULT),
    };

    SkeletonBuilder::new(&src)
        .clang(clang)
        .clang_args(format!("-I{}", out_path.as_ref().display()))
        .generate(&skel_unfiltered)?;

    // Skeletons generated by libbpf-cargo contain inner attributes. Including
    // source files with inner attributes is impossible if it's done with
    // include! macro. But we really want to use include! and get the skeleton
    // from OUT_DIR to not have to commit it to github...
    // So we need to get rid of those inner attributes ourselves and keep them
    // in the file doing !include.
    // TODO(vadorovsky): Solve that problem either by:
    // * making an option in libbpf-cargo to generate skeleton without inner
    //   attributes
    // * switching from libbpf-rs to aya
    // Second option preferred if possible. :)
    let skel_filtered = out_path.as_ref().join("lockc.skel.rs");
    let f_src = fs::File::open(skel_unfiltered)?;
    let f_src_buf = io::BufReader::new(f_src);
    let f_dest = fs::File::create(skel_filtered)?;
    let mut f_dest_buf = io::LineWriter::new(f_dest);
    for line_r in f_src_buf.lines() {
        let line = line_r.unwrap();
        if !line.contains("#![allow(") {
            f_dest_buf.write_all(line.as_bytes())?;
            f_dest_buf.write_all(b"\n")?;
        }
    }
    f_dest_buf.flush()?;

    println!("cargo:rerun-if-changed={}", src.to_str().unwrap());

    Ok(())
}

fn generate_bindings<P: AsRef<path::Path>>(out_path: P) -> Result<()> {
    println!("cargo:rerun-if-changed={}", HEADER_MAP_STRUCTS);

    let bindings = bindgen::Builder::default()
        .header(HEADER_MAP_STRUCTS)
        .parse_callbacks(Box::new(bindgen::CargoCallbacks))
        .generate()
        .unwrap();

    bindings.write_to_file(out_path.as_ref().join("bindings.rs"))?;

    Ok(())
}

fn main() -> Result<()> {
    let out_path = path::PathBuf::from(env::var("OUT_DIR")?);

    generate_btf(out_path.clone())?;
    generate_bpf_skel(out_path.clone())?;
    generate_bindings(out_path)?;

    Ok(())
}