From d31a1805daa1e44b9f91fcbf9230a56feadcf11b Mon Sep 17 00:00:00 2001 From: Michal Rostecki Date: Mon, 20 Mar 2023 12:41:54 +0100 Subject: [PATCH 1/2] integration-test: Add `tokio_integration_test` macro This new macro runs a test in a Tokio runtime and it can be used for asynchronous tests (defined as `async fn`). --- test/integration-test-macros/Cargo.toml | 1 + test/integration-test-macros/src/lib.rs | 29 ++++++++++++++++++++++++- test/integration-test/Cargo.toml | 2 ++ test/integration-test/src/tests/mod.rs | 3 ++- 4 files changed, 33 insertions(+), 2 deletions(-) diff --git a/test/integration-test-macros/Cargo.toml b/test/integration-test-macros/Cargo.toml index e28cdc5c..29788ccb 100644 --- a/test/integration-test-macros/Cargo.toml +++ b/test/integration-test-macros/Cargo.toml @@ -6,6 +6,7 @@ publish = false [dependencies] quote = "1" +proc-macro2 = "1.0" syn = {version = "1.0", features = ["full"]} [lib] diff --git a/test/integration-test-macros/src/lib.rs b/test/integration-test-macros/src/lib.rs index 297159ed..9818b7f5 100644 --- a/test/integration-test-macros/src/lib.rs +++ b/test/integration-test-macros/src/lib.rs @@ -1,6 +1,7 @@ use proc_macro::TokenStream; +use proc_macro2::Span; use quote::quote; -use syn::{parse_macro_input, ItemFn}; +use syn::{parse_macro_input, Ident, ItemFn}; #[proc_macro_attribute] pub fn integration_test(_attr: TokenStream, item: TokenStream) -> TokenStream { @@ -17,3 +18,29 @@ pub fn integration_test(_attr: TokenStream, item: TokenStream) -> TokenStream { }; TokenStream::from(expanded) } + +#[proc_macro_attribute] +pub fn tokio_integration_test(_attr: TokenStream, item: TokenStream) -> TokenStream { + let item = parse_macro_input!(item as ItemFn); + let name = &item.sig.ident; + let name_str = &item.sig.ident.to_string(); + let sync_name_str = format!("sync_{name_str}"); + let sync_name = Ident::new(&sync_name_str, Span::call_site()); + let expanded = quote! { + #item + + fn #sync_name() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + rt.block_on(#name()); + } + + inventory::submit!(crate::IntegrationTest { + name: concat!(module_path!(), "::", #sync_name_str), + test_fn: #sync_name, + }); + }; + TokenStream::from(expanded) +} diff --git a/test/integration-test/Cargo.toml b/test/integration-test/Cargo.toml index 6b02b4a8..c8f7ca02 100644 --- a/test/integration-test/Cargo.toml +++ b/test/integration-test/Cargo.toml @@ -10,6 +10,7 @@ aya = { path = "../../aya" } aya-obj = { path = "../../aya-obj" } clap = { version = "4", features = ["derive"] } env_logger = "0.10" +futures-core = "0.3" inventory = "0.3" integration-test-macros = { path = "../integration-test-macros" } lazy_static = "1" @@ -20,3 +21,4 @@ rbpf = "0.1.0" regex = "1" tempfile = "3.3.0" libtest-mimic = "0.6.0" +tokio = { version = "1.24", features = ["rt", "rt-multi-thread"] } diff --git a/test/integration-test/src/tests/mod.rs b/test/integration-test/src/tests/mod.rs index 127b037d..a3a21c46 100644 --- a/test/integration-test/src/tests/mod.rs +++ b/test/integration-test/src/tests/mod.rs @@ -11,7 +11,8 @@ pub mod rbpf; pub mod relocations; pub mod smoke; -pub use integration_test_macros::integration_test; +pub use integration_test_macros::{integration_test, tokio_integration_test}; + #[derive(Debug)] pub struct IntegrationTest { pub name: &'static str, From 5fa17a192bf295c4977d4a86d64efe70d83e8d51 Mon Sep 17 00:00:00 2001 From: Michal Rostecki Date: Mon, 20 Mar 2023 12:11:34 +0100 Subject: [PATCH 2/2] integration-test: Add tests for aya-log --- test/integration-ebpf/Cargo.toml | 5 + test/integration-ebpf/src/log.rs | 25 +++++ test/integration-test/Cargo.toml | 3 +- test/integration-test/src/tests/log.rs | 136 +++++++++++++++++++++++++ test/integration-test/src/tests/mod.rs | 1 + 5 files changed, 169 insertions(+), 1 deletion(-) create mode 100644 test/integration-ebpf/src/log.rs create mode 100644 test/integration-test/src/tests/log.rs diff --git a/test/integration-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml index 69c60316..0683143a 100644 --- a/test/integration-ebpf/Cargo.toml +++ b/test/integration-ebpf/Cargo.toml @@ -6,6 +6,11 @@ publish = false [dependencies] aya-bpf = { path = "../../bpf/aya-bpf" } +aya-log-ebpf = { path = "../../bpf/aya-log-ebpf" } + +[[bin]] +name = "log" +path = "src/log.rs" [[bin]] name = "map_test" diff --git a/test/integration-ebpf/src/log.rs b/test/integration-ebpf/src/log.rs new file mode 100644 index 00000000..d4570103 --- /dev/null +++ b/test/integration-ebpf/src/log.rs @@ -0,0 +1,25 @@ +#![no_std] +#![no_main] + +use aya_bpf::{macros::uprobe, programs::ProbeContext}; +use aya_log_ebpf::{debug, error, info, trace, warn}; + +#[uprobe] +pub fn test_log(ctx: ProbeContext) { + debug!(&ctx, "Hello from eBPF!"); + error!(&ctx, "{}, {}, {}", 69, 420i32, "wao"); + let ipv4 = 167772161u32; // 10.0.0.1 + let ipv6 = [ + 32u8, 1u8, 13u8, 184u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 1u8, + ]; // 2001:db8::1 + info!(&ctx, "ipv4: {:ipv4}, ipv6: {:ipv6}", ipv4, ipv6); + let mac = [4u8, 32u8, 6u8, 9u8, 0u8, 64u8]; + trace!(&ctx, "mac lc: {:mac}, mac uc: {:MAC}", mac, mac); + let hex = 0x2f; + warn!(&ctx, "hex lc: {:x}, hex uc: {:X}", hex, hex); +} + +#[panic_handler] +fn panic(_info: &core::panic::PanicInfo) -> ! { + unsafe { core::hint::unreachable_unchecked() } +} diff --git a/test/integration-test/Cargo.toml b/test/integration-test/Cargo.toml index c8f7ca02..96a3f793 100644 --- a/test/integration-test/Cargo.toml +++ b/test/integration-test/Cargo.toml @@ -7,6 +7,7 @@ publish = false [dependencies] anyhow = "1" aya = { path = "../../aya" } +aya-log = { path = "../../aya-log" } aya-obj = { path = "../../aya-obj" } clap = { version = "4", features = ["derive"] } env_logger = "0.10" @@ -21,4 +22,4 @@ rbpf = "0.1.0" regex = "1" tempfile = "3.3.0" libtest-mimic = "0.6.0" -tokio = { version = "1.24", features = ["rt", "rt-multi-thread"] } +tokio = { version = "1.24", features = ["rt", "rt-multi-thread", "sync", "time"] } diff --git a/test/integration-test/src/tests/log.rs b/test/integration-test/src/tests/log.rs new file mode 100644 index 00000000..fd99d859 --- /dev/null +++ b/test/integration-test/src/tests/log.rs @@ -0,0 +1,136 @@ +use std::sync::{Arc, LockResult, Mutex, MutexGuard}; + +use aya::{include_bytes_aligned, programs::UProbe, Bpf}; +use aya_log::BpfLogger; +use log::{Level, Log, Record}; +use tokio::time::{sleep, Duration}; + +use super::tokio_integration_test; + +const MAX_ATTEMPTS: usize = 10; +const TIMEOUT_MS: u64 = 10; + +#[no_mangle] +#[inline(never)] +pub extern "C" fn trigger_ebpf_program() {} + +struct CapturedLogs(Arc>>); + +impl CapturedLogs { + fn with_capacity(capacity: usize) -> Self { + Self(Arc::new(Mutex::new(Vec::with_capacity(capacity)))) + } + + fn clone(&self) -> Self { + Self(self.0.clone()) + } + + fn lock(&self) -> LockResult>> { + self.0.lock() + } + + async fn wait_expected_len(&self, expected_len: usize) { + for _ in 0..MAX_ATTEMPTS { + { + let captured_logs = self.0.lock().expect("Failed to lock captured logs"); + if captured_logs.len() == expected_len { + return; + } + } + sleep(Duration::from_millis(TIMEOUT_MS)).await; + } + panic!( + "Expected {} captured logs, but got {}", + expected_len, + self.0.lock().unwrap().len() + ); + } +} + +struct CapturedLog { + pub body: String, + pub level: Level, + pub target: String, +} + +struct TestingLogger { + captured_logs: CapturedLogs, +} + +impl TestingLogger { + pub fn with_capacity(capacity: usize) -> (Self, CapturedLogs) { + let captured_logs = CapturedLogs::with_capacity(capacity); + ( + Self { + captured_logs: captured_logs.clone(), + }, + captured_logs, + ) + } +} + +impl Log for TestingLogger { + fn enabled(&self, _metadata: &log::Metadata) -> bool { + true + } + + fn flush(&self) {} + + fn log(&self, record: &Record) { + let captured_record = CapturedLog { + body: format!("{}", record.args()), + level: record.level(), + target: record.target().to_string(), + }; + self.captured_logs + .lock() + .expect("Failed to acquire a lock for storing a log") + .push(captured_record); + } +} + +#[tokio_integration_test] +async fn log() { + let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/release/log"); + let mut bpf = Bpf::load(bytes).unwrap(); + + let (logger, captured_logs) = TestingLogger::with_capacity(5); + BpfLogger::init_with_logger(&mut bpf, logger).unwrap(); + + let prog: &mut UProbe = bpf.program_mut("test_log").unwrap().try_into().unwrap(); + prog.load().unwrap(); + prog.attach(Some("trigger_ebpf_program"), 0, "/proc/self/exe", None) + .unwrap(); + + // Call the function that the uprobe is attached to, so it starts logging. + trigger_ebpf_program(); + captured_logs.wait_expected_len(5).await; + + let records = captured_logs + .lock() + .expect("Failed to acquire a lock for reading logs"); + assert_eq!(records.len(), 5); + + assert_eq!(records[0].body, "Hello from eBPF!"); + assert_eq!(records[0].level, Level::Debug); + assert_eq!(records[0].target, "log"); + + assert_eq!(records[1].body, "69, 420, wao"); + assert_eq!(records[1].level, Level::Error); + assert_eq!(records[1].target, "log"); + + assert_eq!(records[2].body, "ipv4: 10.0.0.1, ipv6: 2001:db8::1"); + assert_eq!(records[2].level, Level::Info); + assert_eq!(records[2].target, "log"); + + assert_eq!( + records[3].body, + "mac lc: 04:20:06:09:00:40, mac uc: 04:20:06:09:00:40" + ); + assert_eq!(records[3].level, Level::Trace); + assert_eq!(records[3].target, "log"); + + assert_eq!(records[4].body, "hex lc: 2f, hex uc: 2F"); + assert_eq!(records[4].level, Level::Warn); + assert_eq!(records[4].target, "log"); +} diff --git a/test/integration-test/src/tests/mod.rs b/test/integration-test/src/tests/mod.rs index a3a21c46..c26ca5a4 100644 --- a/test/integration-test/src/tests/mod.rs +++ b/test/integration-test/src/tests/mod.rs @@ -7,6 +7,7 @@ use std::{ffi::CStr, mem}; pub mod btf_relocations; pub mod elf; pub mod load; +pub mod log; pub mod rbpf; pub mod relocations; pub mod smoke;