Skip to content

Commit

Permalink
integration-test: add RingBuf tests
Browse files Browse the repository at this point in the history
This simple test exercises sending data from a uprobe. It also adds a
test to demonstrate how to combine the RingBuf with async notifications.
  • Loading branch information
ajwerner committed Jul 2, 2023
1 parent 406e949 commit c843ef2
Show file tree
Hide file tree
Showing 6 changed files with 174 additions and 1 deletion.
4 changes: 4 additions & 0 deletions test/integration-ebpf/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,7 @@ path = "src/test.rs"
[[bin]]
name = "relocations"
path = "src/relocations.rs"

[[bin]]
name = "ring_buf"
path = "src/ring_buf.rs"
38 changes: 38 additions & 0 deletions test/integration-ebpf/src/ring_buf.rs
Original file line number Diff line number Diff line change
@@ -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::<Entry>() + 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<Entry> = ctx.arg(0) else { return };
if let Some(mut entry) = RING_BUF.reserve::<Entry>(0) {
entry.write(arg);
entry.submit(0);
}
}

#[panic_handler]
fn panic(_info: &core::panic::PanicInfo) -> ! {
loop {}
}
1 change: 1 addition & 0 deletions test/integration-test/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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" }
3 changes: 2 additions & 1 deletion test/integration-test/src/tests/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
57 changes: 57 additions & 0 deletions test/integration-test/src/tests/ring_buf.rs
Original file line number Diff line number Diff line change
@@ -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::<u64>::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<u64> {
const DATA_LEN_RANGE: core::ops::RangeInclusive<usize> = 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()
}
72 changes: 72 additions & 0 deletions test/integration-test/src/tests/ring_buf_async.rs
Original file line number Diff line number Diff line change
@@ -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<u64>) {
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);
}
}

0 comments on commit c843ef2

Please sign in to comment.