diff --git a/test/integration-test/src/tests/queue.rs b/test/integration-test/src/tests/queue.rs index e8614ccd..f922e2f2 100644 --- a/test/integration-test/src/tests/queue.rs +++ b/test/integration-test/src/tests/queue.rs @@ -1,5 +1,7 @@ use aya::{EbpfLoader, maps::Array, programs::UProbe}; +use crate::utils::attach_uprobe; + const PEEK_INDEX: u32 = 0; const POP_INDEX: u32 = 1; @@ -27,32 +29,13 @@ pub extern "C" fn trigger_queue_pop(marker: u64) -> u64 { fn queue_basic() { let mut bpf = EbpfLoader::new().load(crate::QUEUE_TEST).unwrap(); - let prog: &mut UProbe = bpf - .program_mut("test_queue_push") - .unwrap() - .try_into() - .unwrap(); - prog.load().unwrap(); - prog.attach("trigger_queue_push", "/proc/self/exe", None, None) - .unwrap(); - - let prog: &mut UProbe = bpf - .program_mut("test_queue_pop") - .unwrap() - .try_into() - .unwrap(); - prog.load().unwrap(); - prog.attach("trigger_queue_pop", "/proc/self/exe", None, None) - .unwrap(); - - let prog: &mut UProbe = bpf - .program_mut("test_queue_peek") - .unwrap() - .try_into() - .unwrap(); - prog.load().unwrap(); - prog.attach("trigger_queue_peek", "/proc/self/exe", None, None) - .unwrap(); + for (probe_name, symbol) in &[ + ("test_queue_push", "trigger_queue_push"), + ("test_queue_pop", "trigger_queue_pop"), + ("test_queue_peek", "trigger_queue_peek"), + ] { + attach_uprobe(bpf, probe_name, symbol); + } let array_map = bpf.map("RESULT").unwrap(); let array = Array::<_, u64>::try_from(array_map).unwrap(); diff --git a/test/integration-test/src/tests/stack.rs b/test/integration-test/src/tests/stack.rs index 4547b6f5..fd2e5606 100644 --- a/test/integration-test/src/tests/stack.rs +++ b/test/integration-test/src/tests/stack.rs @@ -1,5 +1,7 @@ use aya::{EbpfLoader, maps::Array, programs::UProbe}; +use crate::utils::attach_uprobe; + const PEEK_INDEX: u32 = 0; const POP_INDEX: u32 = 1; @@ -27,32 +29,13 @@ pub extern "C" fn trigger_stack_pop(marker: u64) -> u64 { fn stack_basic() { let mut bpf = EbpfLoader::new().load(crate::STACK_TEST).unwrap(); - let prog: &mut UProbe = bpf - .program_mut("test_stack_push") - .unwrap() - .try_into() - .unwrap(); - prog.load().unwrap(); - prog.attach("trigger_stack_push", "/proc/self/exe", None, None) - .unwrap(); - - let prog: &mut UProbe = bpf - .program_mut("test_stack_pop") - .unwrap() - .try_into() - .unwrap(); - prog.load().unwrap(); - prog.attach("trigger_stack_pop", "/proc/self/exe", None, None) - .unwrap(); - - let prog: &mut UProbe = bpf - .program_mut("test_stack_peek") - .unwrap() - .try_into() - .unwrap(); - prog.load().unwrap(); - prog.attach("trigger_stack_peek", "/proc/self/exe", None, None) - .unwrap(); + for (probe_name, symbol) in &[ + ("test_stack_push", "trigger_stack_push"), + ("test_stack_peek", "trigger_stack_peek"), + ("test_stack_pop", "trigger_stack_pop"), + ] { + attach_uprobe(bpf, probe_name, symbol); + } let array_map = bpf.map("RESULT").unwrap(); let array = Array::<_, u64>::try_from(array_map).unwrap(); diff --git a/test/integration-test/src/utils.rs b/test/integration-test/src/utils.rs index 924ab9a6..303c4b64 100644 --- a/test/integration-test/src/utils.rs +++ b/test/integration-test/src/utils.rs @@ -101,3 +101,14 @@ macro_rules! kernel_assert_eq { } pub(crate) use kernel_assert_eq; + +pub(crate) fn attach_uprobe(bpf: &mut Bpf, probe_name: &str, symbol: &str) { + let prog: &mut UProbe = bpf + .program_mut("test_stack_push") + .unwrap() + .try_into() + .unwrap(); + prog.load().unwrap(); + prog.attach("trigger_stack_push", "/proc/self/exe", None, None) + .unwrap(); +}