1 //! Validates uprobestats config protos and adds additional info.
2 use anyhow::{anyhow, Result};
3 use dynamic_instrumentation_manager::{
4 ExecutableMethodFileOffsets, MethodDescriptor, TargetProcess,
5 };
6 use log::debug;
7 use protobuf::Message;
8 use std::collections::HashSet;
9 use std::fs::File;
10 use std::io::Read;
11 use uprobestats_proto::config::{
12 uprobestats_config::task::ProbeConfig, uprobestats_config::Task, UprobestatsConfig,
13 };
14
15 use crate::{art::get_method_offset_from_oatdump, process::get_pid};
16
17 /// Validated probe proto + probe target's code filename and offset.
18 pub struct ResolvedProbe {
19 _probe: ProbeConfig,
20 /// The filename of the code that contains the probe's method.
21 pub filename: String,
22 /// The offset of the probe's method in the code file.
23 pub offset: i32,
24 /// Absolute path to the bpf program.
25 pub bpf_program_path: String,
26 }
27
28 /// Validated task proto + probe target's pid.
29 pub struct ResolvedTask {
30 /// The task proto.
31 pub task: Task,
32 /// The duration of the task in seconds.
33 pub duration_seconds: i32,
34 /// The pid of the task's target process.
35 pub pid: i32,
36 /// The set of absolute bpf map paths used by the task.
37 pub bpf_map_paths: HashSet<String>,
38 }
39
40 /// Validates a single task proto and adds additional info.
resolve_single_task(config: UprobestatsConfig) -> Result<ResolvedTask>41 pub fn resolve_single_task(config: UprobestatsConfig) -> Result<ResolvedTask> {
42 let mut tasks = config.tasks.into_iter();
43 let task = tasks.next().ok_or_else(|| anyhow!("No tasks found in config"))?;
44
45 let duration_seconds =
46 task.duration_seconds.ok_or_else(|| anyhow!("Task duration is required"))?;
47 if duration_seconds <= 0 {
48 return Err(anyhow!("Task duration must be greater than 0"));
49 }
50
51 let target_process_name = task
52 .target_process_name
53 .clone()
54 .ok_or_else(|| anyhow!("Target process name is required"))?;
55 if target_process_name != "system_server" {
56 return Err(anyhow!("system_server is the only target process currently supported"));
57 }
58
59 let pid = get_pid(&target_process_name)
60 .ok_or_else(|| anyhow!("Failed to get pid for process: {target_process_name}"))?;
61
62 let bpf_map_paths = task.bpf_maps.iter().map(|bpf_map| prefix_bpf(bpf_map)).collect();
63
64 Ok(ResolvedTask { duration_seconds, task, pid, bpf_map_paths })
65 }
66
67 /// Validates a single probe proto and adds additional info.
resolve_probes(task: &Task) -> Result<Vec<ResolvedProbe>>68 pub fn resolve_probes(task: &Task) -> Result<Vec<ResolvedProbe>> {
69 let resolved_probes = task.probe_configs.clone().into_iter().map(|probe| {
70 let bpf_name = probe.bpf_name.as_ref().ok_or_else(|| anyhow!("bpf_name is required"))?;
71 let bpf_program_path = prefix_bpf(bpf_name);
72 if let Some(ref fully_qualified_class_name) = probe.fully_qualified_class_name {
73 debug!("using getExecutableMethodFileOffsets to retrieve offsets");
74 let method_name =
75 probe.method_name.clone().ok_or_else(|| anyhow!("method_name is required"))?;
76 let fully_qualified_parameters = probe.fully_qualified_parameters.clone();
77 let offsets = ExecutableMethodFileOffsets::get(
78 &TargetProcess::system_server()?,
79 &MethodDescriptor::new(
80 &fully_qualified_class_name.clone(),
81 &method_name,
82 fully_qualified_parameters,
83 )?,
84 )?;
85 let offsets = offsets.ok_or_else(|| {
86 anyhow!("Failed to get offsets for class: {fully_qualified_class_name}")
87 })?;
88 let offset: i32 = offsets
89 .get_method_offset()
90 .try_into()
91 .map_err(|e| anyhow!("Failed to convert method offset to i32: {e}"))?;
92 Ok(ResolvedProbe {
93 _probe: probe,
94 bpf_program_path,
95 offset,
96 filename: offsets.get_container_path(),
97 })
98 } else {
99 debug!("using oatdump to retrieve offsets");
100 let method_signature =
101 probe.method_signature.clone().ok_or(anyhow!("method_signature is required"))?;
102 let mut offset: i32 = 0;
103 let mut found_file_path: String = "".to_string();
104 for file_path in &probe.file_paths {
105 let found_offset = get_method_offset_from_oatdump(file_path, &method_signature)?;
106 let Some(found_offset) = found_offset else {
107 continue;
108 };
109 if found_offset > 0 {
110 found_file_path = file_path.to_string();
111 offset = found_offset;
112 break;
113 }
114 }
115 if offset > 0 {
116 Ok(ResolvedProbe {
117 _probe: probe,
118 bpf_program_path,
119 filename: found_file_path,
120 offset,
121 })
122 } else {
123 Err(anyhow!("Failed to get offset for method: {method_signature}"))
124 }
125 }
126 });
127
128 resolved_probes.collect()
129 }
130
131 /// Reads a config file and parses it into a UprobestatsConfig proto.
read_config(config_path: &str) -> Result<UprobestatsConfig>132 pub fn read_config(config_path: &str) -> Result<UprobestatsConfig> {
133 let mut file =
134 File::open(config_path).map_err(|e| anyhow!("Failed to open config file: {e}"))?;
135 let mut buffer = Vec::new();
136 file.read_to_end(&mut buffer).map_err(|e| anyhow!("Failed to read config file: {e}"))?;
137 UprobestatsConfig::parse_from_bytes(&buffer)
138 .map_err(|e| anyhow!("Failed to parse config file: {e}"))
139 }
140
141 const BPF_DIR: &str = "/sys/fs/bpf/uprobestats/";
prefix_bpf(path: &str) -> String142 fn prefix_bpf(path: &str) -> String {
143 BPF_DIR.to_string() + path
144 }
145