add attach_uprobe function

Signed-off-by: Xiaobo Liu <cppcoffee@gmail.com>

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
+    for (probe_name, symbol) in &[
-        .program_mut("test_queue_push")
+        ("test_queue_push", "trigger_queue_push"),
-        .unwrap()
+        ("test_queue_pop", "trigger_queue_pop"),
-        .try_into()
+        ("test_queue_peek", "trigger_queue_peek"),
-        .unwrap();
+    ] {
-    prog.load().unwrap();
+        attach_uprobe(bpf, probe_name, symbol);
-    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();
 
     let array_map = bpf.map("RESULT").unwrap();
     let array = Array::<_, u64>::try_from(array_map).unwrap();

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
+    for (probe_name, symbol) in &[
-        .program_mut("test_stack_push")
+        ("test_stack_push", "trigger_stack_push"),
-        .unwrap()
+        ("test_stack_peek", "trigger_stack_peek"),
-        .try_into()
+        ("test_stack_pop", "trigger_stack_pop"),
-        .unwrap();
+    ] {
-    prog.load().unwrap();
+        attach_uprobe(bpf, probe_name, symbol);
-    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();
 
     let array_map = bpf.map("RESULT").unwrap();
     let array = Array::<_, u64>::try_from(array_map).unwrap();

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