1 use anyhow::{bail, Result};
2 use log::debug;
3 use std::{
4 collections::HashMap,
5 fmt::Debug,
6 sync::LazyLock,
7 time::{Duration, Instant},
8 };
9 use uprobestats_bpf::poll_ring_buf;
10 use uprobestats_bpf_bindgen::{
11 CallResult, CallTimestamp, MalwareSignal, SetUidTempAllowlistStateRecord,
12 UpdateDeviceIdleTempAllowlistRecord,
13 };
14 use uprobestats_proto::config::uprobestats_config::Task;
15
16 mod generic_instrumentation;
17 mod malware_signal;
18 mod process_management;
19
poll_and_loop( map_path: &str, now: Instant, duration: Duration, task: Task, ) -> Result<()>20 pub(crate) fn poll_and_loop(
21 map_path: &str,
22 now: Instant,
23 duration: Duration,
24 task: Task,
25 ) -> Result<()> {
26 let duration_millis = duration.as_millis();
27 let mut elapsed_millis = now.elapsed().as_millis();
28 while elapsed_millis <= duration_millis {
29 let timeout_millis = duration_millis - elapsed_millis;
30 let timeout_millis: i32 = timeout_millis.try_into()?;
31 debug!("polling {} for {} seconds", map_path, timeout_millis / 1000);
32 let Some(do_poll) = REGISTRY.get(map_path) else {
33 bail!("unsupported map_path: {}", map_path);
34 };
35 do_poll(map_path, timeout_millis, &task)?;
36 elapsed_millis = now.elapsed().as_millis();
37 }
38 Ok(())
39 }
40
poll<T: OnItem + Debug + Copy>(map_path: &str, timeout_millis: i32, task: &Task) -> Result<()>41 fn poll<T: OnItem + Debug + Copy>(map_path: &str, timeout_millis: i32, task: &Task) -> Result<()> {
42 if map_path != T::MAP_PATH {
43 bail!("map_path mismatch: {} != {}", map_path, T::MAP_PATH)
44 }
45 // SAFETY: we've just checked that the passed `map_path` is the same as the one
46 // expected by the `OnItem` implementation, which encodes how the expected type is mapped to the
47 // ring buffer's path.
48 let result: Result<Vec<T>> = unsafe { poll_ring_buf(map_path, timeout_millis) };
49 let result = result?;
50 debug!("Done polling {}, event count: {}", map_path, result.len());
51 for i in &result {
52 i.on_item(task)?;
53 }
54 Ok(())
55 }
56
57 const JAVA_ARGUMENT_REGISTER_OFFSET: i32 = 2;
58
59 /// Interface for reading items out of a BPF ring buffer.
60 /// # Safety
61 /// There *must* exist a BPF ring buffer at the path represented by `MAP_PATH`
62 /// which holds items of type `T` implementing this trait.
63 unsafe trait OnItem {
64 const MAP_PATH: &'static str;
on_item(&self, task: &Task) -> Result<()>65 fn on_item(&self, task: &Task) -> Result<()>;
66 }
67
68 type Registry = HashMap<&'static str, fn(&str, i32, &Task) -> Result<()>>;
69
register<T: OnItem + Debug + Copy>(registry: &mut Registry)70 fn register<T: OnItem + Debug + Copy>(registry: &mut Registry) {
71 registry.insert(T::MAP_PATH, poll::<T> as _);
72 }
73
74 static REGISTRY: LazyLock<Registry> = LazyLock::new(|| {
75 let mut map = HashMap::new();
76 register::<CallTimestamp>(&mut map);
77 register::<CallResult>(&mut map);
78 register::<MalwareSignal>(&mut map);
79 register::<SetUidTempAllowlistStateRecord>(&mut map);
80 register::<UpdateDeviceIdleTempAllowlistRecord>(&mut map);
81 map
82 });
83