From 974a8cd5baba791fefce8c55ceb3ab5729026fae Mon Sep 17 00:00:00 2001 From: ajwerner Date: Mon, 19 Jun 2023 16:46:36 -0400 Subject: [PATCH] integration-test: add RingBuf tests This simple test exercises sending data from a uprobe. It also adds a test to demonstrate how to combine the RingBuf with async notifications. --- test/integration-ebpf/Cargo.toml | 4 ++ test/integration-ebpf/src/ring_buf.rs | 38 ++++++++++ test/integration-test/Cargo.toml | 1 + test/integration-test/src/tests/mod.rs | 3 +- test/integration-test/src/tests/ring_buf.rs | 57 +++++++++++++++ .../src/tests/ring_buf_async.rs | 72 +++++++++++++++++++ 6 files changed, 174 insertions(+), 1 deletion(-) create mode 100644 test/integration-ebpf/src/ring_buf.rs create mode 100644 test/integration-test/src/tests/ring_buf.rs create mode 100644 test/integration-test/src/tests/ring_buf_async.rs diff --git a/test/integration-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml index 0683143a9..72bae2add 100644 --- a/test/integration-ebpf/Cargo.toml +++ b/test/integration-ebpf/Cargo.toml @@ -31,3 +31,7 @@ path = "src/test.rs" [[bin]] name = "relocations" path = "src/relocations.rs" + +[[bin]] +name = "ring_buf" +path = "src/ring_buf.rs" diff --git a/test/integration-ebpf/src/ring_buf.rs b/test/integration-ebpf/src/ring_buf.rs new file mode 100644 index 000000000..b1c939daa --- /dev/null +++ b/test/integration-ebpf/src/ring_buf.rs @@ -0,0 +1,38 @@ +#![no_std] +#![no_main] + +use aya_bpf::{ + macros::{map, uprobe}, + maps::RingBuf, + programs::ProbeContext, +}; +use core::mem::size_of; + +// Make a buffer large enough to hold MAX_ENTRIES entries at the same time. +// This requires taking into consideration the header size. +type Entry = u64; +const MAX_ENTRIES: usize = 1024; +const HDR_SIZE: usize = aya_bpf::bindings::BPF_RINGBUF_HDR_SZ as usize; + +// Add 1 because the capacity is actually one less than you might think +// because the consumer_pos and producer_pos being equal would mean that +// the buffer is empty. +const RING_BUF_SIZE: usize = ((size_of::() + HDR_SIZE) * MAX_ENTRIES) + 1; + +#[map] +static RING_BUF: RingBuf = RingBuf::with_byte_size(RING_BUF_SIZE as u32, 0); + +#[uprobe] +pub fn ring_buf_test(ctx: ProbeContext) { + // Write the first argument to the function back out to RING_BUF. + let Some(arg): Option = ctx.arg(0) else { return }; + if let Some(mut entry) = RING_BUF.reserve::(0) { + entry.write(arg); + entry.submit(0); + } +} + +#[panic_handler] +fn panic(_info: &core::panic::PanicInfo) -> ! { + loop {} +} diff --git a/test/integration-test/Cargo.toml b/test/integration-test/Cargo.toml index 24975368b..bdbf32fe8 100644 --- a/test/integration-test/Cargo.toml +++ b/test/integration-test/Cargo.toml @@ -23,3 +23,4 @@ regex = "1" tempfile = "3.3.0" libtest-mimic = "0.6.0" tokio = { version = "1.24", features = ["rt", "rt-multi-thread", "sync", "time"] } +rand = { version = "0.8.5" } diff --git a/test/integration-test/src/tests/mod.rs b/test/integration-test/src/tests/mod.rs index c26ca5a44..a1c95ec4e 100644 --- a/test/integration-test/src/tests/mod.rs +++ b/test/integration-test/src/tests/mod.rs @@ -10,8 +10,9 @@ pub mod load; pub mod log; pub mod rbpf; pub mod relocations; +pub mod ring_buf; +pub mod ring_buf_async; pub mod smoke; - pub use integration_test_macros::{integration_test, tokio_integration_test}; #[derive(Debug)] diff --git a/test/integration-test/src/tests/ring_buf.rs b/test/integration-test/src/tests/ring_buf.rs new file mode 100644 index 000000000..74a432f03 --- /dev/null +++ b/test/integration-test/src/tests/ring_buf.rs @@ -0,0 +1,57 @@ +use super::integration_test; +use aya::{include_bytes_aligned, maps::ring_buf::RingBuf, programs::UProbe, Bpf}; + +#[integration_test] +fn ring_buf() { + let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/release/ring_buf"); + let mut bpf = Bpf::load(bytes).unwrap(); + let ring_buf = bpf.take_map("RING_BUF").unwrap(); + let mut ring_buf = RingBuf::try_from(ring_buf).unwrap(); + + let prog: &mut UProbe = bpf + .program_mut("ring_buf_test") + .unwrap() + .try_into() + .unwrap(); + prog.load().unwrap(); + prog.attach( + Some("ring_buf_trigger_ebpf_program"), + 0, + "/proc/self/exe", + None, + ) + .unwrap(); + + // Generate some random data. + let data = gen_data(); + + // Call the function that the uprobe is attached to with randomly generated data. + for val in &data { + ring_buf_trigger_ebpf_program(*val); + } + // Read the data back out of the ring buffer. + let mut seen = Vec::::new(); + while seen.len() < data.len() { + if let Some(item) = ring_buf.next() { + let item: [u8; 8] = (*item).try_into().unwrap(); + let arg = u64::from_ne_bytes(item); + seen.push(arg); + } + } + // Ensure that the data that was read matches what was passed. + assert_eq!(seen, data); +} + +#[no_mangle] +#[inline(never)] +pub extern "C" fn ring_buf_trigger_ebpf_program(_arg: u64) {} + +/// Generate a variable length vector of u64s. The number of values is always small enough to fit +/// into the RING_BUF defined in the probe. +pub(crate) fn gen_data() -> Vec { + const DATA_LEN_RANGE: core::ops::RangeInclusive = 1..=1024; + use rand::Rng as _; + let mut rng = rand::thread_rng(); + let n = rng.gen_range(DATA_LEN_RANGE); + std::iter::repeat_with(|| rng.gen()).take(n).collect() +} diff --git a/test/integration-test/src/tests/ring_buf_async.rs b/test/integration-test/src/tests/ring_buf_async.rs new file mode 100644 index 000000000..d348550ab --- /dev/null +++ b/test/integration-test/src/tests/ring_buf_async.rs @@ -0,0 +1,72 @@ +use std::os::fd::AsRawFd as _; + +use aya::maps::RingBuf; + +use aya::{include_bytes_aligned, programs::UProbe, Bpf}; +use tokio::{ + io::unix::AsyncFd, + time::{sleep, Duration}, +}; + +use super::tokio_integration_test; + +#[tokio_integration_test] +async fn ring_buf_async() { + let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/release/ring_buf"); + let mut bpf = Bpf::load(bytes).unwrap(); + let ring_buf = bpf.take_map("RING_BUF").unwrap(); + let mut ring_buf = RingBuf::try_from(ring_buf).unwrap(); + + let prog: &mut UProbe = bpf + .program_mut("ring_buf_test") + .unwrap() + .try_into() + .unwrap(); + prog.load().unwrap(); + prog.attach( + Some("ring_buf_trigger_ebpf_program"), + 0, + "/proc/self/exe", + None, + ) + .unwrap(); + + // Generate some random data. + let data = super::ring_buf::gen_data(); + let write_handle = + tokio::task::spawn(call_ring_buf_trigger_ebpf_program_over_time(data.clone())); + + // Construct an AsyncFd from the RingBuf in order to receive readiness notifications. + let async_fd = AsyncFd::new(ring_buf.as_raw_fd()).unwrap(); + let seen = { + let mut seen = Vec::with_capacity(data.len()); + while seen.len() < data.len() { + // Wait for readiness, then clear the bit before reading so that no notifications + // are missed. + async_fd.readable().await.unwrap().clear_ready(); + while let Some(data) = ring_buf.next() { + let data: [u8; 8] = (*data).try_into().unwrap(); + let arg = u64::from_ne_bytes(data); + seen.push(arg); + } + } + seen + }; + + // Ensure that the data that was read matches what was passed. + assert_eq!(seen, data); + write_handle.await.unwrap(); +} + +async fn call_ring_buf_trigger_ebpf_program_over_time(data: Vec) { + let random_duration = || { + use rand::Rng as _; + let mut rng = rand::thread_rng(); + let micros = rng.gen_range(0..1_000); + Duration::from_micros(micros) + }; + for value in data { + sleep(random_duration()).await; + super::ring_buf::ring_buf_trigger_ebpf_program(value); + } +}