mirror of https://github.com/aya-rs/aya
				
				
				
			WIP: auto-attach tracepoints based on SEC() data
							parent
							
								
									29b821376e
								
							
						
					
					
						commit
						ef73d65229
					
				| @ -0,0 +1,41 @@ | |||||||
|  | // clang-format off
 | ||||||
|  | #include <vmlinux.h> | ||||||
|  | #include <bpf/bpf_helpers.h> | ||||||
|  | #include <bpf/bpf_tracing.h> | ||||||
|  | // clang-format on
 | ||||||
|  | 
 | ||||||
|  | struct { | ||||||
|  |   __uint(type, BPF_MAP_TYPE_HASH); | ||||||
|  |   __uint(max_entries, 256); | ||||||
|  |   __type(key, char[32]); | ||||||
|  |   __type(value, __u8); | ||||||
|  | } executed_once SEC(".maps"); | ||||||
|  | 
 | ||||||
|  | #define assign_str(target, str) | ||||||
|  | 
 | ||||||
|  | // BPF will not allow us to write out of bounds, so we skip the length checks
 | ||||||
|  | #define mark_executed(key)                                                     \ | ||||||
|  |   {                                                                            \ | ||||||
|  |     __u8 __executed = 1;                                                       \ | ||||||
|  |     char __probe_type[32] = {};                                                \ | ||||||
|  |     __builtin_memcpy(__probe_type, key, sizeof(key));                          \ | ||||||
|  |     bpf_map_update_elem(&executed_once, &__probe_type, &__executed, BPF_ANY);  \ | ||||||
|  |   }                                                                            \ | ||||||
|  |   do {                                                                         \ | ||||||
|  |   } while (0) | ||||||
|  | 
 | ||||||
|  | SEC("tp_btf/sched_switch") | ||||||
|  | int BPF_PROG(sched_switch_tp_btf, bool preempt, struct task_struct *prev, | ||||||
|  |              struct task_struct *next) { | ||||||
|  |   mark_executed("tp_btf"); | ||||||
|  |   return 0; | ||||||
|  | } | ||||||
|  | 
 | ||||||
|  | SEC("tracepoint/sched/sched_switch") | ||||||
|  | int sched_switch_tp(bool preempt, struct task_struct *prev, | ||||||
|  |                     struct task_struct *next) { | ||||||
|  |   mark_executed("tracepoint"); | ||||||
|  |   return 0; | ||||||
|  | } | ||||||
|  | 
 | ||||||
|  | char _license[] SEC("license") = "GPL"; | ||||||
| @ -0,0 +1,42 @@ | |||||||
|  | use std::{ | ||||||
|  |     ffi, | ||||||
|  |     time::{Duration, Instant}, | ||||||
|  | }; | ||||||
|  | 
 | ||||||
|  | use aya::{maps::MapError, Ebpf}; | ||||||
|  | 
 | ||||||
|  | fn to_map_key(key: &str) -> [ffi::c_char; 32] { | ||||||
|  |     let mut padded: Vec<_> = key.bytes().map(|b| b as ffi::c_char).collect(); | ||||||
|  |     padded.resize(32, 0); | ||||||
|  |     padded.try_into().unwrap() | ||||||
|  | } | ||||||
|  | 
 | ||||||
|  | #[test] | ||||||
|  | fn auto_attach_succes() { | ||||||
|  |     let mut bpf = Ebpf::load(crate::PROBE_AUTO_ATTACH).unwrap(); | ||||||
|  |     bpf.auto_attach(); | ||||||
|  | 
 | ||||||
|  |     let executed_map: aya::maps::HashMap<aya::maps::MapData, [ffi::c_char; 32], u8> = | ||||||
|  |         aya::maps::HashMap::try_from(bpf.take_map("executed_once").unwrap()).unwrap(); | ||||||
|  | 
 | ||||||
|  |     let fired_probes = ["tp_btf", "tracepoint"]; | ||||||
|  | 
 | ||||||
|  |     let start = Instant::now(); | ||||||
|  |     const TIMEOUT: Duration = Duration::from_secs(1); | ||||||
|  | 
 | ||||||
|  |     let mut all_fired = false; | ||||||
|  |     while !all_fired && (Instant::now() - start) < TIMEOUT { | ||||||
|  |         all_fired = true; | ||||||
|  |         for probe in fired_probes { | ||||||
|  |             let executed = match executed_map.get(&to_map_key(probe), 0) { | ||||||
|  |                 Ok(fired) => fired, | ||||||
|  |                 Err(MapError::KeyNotFound) => 0, | ||||||
|  |                 e => e.unwrap(), | ||||||
|  |             }; | ||||||
|  |             if executed == 0 { | ||||||
|  |                 all_fired = false; | ||||||
|  |             } | ||||||
|  |         } | ||||||
|  |     } | ||||||
|  |     assert!(all_fired, "Not all expected probes fired"); | ||||||
|  | } | ||||||
					Loading…
					
					
				
		Reference in New Issue