1use serde_json::{Map, Value};
2use std::fs::{self, File, create_dir_all};
3use std::io::{self, Write};
4use std::path::{Path, PathBuf};
5use std::sync::{Mutex, MutexGuard};
6
7use crate::config::AliasProvider;
8use crate::logging::REDACT_KEYS;
9use crate::paths;
10
11#[derive(Debug)]
12pub struct TrafficCapture {
13 root: PathBuf,
14 artifact_counter: Mutex<usize>,
15 event_counter: Mutex<usize>,
16}
17
18pub const MAX_SSE_CAPTURE_BYTES: usize = 8 * 1024 * 1024;
19pub const MAX_STREAM_CAPTURE_EVENT_BYTES: usize = 8 * 1024 * 1024;
20pub const MAX_STREAM_CAPTURE_EVENTS: usize = 1_024;
21pub const MAX_STREAM_CAPTURE_FRAME_BYTES: usize = 64 * 1024;
22
23#[derive(Debug)]
24pub struct TrafficCaptureOptions {
25 pub req_id: String,
26 pub session_id: Option<String>,
27 pub session_seq: Option<u64>,
28 pub provider: Option<String>,
29 pub state_dir_override: Option<PathBuf>,
30}
31
32pub fn traffic_capture_enabled() -> bool {
33 traffic_capture_enabled_for_env(&std::env::vars().collect())
34}
35
36pub fn traffic_capture_enabled_for_env(env: &std::collections::HashMap<String, String>) -> bool {
37 match env.get("CCP_TRAFFIC_LOG").map(String::as_str) {
38 Some(v) => matches!(v, "1" | "true" | "yes"),
39 None => false,
40 }
41}
42
43pub fn create_traffic_capture(opts: TrafficCaptureOptions) -> Option<TrafficCapture> {
44 if !traffic_capture_enabled() {
45 return None;
46 }
47 let state_root = opts
48 .state_dir_override
49 .unwrap_or_else(paths::state_dir)
50 .join("traffic")
51 .join(sanitize_path_part(
52 opts.session_id.as_deref().unwrap_or("no-session"),
53 ))
54 .join(format!(
55 "{:06}-{}-{}",
56 opts.session_seq.unwrap_or(0),
57 sanitize_path_part(opts.provider.as_deref().unwrap_or("unknown-provider")),
58 sanitize_path_part(&opts.req_id),
59 ));
60
61 Some(TrafficCapture {
62 root: state_root,
63 artifact_counter: Mutex::new(0),
64 event_counter: Mutex::new(0),
65 })
66}
67
68impl TrafficCapture {
69 pub fn root(&self) -> &Path {
70 &self.root
71 }
72
73 pub fn write_json(&self, name: &str, value: &Value) {
74 let value = redact_traffic(value);
75 let payload = serde_json::to_string_pretty(&value)
76 .unwrap_or_else(|_| "{}".to_string())
77 .into_bytes();
78 let path = self.next_artifact_path(name, true);
79 let _ = write_bytes(path, &payload);
80 }
81
82 pub fn write_text(&self, name: &str, text: &str) {
83 let file = if name.ends_with(".txt") {
84 name.to_string()
85 } else {
86 format!("{name}.txt")
87 };
88 let path = self.next_artifact_path(&file, false);
89 let _ = write_bytes(path, text.as_bytes());
90 }
91
92 pub fn write_bytes(&self, name: &str, value: &[u8]) {
93 let path = self.next_artifact_path(name, false);
94 let _ = write_bytes(path, value);
95 }
96
97 pub fn write_json_event(&self, name: &str, value: &Value) {
98 let value = redact_traffic(value);
99 let payload = serde_json::to_string_pretty(&value)
100 .unwrap_or_else(|_| "{}".to_string())
101 .into_bytes();
102 let path = self.next_event_path(name, true);
103 let _ = write_bytes(path, &payload);
104 }
105
106 pub fn stream_capture(&self) -> StreamTrafficCapture {
107 StreamTrafficCapture::default()
108 }
109
110 fn next_artifact_path(&self, name: &str, ensure_ext_json: bool) -> PathBuf {
111 let mut counter: MutexGuard<'_, usize> = self
112 .artifact_counter
113 .lock()
114 .unwrap_or_else(|_| self.artifact_counter.lock().unwrap());
115 *counter += 1;
116 let file = if ensure_ext_json && !name.ends_with(".json") {
117 format!("{name}.json")
118 } else {
119 name.to_string()
120 };
121 self.root
122 .join(format!("{:03}-{}", *counter, sanitize_path_part(&file)))
123 }
124
125 fn next_event_path(&self, name: &str, ensure_ext_json: bool) -> PathBuf {
126 let mut counter: MutexGuard<'_, usize> = self
127 .event_counter
128 .lock()
129 .unwrap_or_else(|_| self.event_counter.lock().unwrap());
130 *counter += 1;
131 let file = if ensure_ext_json && !name.ends_with(".json") {
132 format!("{name}.json")
133 } else {
134 name.to_string()
135 };
136 self.root
137 .join("events")
138 .join(format!("{:06}-{}", *counter, sanitize_path_part(&file)))
139 }
140}
141
142#[cfg(test)]
143pub(crate) fn test_capture(root: PathBuf) -> TrafficCapture {
144 TrafficCapture {
145 root,
146 artifact_counter: Mutex::new(0),
147 event_counter: Mutex::new(0),
148 }
149}
150
151fn write_bytes(path: PathBuf, value: &[u8]) -> io::Result<()> {
152 if let Some(parent) = path.parent() {
153 create_dir_all(parent)?;
154 if let Ok(meta) = fs::metadata(parent) {
155 set_mode(parent, 0o700);
156 if meta.is_dir() {
157 #[cfg(unix)]
158 {
159 use std::os::unix::fs::PermissionsExt;
160 let mut perm = meta.permissions();
161 perm.set_mode(0o700);
162 let _ = fs::set_permissions(parent, perm);
163 }
164 }
165 }
166 }
167 let mut out = File::create(&path)?;
168 out.write_all(value)?;
169 #[cfg(unix)]
170 {
171 use std::os::unix::fs::PermissionsExt;
172 let mut perm = out.metadata()?.permissions();
173 perm.set_mode(0o600);
174 let _ = fs::set_permissions(&path, perm);
175 }
176 Ok(())
177}
178
179pub struct StreamTrafficCapture {
180 upstream_sse: Vec<u8>,
181 upstream_events: Vec<Value>,
182 downstream_events: Vec<Value>,
183 malformed: Vec<Value>,
184 upstream_event_bytes: usize,
185 downstream_event_bytes: usize,
186 upstream_sse_truncated: u64,
187 upstream_events_truncated: u64,
188 downstream_events_truncated: u64,
189 malformed_truncated: u64,
190 upstream_frames_truncated: u64,
191}
192
193impl Default for StreamTrafficCapture {
194 fn default() -> Self {
195 Self {
196 upstream_sse: Vec::with_capacity(MAX_SSE_CAPTURE_BYTES.min(64 * 1024)),
197 upstream_events: Vec::new(),
198 downstream_events: Vec::new(),
199 malformed: Vec::new(),
200 upstream_event_bytes: 0,
201 downstream_event_bytes: 0,
202 upstream_sse_truncated: 0,
203 upstream_events_truncated: 0,
204 downstream_events_truncated: 0,
205 malformed_truncated: 0,
206 upstream_frames_truncated: 0,
207 }
208 }
209}
210
211impl StreamTrafficCapture {
212 pub fn upstream_event(&mut self, event: Option<&str>, value: &Value) {
213 let value = redact_traffic(value);
214 let frame = serde_json::to_vec(&value).unwrap_or_default();
215 let event = event.unwrap_or("message");
216 let frame_len = event.len().saturating_add(frame.len()).saturating_add(16);
217 if frame_len > MAX_STREAM_CAPTURE_FRAME_BYTES {
218 self.upstream_frames_truncated = self.upstream_frames_truncated.saturating_add(1);
219 } else if self.upstream_sse.len().saturating_add(frame_len) <= MAX_SSE_CAPTURE_BYTES {
220 self.upstream_sse.extend_from_slice(b"event: ");
221 self.upstream_sse.extend_from_slice(event.as_bytes());
222 self.upstream_sse.extend_from_slice(b"\ndata: ");
223 self.upstream_sse.extend_from_slice(&frame);
224 self.upstream_sse.extend_from_slice(b"\n\n");
225 } else {
226 self.upstream_sse_truncated = self.upstream_sse_truncated.saturating_add(1);
227 }
228 self.push_event(true, serde_json::json!({"event":event,"data":value}));
229 }
230
231 pub fn malformed(&mut self, stage: &str, kind: &str) {
232 if self.malformed.len() < MAX_STREAM_CAPTURE_EVENTS {
233 self.malformed
234 .push(serde_json::json!({"stage":stage,"kind":kind}));
235 } else {
236 self.malformed_truncated = self.malformed_truncated.saturating_add(1);
237 }
238 }
239
240 pub fn downstream_event(&mut self, event: &str, data: Value) {
241 self.push_event(
242 false,
243 serde_json::json!({"event":event,"data":redact_traffic(&data)}),
244 );
245 }
246
247 fn push_event(&mut self, upstream: bool, value: Value) {
248 let bytes = serde_json::to_vec(&value).map_or(0, |value| value.len());
249 let (events, total, truncated) = if upstream {
250 (
251 &mut self.upstream_events,
252 &mut self.upstream_event_bytes,
253 &mut self.upstream_events_truncated,
254 )
255 } else {
256 (
257 &mut self.downstream_events,
258 &mut self.downstream_event_bytes,
259 &mut self.downstream_events_truncated,
260 )
261 };
262 if events.len() < MAX_STREAM_CAPTURE_EVENTS
263 && total.saturating_add(bytes) <= MAX_STREAM_CAPTURE_EVENT_BYTES
264 {
265 *total += bytes;
266 events.push(value);
267 } else {
268 *truncated = truncated.saturating_add(1);
269 }
270 }
271
272 pub fn finish(self, traffic: &TrafficCapture, completion: Value) {
273 let upstream_event_count = self.upstream_events.len();
274 let downstream_event_count = self.downstream_events.len();
275 if !self.upstream_sse.is_empty() {
276 traffic.write_bytes("032-upstream-response-body.sse", &self.upstream_sse);
277 }
278 traffic.write_json(
279 "033-upstream-response-capture",
280 &serde_json::json!({
281 "truncated": self.upstream_sse_truncated > 0 || self.upstream_frames_truncated > 0 || self.upstream_events_truncated > 0,
282 "captured_bytes": self.upstream_sse.len(),
283 "truncated_frames": self.upstream_sse_truncated,
284 "oversized_frames": self.upstream_frames_truncated,
285 "captured_events": self.upstream_events.len(),
286 "captured_event_bytes": self.upstream_event_bytes,
287 "truncated_events": self.upstream_events_truncated,
288 "malformed": self.malformed,
289 "truncated_malformed": self.malformed_truncated,
290 }),
291 );
292 for value in self.upstream_events {
293 traffic.write_json_event("040-upstream-event", &value);
294 }
295 for value in self.downstream_events {
296 traffic.write_json_event("050-downstream-event", &value);
297 }
298 traffic.write_json(
299 "061-grok-stream-summary",
300 &serde_json::json!({
301 "completion": completion,
302 "upstream_sse": {
303 "captured_bytes": self.upstream_sse.len(),
304 "truncated_frames": self.upstream_sse_truncated,
305 "oversized_frames": self.upstream_frames_truncated,
306 },
307 "upstream_events": {
308 "captured": upstream_event_count,
309 "captured_bytes": self.upstream_event_bytes,
310 "truncated": self.upstream_events_truncated,
311 },
312 "downstream_events": {
313 "captured": downstream_event_count,
314 "captured_bytes": self.downstream_event_bytes,
315 "truncated": self.downstream_events_truncated,
316 },
317 }),
318 );
319 }
320}
321
322pub fn sanitize_path_part(input: &str) -> String {
323 let cleaned: String = input
324 .chars()
325 .map(|ch| {
326 if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' || ch == '.' {
327 ch
328 } else {
329 '_'
330 }
331 })
332 .collect();
333
334 let truncated = if cleaned.len() > 160 {
335 &cleaned[..160]
336 } else {
337 &cleaned
338 };
339 if truncated.is_empty() {
340 "unknown".to_string()
341 } else {
342 truncated.to_string()
343 }
344}
345
346pub fn redact_traffic(value: &Value) -> Value {
347 redact_traffic_with_depth(value, 0)
348}
349
350fn redact_traffic_with_depth(value: &Value, depth: u16) -> Value {
351 if depth > 100 {
352 return Value::String("[depth-limit]".to_string());
353 }
354
355 match value {
356 Value::Object(map) => {
357 let mut out = Map::new();
358 for (key, value) in map {
359 let normalized = key.to_lowercase();
360 if REDACT_KEYS.contains(&normalized.as_str())
361 || matches!(
362 normalized.as_str(),
363 "token"
364 | "bearer_token"
365 | "oauth_token"
366 | "oauth_access_token"
367 | "oauth_refresh_token"
368 | "client_secret"
369 | "secret"
370 | "password"
371 | "email"
372 | "user_id"
373 | "account_id"
374 | "identity"
375 | "identity_id"
376 | "subject"
377 | "sub"
378 )
379 {
380 out.insert(key.clone(), redact_traffic_value(value));
381 } else {
382 out.insert(key.clone(), redact_traffic_with_depth(value, depth + 1));
383 }
384 }
385 Value::Object(out)
386 }
387 Value::Array(values) => Value::Array(
388 values
389 .iter()
390 .map(|value| redact_traffic_with_depth(value, depth + 1))
391 .collect(),
392 ),
393 _ => value.clone(),
394 }
395}
396
397fn redact_traffic_value(value: &Value) -> Value {
398 match value {
399 Value::String(s) => Value::String(format!("[redacted len={}]", s.len())),
400 Value::Object(_) | Value::Array(_) => Value::String("[redacted]".to_string()),
401 _ => Value::String("[redacted]".to_string()),
402 }
403}
404
405fn set_mode(path: &Path, mode: u32) {
406 #[cfg(unix)]
407 {
408 use std::os::unix::fs::PermissionsExt;
409 if let Ok(meta) = fs::metadata(path) {
410 let mut perm = meta.permissions();
411 perm.set_mode(mode);
412 let _ = fs::set_permissions(path, perm);
413 }
414 }
415}
416
417#[allow(dead_code)]
418fn _provider_alias(_provider: &str) -> Option<AliasProvider> {
419 None
420}