diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..f0ccbc9 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,2 @@ +[alias] +xtask = "run --package xtask --" \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d6e6ce2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,16 @@ +### https://raw.github.com/github/gitignore/master/Rust.gitignore + +# Generated by Cargo +# will have compiled files and executables +debug/ +target/ + +# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries +# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html +Cargo.lock + +# These are backup files generated by rustfmt +**/*.rs.bk + +.idea +.vscode diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..e3bae2a --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,2 @@ +[workspace] +members = ["tcbpftest", "tcbpftest-common", "xtask"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..f5e6a17 --- /dev/null +++ b/README.md @@ -0,0 +1,28 @@ +# tcbpftest + +## Prerequisites + +1. Install a rust stable toolchain: `rustup install stable` +1. Install a rust nightly toolchain: `rustup install nightly` +1. Install bpf-linker: `cargo install bpf-linker` + +## Build eBPF + +```bash +cargo xtask build-ebpf +``` + +To perform a release build you can use the `--release` flag. +You may also change the target architecture with the `--target` flag + +## Build Userspace + +```bash +cargo build +``` + +## Run + +```bash +sudo target/debug/tcbpftest +``` diff --git a/tcbpftest-common/Cargo.toml b/tcbpftest-common/Cargo.toml new file mode 100644 index 0000000..e4a1647 --- /dev/null +++ b/tcbpftest-common/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "tcbpftest-common" +version = "0.1.0" +edition = "2018" + +[features] +default = [] +userspace = [ "aya" ] + +[dependencies] +aya = { git = "https://github.com/aya-rs/aya", branch="main", optional=true } + +[lib] +path = "src/lib.rs" \ No newline at end of file diff --git a/tcbpftest-common/src/lib.rs b/tcbpftest-common/src/lib.rs new file mode 100644 index 0000000..2cbce4c --- /dev/null +++ b/tcbpftest-common/src/lib.rs @@ -0,0 +1,10 @@ +#![no_std] + +#[repr(C)] +pub struct PacketLog { + pub len: u32, // packet length + pub src_addr: u32, // ipv4 source IP address +} + +#[cfg(feature = "user")] +unsafe impl aya::Pod for PacketLog {} diff --git a/tcbpftest-ebpf/.cargo/config.toml b/tcbpftest-ebpf/.cargo/config.toml new file mode 100644 index 0000000..5d7e591 --- /dev/null +++ b/tcbpftest-ebpf/.cargo/config.toml @@ -0,0 +1,6 @@ +[build] +target-dir = "../target" +target = "bpfel-unknown-none" + +[unstable] +build-std = ["core"] \ No newline at end of file diff --git a/tcbpftest-ebpf/Cargo.toml b/tcbpftest-ebpf/Cargo.toml new file mode 100644 index 0000000..d3b16b6 --- /dev/null +++ b/tcbpftest-ebpf/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "tcbpftest-ebpf" +version = "0.1.0" +edition = "2018" + +[dependencies] +aya-bpf = { git = "http://github.com/aya-rs/aya", branch = "main" } +memoffset = "0.6.1" +tcbpftest-common = { path = "../tcbpftest-common" } + +[[bin]] +name = "tcbpftest" +path = "src/main.rs" + +[profile.dev] +panic = "abort" +debug = 1 +opt-level = 2 +overflow-checks = false + +[profile.release] +panic = "abort" + +[workspace] +members = [] diff --git a/tcbpftest-ebpf/src/bindings.rs b/tcbpftest-ebpf/src/bindings.rs new file mode 100644 index 0000000..4ddab4c --- /dev/null +++ b/tcbpftest-ebpf/src/bindings.rs @@ -0,0 +1,189 @@ +/* automatically generated by rust-bindgen 0.59.2 */ + +#[repr(C)] +#[derive(Copy, Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub struct __BindgenBitfieldUnit { + storage: Storage, +} +impl __BindgenBitfieldUnit { + #[inline] + pub const fn new(storage: Storage) -> Self { + Self { storage } + } +} +impl __BindgenBitfieldUnit +where + Storage: AsRef<[u8]> + AsMut<[u8]>, +{ + #[inline] + pub fn get_bit(&self, index: usize) -> bool { + debug_assert!(index / 8 < self.storage.as_ref().len()); + let byte_index = index / 8; + let byte = self.storage.as_ref()[byte_index]; + let bit_index = if cfg!(target_endian = "big") { + 7 - (index % 8) + } else { + index % 8 + }; + let mask = 1 << bit_index; + byte & mask == mask + } + #[inline] + pub fn set_bit(&mut self, index: usize, val: bool) { + debug_assert!(index / 8 < self.storage.as_ref().len()); + let byte_index = index / 8; + let byte = &mut self.storage.as_mut()[byte_index]; + let bit_index = if cfg!(target_endian = "big") { + 7 - (index % 8) + } else { + index % 8 + }; + let mask = 1 << bit_index; + if val { + *byte |= mask; + } else { + *byte &= !mask; + } + } + #[inline] + pub fn get(&self, bit_offset: usize, bit_width: u8) -> u64 { + debug_assert!(bit_width <= 64); + debug_assert!(bit_offset / 8 < self.storage.as_ref().len()); + debug_assert!((bit_offset + (bit_width as usize)) / 8 <= self.storage.as_ref().len()); + let mut val = 0; + for i in 0..(bit_width as usize) { + if self.get_bit(i + bit_offset) { + let index = if cfg!(target_endian = "big") { + bit_width as usize - 1 - i + } else { + i + }; + val |= 1 << index; + } + } + val + } + #[inline] + pub fn set(&mut self, bit_offset: usize, bit_width: u8, val: u64) { + debug_assert!(bit_width <= 64); + debug_assert!(bit_offset / 8 < self.storage.as_ref().len()); + debug_assert!((bit_offset + (bit_width as usize)) / 8 <= self.storage.as_ref().len()); + for i in 0..(bit_width as usize) { + let mask = 1 << i; + let val_bit_is_set = val & mask == mask; + let index = if cfg!(target_endian = "big") { + bit_width as usize - 1 - i + } else { + i + }; + self.set_bit(index + bit_offset, val_bit_is_set); + } + } +} +pub type __u8 = ::aya_bpf::cty::c_uchar; +pub type __u16 = ::aya_bpf::cty::c_ushort; +pub type __u32 = ::aya_bpf::cty::c_uint; +pub type __be16 = __u16; +pub type __be32 = __u32; +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct ethhdr { + pub h_dest: [::aya_bpf::cty::c_uchar; 6usize], + pub h_source: [::aya_bpf::cty::c_uchar; 6usize], + pub h_proto: __be16, +} +pub type __sum16 = __u16; +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct iphdr { + pub _bitfield_align_1: [u8; 0], + pub _bitfield_1: __BindgenBitfieldUnit<[u8; 1usize]>, + pub tos: __u8, + pub tot_len: __be16, + pub id: __be16, + pub frag_off: __be16, + pub ttl: __u8, + pub protocol: __u8, + pub check: __sum16, + pub saddr: __be32, + pub daddr: __be32, +} +impl iphdr { + #[inline] + pub fn ihl(&self) -> __u8 { + unsafe { ::core::mem::transmute(self._bitfield_1.get(0usize, 4u8) as u8) } + } + #[inline] + pub fn set_ihl(&mut self, val: __u8) { + unsafe { + let val: u8 = ::core::mem::transmute(val); + self._bitfield_1.set(0usize, 4u8, val as u64) + } + } + #[inline] + pub fn version(&self) -> __u8 { + unsafe { ::core::mem::transmute(self._bitfield_1.get(4usize, 4u8) as u8) } + } + #[inline] + pub fn set_version(&mut self, val: __u8) { + unsafe { + let val: u8 = ::core::mem::transmute(val); + self._bitfield_1.set(4usize, 4u8, val as u64) + } + } + #[inline] + pub fn new_bitfield_1(ihl: __u8, version: __u8) -> __BindgenBitfieldUnit<[u8; 1usize]> { + let mut __bindgen_bitfield_unit: __BindgenBitfieldUnit<[u8; 1usize]> = Default::default(); + __bindgen_bitfield_unit.set(0usize, 4u8, { + let ihl: u8 = unsafe { ::core::mem::transmute(ihl) }; + ihl as u64 + }); + __bindgen_bitfield_unit.set(4usize, 4u8, { + let version: u8 = unsafe { ::core::mem::transmute(version) }; + version as u64 + }); + __bindgen_bitfield_unit + } +} + +impl __BindgenBitfieldUnit {} +impl ethhdr { + pub fn h_dest(&self) -> Option<[::aya_bpf::cty::c_uchar; 6usize]> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.h_dest) }.ok() + } + pub fn h_source(&self) -> Option<[::aya_bpf::cty::c_uchar; 6usize]> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.h_source) }.ok() + } + pub fn h_proto(&self) -> Option<__be16> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.h_proto) }.ok() + } +} +impl iphdr { + pub fn tos(&self) -> Option<__u8> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.tos) }.ok() + } + pub fn tot_len(&self) -> Option<__be16> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.tot_len) }.ok() + } + pub fn id(&self) -> Option<__be16> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.id) }.ok() + } + pub fn frag_off(&self) -> Option<__be16> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.frag_off) }.ok() + } + pub fn ttl(&self) -> Option<__u8> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.ttl) }.ok() + } + pub fn protocol(&self) -> Option<__u8> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.protocol) }.ok() + } + pub fn check(&self) -> Option<__sum16> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.check) }.ok() + } + pub fn saddr(&self) -> Option<__be32> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.saddr) }.ok() + } + pub fn daddr(&self) -> Option<__be32> { + unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.daddr) }.ok() + } +} diff --git a/tcbpftest-ebpf/src/main.rs b/tcbpftest-ebpf/src/main.rs new file mode 100644 index 0000000..238d666 --- /dev/null +++ b/tcbpftest-ebpf/src/main.rs @@ -0,0 +1,68 @@ +#![no_std] +#![no_main] + +use aya_bpf::{ + BpfContext, + macros::{map, classifier}, + maps::PerfEventArray, + programs::SkBuffContext, +}; +use aya_bpf::bindings::__sk_buff; + +use core::convert::TryInto; +use core::mem; +use memoffset::offset_of; + +use tcbpftest_common::PacketLog; + +mod bindings; +use bindings::{ethhdr, iphdr}; + +#[panic_handler] +fn panic(_info: &core::panic::PanicInfo) -> ! { + unreachable!() +} + +#[map(name = "EVENTS")] +static mut EVENTS: PerfEventArray = PerfEventArray::::with_max_entries(1024, 0); + +#[classifier(name="tcbpftest")] +pub fn tcbpftest(ctx: SkBuffContext) -> i32 { + match unsafe { try_tcbpftest(ctx) } { + Ok(ret) => ret, + Err(_) => 123, + } +} + +#[inline(always)] +unsafe fn ptr_at(ctx: &SkBuffContext, offset: usize) -> Result<*const T, ()> { + 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 { + return Err(()); + } + + Ok((start + offset) as *const T) +} + +unsafe fn try_tcbpftest(ctx: SkBuffContext) -> Result { + let skb = ctx.as_ptr() as *const __sk_buff; + let offset : usize = 8; + let val = match ptr_at::(&ctx, offset) { + Err(_) => return Err(123), + Ok(v) => v, + }; + let proto_bytes = u16::from_be(*val); + let log_entry = PacketLog { + len: u32::from_be((*skb).len), + proto: proto_bytes as u32, + }; + + unsafe { + EVENTS.output(&ctx, &log_entry, 0); + } + Ok(0) +} diff --git a/tcbpftest/Cargo.toml b/tcbpftest/Cargo.toml new file mode 100644 index 0000000..a8060c8 --- /dev/null +++ b/tcbpftest/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "tcbpftest" +version = "0.1.0" +edition = "2018" +publish = false + +[dependencies] +aya = { git = "https://github.com/aya-rs/aya", branch="main", features=["async_tokio"] } +tcbpftest-common = { path = "../tcbpftest-common", features=["userspace"] } +anyhow = "1.0.42" +bytes = "1.1" +ctrlc = "3.2" + +structopt = { version = "0.3"} +tokio = { version = "1.5.0", features = ["macros", "rt", "rt-multi-thread", "net", "signal"] } +simplelog = "0.11" +log = "0.4" + +[[bin]] +name = "tcbpftest" +path = "src/main.rs" diff --git a/tcbpftest/src/main.rs b/tcbpftest/src/main.rs new file mode 100644 index 0000000..1d2a312 --- /dev/null +++ b/tcbpftest/src/main.rs @@ -0,0 +1,94 @@ +use anyhow::Context; +use aya::programs::{tc, SchedClassifier, TcAttachType}; +use aya::{ + include_bytes_aligned, + maps::perf::{AsyncPerfEventArray, PerfBufferError}, + util::online_cpus, + Bpf, +}; +use bytes::BytesMut; +use log::info; +use simplelog::{ColorChoice, ConfigBuilder, LevelFilter, TermLogger, TerminalMode}; +use std::net::Ipv4Addr; +use std::{ + convert::{TryFrom, TryInto}, + sync::atomic::{AtomicBool, Ordering}, + sync::Arc, +}; +use structopt::StructOpt; +use tokio::{signal, task}; + +use tcbpftest_common::PacketLog; + +#[derive(Debug, StructOpt)] +struct Opt { + #[structopt(short, long, default_value = "eth0")] + iface: String, +} + +#[tokio::main] +async fn main() -> Result<(), anyhow::Error> { + let opt = Opt::from_args(); + + TermLogger::init( + LevelFilter::Debug, + ConfigBuilder::new() + .set_target_level(LevelFilter::Error) + .set_location_level(LevelFilter::Error) + .build(), + TerminalMode::Mixed, + ColorChoice::Auto, + )?; + + #[cfg(debug_assertions)] + let mut bpf = Bpf::load(include_bytes_aligned!( + "../../target/bpfel-unknown-none/debug/tcbpftest" + ))?; + #[cfg(not(debug_assertions))] + let mut bpf = Bpf::load(include_bytes_aligned!( + "../../target/bpfel-unknown-none/release/tcbpftest}" + ))?; + + // This will include your eBPF object file as raw bytes at compile-time and load it at + // runtime. This approach is recommended for most real-world use cases. If you would + // like to specify the eBPF program at runtime rather than at compile-time, you can + // reach for `Bpf::load_file` instead. + #[cfg(debug_assertions)] + let mut bpf = Bpf::load(include_bytes_aligned!( + "../../target/bpfel-unknown-none/debug/tcbpftest" + ))?; + #[cfg(not(debug_assertions))] + let mut bpf = Bpf::load(include_bytes_aligned!( + "../../target/bpfel-unknown-none/release/tcbpftest" + ))?; + // error adding clsact to the interface if it is already added is harmless + // the full cleanup can be done with 'sudo tc qdisc del dev eth0 clsact'. + let _ = tc::qdisc_add_clsact(&opt.iface); + let program: &mut SchedClassifier = bpf.program_mut("tcbpftest").unwrap().try_into()?; + program.load()?; + program.attach(&opt.iface, TcAttachType::Ingress)?; + + let mut perf_array = AsyncPerfEventArray::try_from(bpf.map_mut("EVENTS")?)?; + for cpu_id in online_cpus()? { + let mut buf = perf_array.open(cpu_id, None)?; + + task::spawn(async move { + let mut buffers = (0..10) + .map(|_| BytesMut::with_capacity(1024)) + .collect::>(); + + loop { + let events = buf.read_events(&mut buffers).await.unwrap(); + for i in 0..events.read { + let buf = &mut buffers[i]; + let ptr = buf.as_ptr() as *const PacketLog; + let data = unsafe { ptr.read_unaligned() }; + let src_addr = Ipv4Addr::from(data.ipv4_address); + println!("LOG: LEN {}, SRC_IP {}", data.len, src_addr); + } + } + }); + } + signal::ctrl_c().await.expect("failed to listen for event"); + Ok::<_, anyhow::Error>(()) +} diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml new file mode 100644 index 0000000..0185ca1 --- /dev/null +++ b/xtask/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "xtask" +version = "0.1.0" +edition = "2018" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +aya-gen = { git = "https://github.com/aya-rs/aya", branch = "main" } +structopt = {version = "0.3", default-features = false } +anyhow = "1" diff --git a/xtask/src/build_ebpf.rs b/xtask/src/build_ebpf.rs new file mode 100644 index 0000000..768fc0e --- /dev/null +++ b/xtask/src/build_ebpf.rs @@ -0,0 +1,62 @@ +use std::path::PathBuf; +use std::process::Command; + +use structopt::StructOpt; + +#[derive(Debug, Copy, Clone)] +pub enum Architecture { + BpfEl, + BpfEb, +} + +impl std::str::FromStr for Architecture { + type Err = String; + + fn from_str(s: &str) -> Result { + Ok(match s { + "bpfel-unknown-none" => Architecture::BpfEl, + "bpfeb-unknown-none" => Architecture::BpfEb, + _ => return Err("invalid target".to_owned()), + }) + } +} + +impl std::fmt::Display for Architecture { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Architecture::BpfEl => "bpfel-unknown-none", + Architecture::BpfEb => "bpfeb-unknown-none", + }) + } +} + +#[derive(StructOpt)] +pub struct Options { + #[structopt(default_value = "bpfel-unknown-none", long)] + target: Architecture, + #[structopt(long)] + release: bool, +} + +pub fn build(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(), + "-Z", + "build-std=core", + ]; + if opts.release { + args.push("--release") + } + let status = Command::new("cargo") + .current_dir(&dir) + .args(&args) + .status() + .expect("failed to build bpf examples"); + assert!(status.success()); + Ok(()) +} diff --git a/xtask/src/codegen.rs b/xtask/src/codegen.rs new file mode 100644 index 0000000..6ffd581 --- /dev/null +++ b/xtask/src/codegen.rs @@ -0,0 +1,16 @@ +use aya_gen::btf_types; +use std::{ + fs::File, + io::Write, + path::{Path, PathBuf}, +}; + +pub fn generate() -> Result<(), anyhow::Error> { + let dir = PathBuf::from("tcbpftest-ebpf/src"); + let names: Vec<&str> = vec!["ethhdr", "iphdr"]; + let bindings = btf_types::generate(Path::new("/sys/kernel/btf/vmlinux"), &names, true)?; + // Write the bindings to the $OUT_DIR/bindings.rs file. + let mut out = File::create(dir.join("bindings.rs"))?; + write!(out, "{}", bindings)?; + Ok(()) +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs new file mode 100644 index 0000000..628dafd --- /dev/null +++ b/xtask/src/main.rs @@ -0,0 +1,32 @@ +mod build_ebpf; +mod codegen; + +use std::process::exit; + +use structopt::StructOpt; +#[derive(StructOpt)] +pub struct Options { + #[structopt(subcommand)] + command: Command, +} + +#[derive(StructOpt)] +enum Command { + BuildEbpf(build_ebpf::Options), + Codegen, +} + +fn main() { + let opts = Options::from_args(); + + use Command::*; + let ret = match opts.command { + BuildEbpf(opts) => build_ebpf::build(opts), + Codegen => codegen::generate(), + }; + + if let Err(e) = ret { + eprintln!("{:#}", e); + exit(1); + } +}