Skip to content

Commit

Permalink
rewrote the bpf example (#1877)
Browse files Browse the repository at this point in the history
Signed-off-by: yihuaf <yihuaf@unkies.org>
  • Loading branch information
yihuaf authored May 4, 2023
1 parent b51cf1a commit 5f31428
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 86 deletions.
2 changes: 1 addition & 1 deletion crates/libcgroups/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ readme = "README.md"
authors = ["youki team"]
edition = "2021"
rust-version = "1.58.1"
autoexamples = false
autoexamples = true
keywords = ["youki", "container", "cgroups"]

[features]
Expand Down
187 changes: 102 additions & 85 deletions crates/libcgroups/examples/bpf.rs
Original file line number Diff line number Diff line change
@@ -1,100 +1,117 @@
use anyhow::{bail, Result};
use clap::{Arg, SubCommand};
use std::os::unix::io::AsRawFd;
use std::path::Path;
use anyhow::Result;

use nix::fcntl::OFlag;
use nix::sys::stat::Mode;
#[cfg(feature = "cgroupsv2_devices")]
mod bpf {
use anyhow::{bail, Result};
use clap::Arg;
use clap::Command;
use oci_spec::runtime::LinuxDeviceCgroup;
use std::os::unix::io::AsRawFd;
use std::path::Path;

use cgroups::v2::devices::bpf;
use cgroups::v2::devices::emulator;
use cgroups::v2::devices::program;
use oci_spec::*;
use nix::fcntl::OFlag;
use nix::sys::stat::Mode;

const LICENSE: &'static str = &"Apache";
use libcgroups::v2::devices::bpf;
use libcgroups::v2::devices::emulator;
use libcgroups::v2::devices::program;

fn main() -> Result<()> {
env_logger::init();
const LICENSE: &str = "Apache";
fn cli() -> Command {
clap::Command::new("bpf")
.version("0.1")
.about("tools to test BPF program for cgroups v2 devices")
.arg(Arg::new("cgroup_dir").short('c').value_name("CGROUP_DIR"))
.subcommand(
Command::new("query").about("query list of BPF programs attached to cgroup dir"),
)
.subcommand(
Command::new("detach")
.about("detach BPF program by id")
.arg(
Arg::new("id")
.value_name("PROG_ID")
.required(true)
.help("ID of BPF program returned by query command"),
),
)
.subcommand(
Command::new("attach")
.about("compile rules to BPF and attach to cgroup dir")
.arg(
Arg::new("input_file")
.value_name("INPUT_FILE")
.required(true)
.help("File contains Vec<LinuxDeviceCgroup> in json format"),
),
)
}

let matches = clap::App::new("bpf")
.version("0.1")
.about("tools to test BPF program for cgroups v2 devices")
.arg(
Arg::with_name("cgroup_dir")
.short("c")
.value_name("CGROUP_DIR"),
)
.subcommand(
SubCommand::with_name("query")
.help("query list of BPF programs attached to cgroup dir"),
)
.subcommand(
SubCommand::with_name("detach")
.help("detach BPF program by id")
.arg(
Arg::with_name("id")
.value_name("PROG_ID")
.required(true)
.help("ID of BPF program returned by query command"),
),
)
.subcommand(
SubCommand::with_name("attach")
.help("compile rules to BPF and attach to cgroup dir")
.arg(
Arg::with_name("input_file")
.value_name("INPUT_FILE")
.required(true)
.help("File contains Vec<LinuxDeviceCgroup> in json format"),
),
)
.get_matches_safe()?;
fn parse_cgroupv1_device_rules<P: AsRef<Path>>(path: P) -> Result<Vec<LinuxDeviceCgroup>> {
let content = std::fs::read_to_string(path)?;
let devices = serde_json::from_str(&content)?;
Ok(devices)
}

let cgroup_dir = matches.value_of("cgroup_dir").unwrap();

let cgroup_fd = nix::dir::Dir::open(
cgroup_dir,
OFlag::O_RDONLY | OFlag::O_DIRECTORY,
Mode::from_bits(0o600).unwrap(),
)?;
pub fn run() -> Result<()> {
let matches = cli().get_matches();
let cgroup_dir = matches.get_one::<String>("cgroup_dir").unwrap();
let cgroup_fd = nix::dir::Dir::open(
cgroup_dir.as_str(),
OFlag::O_RDONLY | OFlag::O_DIRECTORY,
Mode::from_bits(0o600).unwrap(),
)?;
match matches.subcommand() {
Some(("query", _)) => {
let progs = bpf::prog::query(cgroup_fd.as_raw_fd())?;
for prog in &progs {
println!("prog: id={}, fd={}", prog.id, prog.fd);
}
}
Some(("detach", submatch)) => {
let prog_id = submatch.get_one::<String>("id").unwrap().parse::<u32>()?;
let progs = bpf::prog::query(cgroup_fd.as_raw_fd())?;
let prog = progs.iter().find(|v| v.id == prog_id);
if prog.is_none() {
bail!("can't get prog fd by prog id");
}

match matches.subcommand() {
("query", Some(_)) => {
let progs = bpf::prog_query(cgroup_fd.as_raw_fd())?;
for prog in &progs {
println!("prog: id={}, fd={}", prog.id, prog.fd);
bpf::prog::detach2(prog.unwrap().fd, cgroup_fd.as_raw_fd())?;
println!("detach ok");
}
}
("detach", Some(submatch)) => {
let prog_id = submatch.value_of("id").unwrap().parse::<u32>()?;
let progs = bpf::prog_query(cgroup_fd.as_raw_fd())?;
let prog = progs.iter().find(|v| v.id == prog_id);
if prog.is_none() {
bail!("can't get prog fd by prog id");
Some(("attach", submatch)) => {
let input_file = submatch.get_one::<String>("input_file").unwrap();
let rules = parse_cgroupv1_device_rules(input_file)?;
let mut emulator = emulator::Emulator::with_default_allow(false);
emulator.add_rules(&rules);
let prog = program::Program::from_rules(&emulator.rules, emulator.default_allow)?;
let prog_fd = bpf::prog::load(LICENSE, prog.bytecodes())?;
bpf::prog::attach(prog_fd, cgroup_fd.as_raw_fd())?;
println!("attach ok");
}

bpf::prog_detach2(prog.unwrap().fd, cgroup_fd.as_raw_fd())?;
println!("detach ok");
}
("attach", Some(submatch)) => {
let input_file = submatch.value_of("input_file").unwrap();
let rules = parse_cgroupv1_device_rules(&input_file)?;
let mut emulator = emulator::Emulator::with_default_allow(false);
emulator.add_rules(&rules)?;
let prog = program::Program::from_rules(&emulator.rules, emulator.default_allow)?;
let prog_fd = bpf::prog_load(LICENSE, prog.bytecodes())?;
bpf::prog_attach(prog_fd, cgroup_fd.as_raw_fd())?;
println!("attach ok");
}
_ => unreachable!(),
};
Ok(())
}
}

(_, _) => {}
};
#[cfg(not(feature = "cgroupsv2_devices"))]
mod bpf {
use anyhow::{bail, Result};

Ok(())
pub fn run() -> Result<()> {
if !cfg!(feature = "cgroupsv2_devices") {
bail!("cgroupsv2_devices feature is not enabled");
}

unreachable!()
}
}

fn parse_cgroupv1_device_rules<P: AsRef<Path>>(path: P) -> Result<Vec<LinuxDeviceCgroup>> {
let content = std::fs::read_to_string(path)?;
let devices = serde_json::from_str(&content)?;
Ok(devices)
fn main() -> Result<()> {
env_logger::init();
bpf::run()?;

Ok(())
}

0 comments on commit 5f31428

Please sign in to comment.