diff --git a/README.md b/README.md index eec82f8..9c49279 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ sudo ./target/release/tcbpftest sudo ./target/debug/tcbpftest ``` -You can also use `cargo xtask run [--release]` to build and run the program with one command. +You can also use `RUST_LOG=info cargo xtask run [--release]` to build and run the program with one command. The `tcbpftest` executable loads into the kernel the eBPF object file and attaches it to the `tc` hook, then runs the user-space program reading data from the maps and printing to stdout: diff --git a/tcbpftest-common/Cargo.toml b/tcbpftest-common/Cargo.toml index 815877b..4fb5d74 100644 --- a/tcbpftest-common/Cargo.toml +++ b/tcbpftest-common/Cargo.toml @@ -5,10 +5,10 @@ edition = "2021" [features] default = [] -userspace = [ "aya" ] +user = ["aya"] [dependencies] -aya = { git = "https://github.com/aya-rs/aya", branch="main", optional=true } +aya = { git = "https://github.com/aya-rs/aya", branch = "main", optional = true } [lib] path = "src/lib.rs" diff --git a/tcbpftest-common/src/lib.rs b/tcbpftest-common/src/lib.rs index 2ffffe8..dc8f69a 100644 --- a/tcbpftest-common/src/lib.rs +++ b/tcbpftest-common/src/lib.rs @@ -1,6 +1,7 @@ #![no_std] #[repr(C)] +#[derive(Clone, Copy)] pub struct PacketLog { pub len: u32, // packet length pub ctx_len: u32, // skb length diff --git a/tcbpftest-ebpf/Cargo.toml b/tcbpftest-ebpf/Cargo.toml index 39c627f..13bb57a 100644 --- a/tcbpftest-ebpf/Cargo.toml +++ b/tcbpftest-ebpf/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" [dependencies] aya-bpf = { git = "http://github.com/aya-rs/aya", branch = "main" } -memoffset = "0.6" +memoffset = "0.8" tcbpftest-common = { path = "../tcbpftest-common" } [[bin]] @@ -29,4 +29,4 @@ panic = "abort" codegen-units = 1 [workspace] -members = [] \ No newline at end of file +members = [] diff --git a/tcbpftest-ebpf/rust-toolchain.toml b/tcbpftest-ebpf/rust-toolchain.toml index c046a09..5d56faf 100644 --- a/tcbpftest-ebpf/rust-toolchain.toml +++ b/tcbpftest-ebpf/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel="nightly" +channel = "nightly" diff --git a/tcbpftest-ebpf/src/main.rs b/tcbpftest-ebpf/src/main.rs index 173892a..af374f0 100644 --- a/tcbpftest-ebpf/src/main.rs +++ b/tcbpftest-ebpf/src/main.rs @@ -1,14 +1,14 @@ #![no_std] #![no_main] +use aya_bpf::bindings::__sk_buff; +use aya_bpf::helpers::bpf_skb_pull_data; use aya_bpf::{ - BpfContext, macros::{classifier, map}, maps::PerfEventArray, programs::TcContext, + BpfContext, }; -use aya_bpf::bindings:: __sk_buff; -use aya_bpf::helpers::bpf_skb_pull_data; use core::mem; use memoffset::offset_of; @@ -27,9 +27,10 @@ fn panic(_info: &core::panic::PanicInfo) -> ! { } #[map(name = "EVENTS")] -static mut EVENTS: PerfEventArray = PerfEventArray::::with_max_entries(1024, 0); +static mut EVENTS: PerfEventArray = + PerfEventArray::::with_max_entries(1024, 0); -#[classifier(name="tcbpftest")] +#[classifier(name = "tcbpftest")] pub fn tcbpftest(ctx: TcContext) -> i32 { match unsafe { try_tcbpftest(ctx) } { Ok(ret) => ret, @@ -40,13 +41,13 @@ pub fn tcbpftest(ctx: TcContext) -> i32 { const ETH_P_IP: u16 = 0x0800; const ETH_HDR_LEN: usize = mem::size_of::(); const IP_HDR_LEN: usize = mem::size_of::(); -const IPPROTO_TCP : u8 = 6; -const IPPROTO_UDP : u8 = 17; +const IPPROTO_TCP: u8 = 6; +const IPPROTO_UDP: u8 = 17; unsafe fn try_tcbpftest(ctx: TcContext) -> Result { let ctx_len = ctx.len(); let skb = ctx.as_ptr() as *mut __sk_buff; - if bpf_skb_pull_data(skb, ctx_len) != 0 { + if bpf_skb_pull_data(skb, ctx_len) != 0 { return Err(199); } // get the ethernet header proto field as well as the IP protocol one @@ -63,37 +64,37 @@ unsafe fn try_tcbpftest(ctx: TcContext) -> Result { let rem_port2 = u16::from_be(ctx.load(ETH_HDR_LEN + IP_HDR_LEN + offset_of!(tcphdr, source))?); let loc_port2 = u16::from_be(ctx.load(ETH_HDR_LEN + IP_HDR_LEN + offset_of!(tcphdr, dest))?); - let rem_port_val : u16; - let loc_port_val : u16; + let rem_port_val: u16; + let loc_port_val: u16; unsafe { - rem_port_val = match ptr_at(&ctx, ETH_HDR_LEN + IP_HDR_LEN + offset_of!(tcphdr, source)) { + rem_port_val = match ptr_at(&ctx, ETH_HDR_LEN + IP_HDR_LEN + offset_of!(tcphdr, source)) { Err(_) => return Err(197), Ok(val) => *val, - }; - loc_port_val = match ptr_at(&ctx, ETH_HDR_LEN + IP_HDR_LEN + offset_of!(tcphdr, dest)) { - Err(_) => return Err(198), - Ok(val) => *val, - }; + }; + loc_port_val = match ptr_at(&ctx, ETH_HDR_LEN + IP_HDR_LEN + offset_of!(tcphdr, dest)) { + Err(_) => return Err(198), + Ok(val) => *val, + }; } let rem_port = u16::from_be(rem_port_val); let loc_port = u16::from_be(loc_port_val); - let mut udp_len_val : u16 = 0; + let mut udp_len_val: u16 = 0; if ip_proto == IPPROTO_UDP { - unsafe { - udp_len_val = match ptr_at(&ctx, ETH_HDR_LEN + IP_HDR_LEN + offset_of!(udphdr, len)) { - Err(_) => return Err(197), - Ok(val) => *val, - }; - } + unsafe { + udp_len_val = match ptr_at(&ctx, ETH_HDR_LEN + IP_HDR_LEN + offset_of!(udphdr, len)) { + Err(_) => return Err(197), + Ok(val) => *val, + }; + } } let log_entry = PacketLog { len: length as u32, ctx_len: ctx_len, src_addr: saddr, dest_addr: daddr, - eth_proto: eth_proto as u32, - eth_proto2: eth_proto2 as u32, + eth_proto: eth_proto as u32, + eth_proto2: eth_proto2 as u32, ip_proto: ip_proto as u32, remote_port: rem_port as u32, remote_port2: rem_port2 as u32, @@ -112,12 +113,12 @@ unsafe fn try_tcbpftest(ctx: TcContext) -> Result { // #[inline(always)] unsafe fn ptr_at(ctx: &TcContext, offset: usize) -> Result<*const T, ()> { - let raw_skb = ctx.as_ptr() as *const __sk_buff; + let raw_skb = ctx.as_ptr() as *const __sk_buff; let start = (*raw_skb).data as usize; let end = (*raw_skb).data_end as usize; let len = mem::size_of::(); - if start + offset + len > end { + if start + offset + len > end { return Err(()); } diff --git a/tcbpftest/Cargo.toml b/tcbpftest/Cargo.toml index d35c251..975b474 100644 --- a/tcbpftest/Cargo.toml +++ b/tcbpftest/Cargo.toml @@ -5,15 +5,21 @@ edition = "2021" publish = false [dependencies] -aya = { version = ">=0.11", features=["async_tokio"] } -tcbpftest-common = { path = "../tcbpftest-common", features=["userspace"] } -anyhow = "1.0.42" +aya = { version = ">=0.11", features = ["async_tokio"] } +tcbpftest-common = { path = "../tcbpftest-common", features = ["user"] } +anyhow = "1" bytes = "1" ctrlc = "3.2" log = "0.4" simplelog = "0.12" -tokio = { version = "1.18", features = ["macros", "rt", "rt-multi-thread", "net", "signal"] } -clap = { version = "4.0.18", features = ["derive"] } +tokio = { version = "1.25", features = [ + "macros", + "rt", + "rt-multi-thread", + "net", + "signal", +] } +clap = { version = "4.1", features = ["derive"] } [[bin]] name = "tcbpftest" diff --git a/tcbpftest/src/main.rs b/tcbpftest/src/main.rs index 484d1a6..efd8eb4 100644 --- a/tcbpftest/src/main.rs +++ b/tcbpftest/src/main.rs @@ -1,5 +1,10 @@ -use aya::programs::{tc, SchedClassifier, TcAttachType}; -use aya::{include_bytes_aligned, maps::perf::AsyncPerfEventArray, util::online_cpus, Bpf}; +use aya::{ + include_bytes_aligned, + maps::perf::AsyncPerfEventArray, + programs::{tc, SchedClassifier, TcAttachType}, + util::online_cpus, + Bpf, +}; use bytes::BytesMut; use clap::Parser; use log::info; @@ -50,12 +55,15 @@ async fn main() -> Result<(), anyhow::Error> { program.attach(&args.iface, TcAttachType::Ingress)?; let mut perf_array = AsyncPerfEventArray::try_from(bpf.map_mut("EVENTS")?)?; - for cpu_id in online_cpus()? { + + let cpus = online_cpus()?; + let num_cpus = cpus.len(); + for cpu_id in cpus { let mut buf = perf_array.open(cpu_id, None)?; task::spawn(async move { - let mut buffers = (0..10) - .map(|_| BytesMut::with_capacity(1024)) + let mut buffers = (0..num_cpus) + .map(|_| BytesMut::with_capacity(std::mem::size_of::())) .collect::>(); loop { diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml index e12df4a..bab7910 100644 --- a/xtask/Cargo.toml +++ b/xtask/Cargo.toml @@ -3,9 +3,7 @@ name = "xtask" version = "0.1.0" edition = "2021" -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - [dependencies] aya-tool = { git = "https://github.com/aya-rs/aya", branch = "main" } -clap = { version = "4", features = ["derive"] } anyhow = "1" +clap = { version = "4.1", features = ["derive"] } diff --git a/xtask/src/build_ebpf.rs b/xtask/src/build_ebpf.rs index 9fb4497..481612c 100644 --- a/xtask/src/build_ebpf.rs +++ b/xtask/src/build_ebpf.rs @@ -1,5 +1,4 @@ -use std::path::PathBuf; -use std::process::Command; +use std::{path::PathBuf, process::Command}; use clap::Parser; @@ -44,7 +43,6 @@ pub fn build_ebpf(opts: Options) -> Result<(), anyhow::Error> { let dir = PathBuf::from("tcbpftest-ebpf"); let target = format!("--target={}", opts.target); let mut args = vec![ - "+nightly", "build", "--verbose", target.as_str(), @@ -54,8 +52,14 @@ pub fn build_ebpf(opts: Options) -> Result<(), anyhow::Error> { if opts.release { args.push("--release") } + + // Command::new creates a child process which inherits all env variables. This means env + // vars set by the cargo xtask command are also inherited. RUSTUP_TOOLCHAIN is removed + // so the rust-toolchain.toml file in the -ebpf folder is honored. + let status = Command::new("cargo") - .current_dir(&dir) + .current_dir(dir) + .env_remove("RUSTUP_TOOLCHAIN") .args(&args) .status() .expect("failed to build bpf program"); diff --git a/xtask/src/codegen.rs b/xtask/src/codegen.rs index b39596b..b3e9a60 100644 --- a/xtask/src/codegen.rs +++ b/xtask/src/codegen.rs @@ -11,6 +11,6 @@ pub fn generate() -> Result<(), anyhow::Error> { )?; // Write the bindings to the $OUT_DIR/bindings.rs file. let mut out = File::create(dir.join("bindings.rs"))?; - write!(out, "{}", bindings)?; + write!(out, "{bindings}")?; Ok(()) } diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 6ad1f49..c3dd660 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -30,7 +30,7 @@ fn main() { }; if let Err(e) = ret { - eprintln!("{:#}", e); + eprintln!("{e:#}"); exit(1); } } diff --git a/xtask/src/run.rs b/xtask/src/run.rs index 8d1de3c..01470fe 100644 --- a/xtask/src/run.rs +++ b/xtask/src/run.rs @@ -1,4 +1,4 @@ -use std::{os::unix::process::CommandExt, process::Command}; +use std::process::Command; use anyhow::Context as _; use clap::Parser; @@ -47,7 +47,7 @@ pub fn run(opts: Options) -> Result<(), anyhow::Error> { // profile we are building (release or debug) let profile = if opts.release { "release" } else { "debug" }; - let bin_path = format!("target/{}/tcbpftest", profile); + let bin_path = format!("target/{profile}/tcbpftest"); // arguments to pass to the application let mut run_args: Vec<_> = opts.run_args.iter().map(String::as_str).collect(); @@ -57,11 +57,14 @@ pub fn run(opts: Options) -> Result<(), anyhow::Error> { args.push(bin_path.as_str()); args.append(&mut run_args); - // spawn the command - let err = Command::new(args.get(0).expect("No first argument")) + // run the command + let status = Command::new(args.first().expect("No first argument")) .args(args.iter().skip(1)) - .exec(); + .status() + .expect("failed to run the command"); - // we shouldn't get here unless the command failed to spawn - Err(anyhow::Error::from(err).context(format!("Failed to run `{}`", args.join(" ")))) + if !status.success() { + anyhow::bail!("Failed to run `{}`", args.join(" ")); + } + Ok(()) }