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

@@ -31,3 +31,7 @@ path = "src/test.rs"
 [[bin]]
 name = "relocations"
 path = "src/relocations.rs"
+
+[[bin]]
+name = "ring_buf"
+path = "src/ring_buf.rs"

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::<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 {}
+}

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" }

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)]

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::<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()
+}

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<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);
+    }
+}