1use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
7use std::path::PathBuf;
8
9use crate::parser::{normalize_session_log_path, session_log_path_from_str};
10use crate::{
11 SOURCE_SESSION_PROCESS_MATCH, TRACE_EBPF_FILE, TRACE_PROC_FD, TRACE_RECENT_CWD,
12 TRACE_STICKY_BINDING,
13};
14
15const SESSION_PROCESS_START_SKEW_MS: u64 = 30_000;
16
17#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
19pub struct ProcessKey {
20 pub pid: u32,
21 pub starttime_ticks: u64,
22}
23
24#[derive(Debug, Clone, Default)]
26pub struct ProcessTree {
27 pub root: ProcessKey,
28 pub members: Vec<ProcessKey>,
29}
30
31#[derive(Debug, Clone, Default)]
33pub struct LiveProcessCandidate {
34 pub tree: ProcessTree,
35 pub agent: String,
36 pub age_s: Option<f64>,
37 pub cwd: Option<String>,
38}
39
40#[derive(Debug, Clone)]
42pub struct SessionProcessInput {
43 pub id: String,
44 pub agent: String,
45 pub path: PathBuf,
46 pub start_timestamp_ms: Option<u64>,
47 pub end_timestamp_ms: Option<u64>,
48 pub cwd: Option<String>,
49}
50
51#[derive(Debug, Clone)]
53pub struct SessionProcessMatch {
54 pub session_id: String,
55 pub root_pid: u32,
56 pub matched_pids: Vec<u32>,
57 pub pid_starttime_ticks: u64,
58 pub source: &'static str,
59 pub confidence: f32,
60 pub evidence: &'static str,
61}
62
63#[derive(Debug, Default)]
65pub struct SessionProcessMatches {
66 pub by_session_id: HashMap<String, SessionProcessMatch>,
67 pub by_pid: HashMap<u32, String>,
68 pub used_root_pids: HashSet<u32>,
69}
70
71impl SessionProcessMatches {
72 pub fn session_for_pid(&self, pid: u32) -> Option<&SessionProcessMatch> {
73 self.by_pid
74 .get(&pid)
75 .and_then(|session_id| self.by_session_id.get(session_id))
76 }
77}
78
79#[derive(Default)]
81pub struct SessionProcessMatcher {
82 bindings: HashMap<u32, LiveSessionBinding>,
83}
84
85struct LiveSessionBinding {
86 starttime_ticks: u64,
87 session_path: PathBuf,
88}
89
90impl SessionProcessMatcher {
91 pub fn match_sessions(
92 &mut self,
93 sessions: &[SessionProcessInput],
94 processes: &[LiveProcessCandidate],
95 fd_paths_by_process: &HashMap<ProcessKey, BTreeSet<PathBuf>>,
96 observed_path_by_process: &HashMap<ProcessKey, PathBuf>,
97 now_ms: u64,
98 ) -> SessionProcessMatches {
99 let path_evidence =
100 collect_path_evidence(processes, fd_paths_by_process, observed_path_by_process);
101 self.retain_live(processes);
102
103 let mut out = SessionProcessMatches::default();
104 for session in sessions {
105 let Some((process, evidence)) = processes.iter().find_map(|process| {
106 if out.used_root_pids.contains(&process.tree.root.pid)
107 || process.agent != session.agent
108 || !session_is_fresh_enough_for_process(session, process, now_ms)
109 {
110 return None;
111 }
112 self.link_trace(&session.path, process, &path_evidence)
113 .map(|evidence| (process, evidence))
114 }) else {
115 continue;
116 };
117 record_match(&mut out, session, process, evidence);
118 }
119
120 let mut cwd_candidates = Vec::new();
121 for (session_index, session) in sessions.iter().enumerate() {
122 if out.by_session_id.contains_key(&session.id) {
123 continue;
124 }
125 for (process_index, process) in processes.iter().enumerate() {
126 if out.used_root_pids.contains(&process.tree.root.pid)
127 || process.agent != session.agent
128 || !self.can_use_cwd_trace(&session.path, process, &path_evidence)
129 {
130 continue;
131 }
132 let Some(distance_ms) = recent_cwd_distance_ms(session, process, now_ms) else {
133 continue;
134 };
135 cwd_candidates.push((
136 distance_ms,
137 std::cmp::Reverse(session_end_ms(session)),
138 session_index,
139 process_index,
140 ));
141 }
142 }
143 cwd_candidates.sort_unstable();
144 for (_, _, session_index, process_index) in cwd_candidates {
145 let session = &sessions[session_index];
146 let process = &processes[process_index];
147 if out.by_session_id.contains_key(&session.id)
148 || out.used_root_pids.contains(&process.tree.root.pid)
149 {
150 continue;
151 }
152 record_match(&mut out, session, process, TRACE_RECENT_CWD);
153 }
154 out
155 }
156
157 fn retain_live(&mut self, processes: &[LiveProcessCandidate]) {
158 self.bindings.retain(|pid, binding| {
159 processes.iter().any(|process| {
160 process.tree.root.pid == *pid
161 && process.tree.root.starttime_ticks == binding.starttime_ticks
162 })
163 });
164 }
165
166 fn link_trace(
167 &mut self,
168 session_path: &PathBuf,
169 process: &LiveProcessCandidate,
170 path_evidence: &HashMap<u32, BTreeMap<PathBuf, &'static str>>,
171 ) -> Option<&'static str> {
172 let pid = process.tree.root.pid;
173 let path = normalize_session_log_path(session_path);
174 if let Some(evidence) = path_evidence.get(&pid) {
175 if let Some(trace) = evidence.get(&path).copied() {
176 self.bindings.insert(
177 pid,
178 LiveSessionBinding {
179 starttime_ticks: process.tree.root.starttime_ticks,
180 session_path: path,
181 },
182 );
183 return Some(trace);
184 }
185 self.bindings.remove(&pid);
186 return None;
187 }
188 self.bindings
189 .get(&pid)
190 .filter(|binding| {
191 binding.starttime_ticks == process.tree.root.starttime_ticks
192 && binding.session_path == path
193 })
194 .map(|_| TRACE_STICKY_BINDING)
195 }
196
197 fn can_use_cwd_trace(
198 &self,
199 session_path: &PathBuf,
200 process: &LiveProcessCandidate,
201 path_evidence: &HashMap<u32, BTreeMap<PathBuf, &'static str>>,
202 ) -> bool {
203 let pid = process.tree.root.pid;
204 if path_evidence.contains_key(&pid) {
205 return false;
206 }
207 let path = normalize_session_log_path(session_path);
208 !self.bindings.get(&pid).is_some_and(|binding| {
209 binding.starttime_ticks == process.tree.root.starttime_ticks
210 && binding.session_path != path
211 })
212 }
213}
214
215fn record_match(
216 out: &mut SessionProcessMatches,
217 session: &SessionProcessInput,
218 process: &LiveProcessCandidate,
219 evidence: &'static str,
220) {
221 let matched_pids = process
222 .tree
223 .members
224 .iter()
225 .map(|key| key.pid)
226 .collect::<Vec<_>>();
227 out.used_root_pids.insert(process.tree.root.pid);
228 for pid in &matched_pids {
229 out.by_pid.insert(*pid, session.id.clone());
230 }
231 out.by_session_id.insert(
232 session.id.clone(),
233 SessionProcessMatch {
234 session_id: session.id.clone(),
235 root_pid: process.tree.root.pid,
236 matched_pids,
237 pid_starttime_ticks: process.tree.root.starttime_ticks,
238 source: SOURCE_SESSION_PROCESS_MATCH,
239 confidence: confidence_for_evidence(evidence),
240 evidence,
241 },
242 );
243}
244
245fn collect_path_evidence(
246 processes: &[LiveProcessCandidate],
247 fd_paths_by_process: &HashMap<ProcessKey, BTreeSet<PathBuf>>,
248 observed_path_by_process: &HashMap<ProcessKey, PathBuf>,
249) -> HashMap<u32, BTreeMap<PathBuf, &'static str>> {
250 let mut out = HashMap::new();
251 for process in processes {
252 let mut evidence = BTreeMap::new();
253 for key in &process.tree.members {
254 if let Some(paths) = fd_paths_by_process.get(key) {
255 for path in paths {
256 if let Some(session_path) = session_log_path_from_str(&path.to_string_lossy()) {
257 evidence.entry(session_path).or_insert(TRACE_PROC_FD);
258 }
259 }
260 }
261 if let Some(path) = observed_path_by_process.get(key)
262 && let Some(session_path) = session_log_path_from_str(&path.to_string_lossy())
263 {
264 evidence.insert(session_path, TRACE_EBPF_FILE);
265 }
266 }
267 if !evidence.is_empty() {
268 out.insert(process.tree.root.pid, evidence);
269 }
270 }
271 out
272}
273
274fn session_is_fresh_enough_for_process(
275 session: &SessionProcessInput,
276 process: &LiveProcessCandidate,
277 now_ms: u64,
278) -> bool {
279 let Some(process_start_ms) = process_start_ms(process, now_ms) else {
280 return true;
281 };
282 session_end_ms(session).saturating_add(SESSION_PROCESS_START_SKEW_MS) >= process_start_ms
283}
284
285fn recent_cwd_distance_ms(
286 session: &SessionProcessInput,
287 process: &LiveProcessCandidate,
288 now_ms: u64,
289) -> Option<u64> {
290 let session_cwd = session.cwd.as_deref().filter(|value| !value.is_empty())?;
291 let process_cwd = process.cwd.as_deref().filter(|value| !value.is_empty())?;
292 if session_cwd != process_cwd {
293 return None;
294 }
295 let process_start_ms = process_start_ms(process, now_ms)?;
296 let session_end_ms = session_end_ms(session);
297 (session_end_ms.saturating_add(SESSION_PROCESS_START_SKEW_MS) >= process_start_ms)
298 .then_some(session_end_ms.abs_diff(process_start_ms))
299}
300
301fn process_start_ms(process: &LiveProcessCandidate, now_ms: u64) -> Option<u64> {
302 process
303 .age_s
304 .map(|age_s| now_ms.saturating_sub((age_s.max(0.0) * 1000.0).round() as u64))
305}
306
307fn session_end_ms(session: &SessionProcessInput) -> u64 {
308 session
309 .end_timestamp_ms
310 .or(session.start_timestamp_ms)
311 .unwrap_or_default()
312}
313
314fn confidence_for_evidence(evidence: &str) -> f32 {
315 match evidence {
316 TRACE_EBPF_FILE => 0.95,
317 TRACE_PROC_FD => 0.90,
318 TRACE_STICKY_BINDING => 0.70,
319 TRACE_RECENT_CWD => 0.55,
320 _ => 0.50,
321 }
322}