use serde_json::{Map, Value};
use std::fs::{self, File, create_dir_all};
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::{Mutex, MutexGuard};
use crate::config::AliasProvider;
use crate::logging::REDACT_KEYS;
use crate::paths;
#[derive(Debug)]
pub struct TrafficCapture {
root: PathBuf,
artifact_counter: Mutex<usize>,
event_counter: Mutex<usize>,
}
pub const MAX_SSE_CAPTURE_BYTES: usize = 8 * 1024 * 1024;
pub const MAX_STREAM_CAPTURE_EVENT_BYTES: usize = 8 * 1024 * 1024;
pub const MAX_STREAM_CAPTURE_EVENTS: usize = 1_024;
pub const MAX_STREAM_CAPTURE_FRAME_BYTES: usize = 64 * 1024;
#[derive(Debug)]
pub struct TrafficCaptureOptions {
pub req_id: String,
pub session_id: Option<String>,
pub session_seq: Option<u64>,
pub provider: Option<String>,
pub state_dir_override: Option<PathBuf>,
}
pub fn traffic_capture_enabled() -> bool {
traffic_capture_enabled_for_env(&std::env::vars().collect())
}
pub fn traffic_capture_enabled_for_env(env: &std::collections::HashMap<String, String>) -> bool {
match env.get("CCP_TRAFFIC_LOG").map(String::as_str) {
Some(v) => matches!(v, "1" | "true" | "yes"),
None => false,
}
}
pub fn create_traffic_capture(opts: TrafficCaptureOptions) -> Option<TrafficCapture> {
if !traffic_capture_enabled() {
return None;
}
let state_root = opts
.state_dir_override
.unwrap_or_else(paths::state_dir)
.join("traffic")
.join(sanitize_path_part(
opts.session_id.as_deref().unwrap_or("no-session"),
))
.join(format!(
"{:06}-{}-{}",
opts.session_seq.unwrap_or(0),
sanitize_path_part(opts.provider.as_deref().unwrap_or("unknown-provider")),
sanitize_path_part(&opts.req_id),
));
Some(TrafficCapture {
root: state_root,
artifact_counter: Mutex::new(0),
event_counter: Mutex::new(0),
})
}
impl TrafficCapture {
pub fn root(&self) -> &Path {
&self.root
}
pub fn write_json(&self, name: &str, value: &Value) {
let value = redact_traffic(value);
let payload = serde_json::to_string_pretty(&value)
.unwrap_or_else(|_| "{}".to_string())
.into_bytes();
let path = self.next_artifact_path(name, true);
let _ = write_bytes(path, &payload);
}
pub fn write_text(&self, name: &str, text: &str) {
let file = if name.ends_with(".txt") {
name.to_string()
} else {
format!("{name}.txt")
};
let path = self.next_artifact_path(&file, false);
let _ = write_bytes(path, text.as_bytes());
}
pub fn write_bytes(&self, name: &str, value: &[u8]) {
let path = self.next_artifact_path(name, false);
let _ = write_bytes(path, value);
}
pub fn write_json_event(&self, name: &str, value: &Value) {
let value = redact_traffic(value);
let payload = serde_json::to_string_pretty(&value)
.unwrap_or_else(|_| "{}".to_string())
.into_bytes();
let path = self.next_event_path(name, true);
let _ = write_bytes(path, &payload);
}
pub fn stream_capture(&self) -> StreamTrafficCapture {
StreamTrafficCapture::default()
}
fn next_artifact_path(&self, name: &str, ensure_ext_json: bool) -> PathBuf {
let mut counter: MutexGuard<'_, usize> = self
.artifact_counter
.lock()
.unwrap_or_else(|_| self.artifact_counter.lock().unwrap());
*counter += 1;
let file = if ensure_ext_json && !name.ends_with(".json") {
format!("{name}.json")
} else {
name.to_string()
};
self.root
.join(format!("{:03}-{}", *counter, sanitize_path_part(&file)))
}
fn next_event_path(&self, name: &str, ensure_ext_json: bool) -> PathBuf {
let mut counter: MutexGuard<'_, usize> = self
.event_counter
.lock()
.unwrap_or_else(|_| self.event_counter.lock().unwrap());
*counter += 1;
let file = if ensure_ext_json && !name.ends_with(".json") {
format!("{name}.json")
} else {
name.to_string()
};
self.root
.join("events")
.join(format!("{:06}-{}", *counter, sanitize_path_part(&file)))
}
}
#[cfg(test)]
pub(crate) fn test_capture(root: PathBuf) -> TrafficCapture {
TrafficCapture {
root,
artifact_counter: Mutex::new(0),
event_counter: Mutex::new(0),
}
}
fn write_bytes(path: PathBuf, value: &[u8]) -> io::Result<()> {
if let Some(parent) = path.parent() {
create_dir_all(parent)?;
if let Ok(meta) = fs::metadata(parent) {
set_mode(parent, 0o700);
if meta.is_dir() {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mut perm = meta.permissions();
perm.set_mode(0o700);
let _ = fs::set_permissions(parent, perm);
}
}
}
}
let mut out = File::create(&path)?;
out.write_all(value)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mut perm = out.metadata()?.permissions();
perm.set_mode(0o600);
let _ = fs::set_permissions(&path, perm);
}
Ok(())
}
pub struct StreamTrafficCapture {
upstream_sse: Vec<u8>,
upstream_events: Vec<Value>,
downstream_events: Vec<Value>,
malformed: Vec<Value>,
upstream_event_bytes: usize,
downstream_event_bytes: usize,
upstream_sse_truncated: u64,
upstream_events_truncated: u64,
downstream_events_truncated: u64,
malformed_truncated: u64,
upstream_frames_truncated: u64,
}
impl Default for StreamTrafficCapture {
fn default() -> Self {
Self {
upstream_sse: Vec::with_capacity(MAX_SSE_CAPTURE_BYTES.min(64 * 1024)),
upstream_events: Vec::new(),
downstream_events: Vec::new(),
malformed: Vec::new(),
upstream_event_bytes: 0,
downstream_event_bytes: 0,
upstream_sse_truncated: 0,
upstream_events_truncated: 0,
downstream_events_truncated: 0,
malformed_truncated: 0,
upstream_frames_truncated: 0,
}
}
}
impl StreamTrafficCapture {
pub fn upstream_event(&mut self, event: Option<&str>, value: &Value) {
let value = redact_traffic(value);
let frame = serde_json::to_vec(&value).unwrap_or_default();
let event = event.unwrap_or("message");
let frame_len = event.len().saturating_add(frame.len()).saturating_add(16);
if frame_len > MAX_STREAM_CAPTURE_FRAME_BYTES {
self.upstream_frames_truncated = self.upstream_frames_truncated.saturating_add(1);
} else if self.upstream_sse.len().saturating_add(frame_len) <= MAX_SSE_CAPTURE_BYTES {
self.upstream_sse.extend_from_slice(b"event: ");
self.upstream_sse.extend_from_slice(event.as_bytes());
self.upstream_sse.extend_from_slice(b"\ndata: ");
self.upstream_sse.extend_from_slice(&frame);
self.upstream_sse.extend_from_slice(b"\n\n");
} else {
self.upstream_sse_truncated = self.upstream_sse_truncated.saturating_add(1);
}
self.push_event(true, serde_json::json!({"event":event,"data":value}));
}
pub fn malformed(&mut self, stage: &str, kind: &str) {
if self.malformed.len() < MAX_STREAM_CAPTURE_EVENTS {
self.malformed
.push(serde_json::json!({"stage":stage,"kind":kind}));
} else {
self.malformed_truncated = self.malformed_truncated.saturating_add(1);
}
}
pub fn downstream_event(&mut self, event: &str, data: Value) {
self.push_event(
false,
serde_json::json!({"event":event,"data":redact_traffic(&data)}),
);
}
fn push_event(&mut self, upstream: bool, value: Value) {
let bytes = serde_json::to_vec(&value).map_or(0, |value| value.len());
let (events, total, truncated) = if upstream {
(
&mut self.upstream_events,
&mut self.upstream_event_bytes,
&mut self.upstream_events_truncated,
)
} else {
(
&mut self.downstream_events,
&mut self.downstream_event_bytes,
&mut self.downstream_events_truncated,
)
};
if events.len() < MAX_STREAM_CAPTURE_EVENTS
&& total.saturating_add(bytes) <= MAX_STREAM_CAPTURE_EVENT_BYTES
{
*total += bytes;
events.push(value);
} else {
*truncated = truncated.saturating_add(1);
}
}
pub fn finish(self, traffic: &TrafficCapture, completion: Value) {
let upstream_event_count = self.upstream_events.len();
let downstream_event_count = self.downstream_events.len();
if !self.upstream_sse.is_empty() {
traffic.write_bytes("032-upstream-response-body.sse", &self.upstream_sse);
}
traffic.write_json(
"033-upstream-response-capture",
&serde_json::json!({
"truncated": self.upstream_sse_truncated > 0 || self.upstream_frames_truncated > 0 || self.upstream_events_truncated > 0,
"captured_bytes": self.upstream_sse.len(),
"truncated_frames": self.upstream_sse_truncated,
"oversized_frames": self.upstream_frames_truncated,
"captured_events": self.upstream_events.len(),
"captured_event_bytes": self.upstream_event_bytes,
"truncated_events": self.upstream_events_truncated,
"malformed": self.malformed,
"truncated_malformed": self.malformed_truncated,
}),
);
for value in self.upstream_events {
traffic.write_json_event("040-upstream-event", &value);
}
for value in self.downstream_events {
traffic.write_json_event("050-downstream-event", &value);
}
traffic.write_json(
"061-grok-stream-summary",
&serde_json::json!({
"completion": completion,
"upstream_sse": {
"captured_bytes": self.upstream_sse.len(),
"truncated_frames": self.upstream_sse_truncated,
"oversized_frames": self.upstream_frames_truncated,
},
"upstream_events": {
"captured": upstream_event_count,
"captured_bytes": self.upstream_event_bytes,
"truncated": self.upstream_events_truncated,
},
"downstream_events": {
"captured": downstream_event_count,
"captured_bytes": self.downstream_event_bytes,
"truncated": self.downstream_events_truncated,
},
}),
);
}
}
pub fn sanitize_path_part(input: &str) -> String {
let cleaned: String = input
.chars()
.map(|ch| {
if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' || ch == '.' {
ch
} else {
'_'
}
})
.collect();
let truncated = if cleaned.len() > 160 {
&cleaned[..160]
} else {
&cleaned
};
if truncated.is_empty() {
"unknown".to_string()
} else {
truncated.to_string()
}
}
pub fn redact_traffic(value: &Value) -> Value {
redact_traffic_with_depth(value, 0)
}
fn redact_traffic_with_depth(value: &Value, depth: u16) -> Value {
if depth > 100 {
return Value::String("[depth-limit]".to_string());
}
match value {
Value::Object(map) => {
let mut out = Map::new();
for (key, value) in map {
let normalized = key.to_lowercase();
if REDACT_KEYS.contains(&normalized.as_str())
|| matches!(
normalized.as_str(),
"token"
| "bearer_token"
| "oauth_token"
| "oauth_access_token"
| "oauth_refresh_token"
| "client_secret"
| "secret"
| "password"
| "email"
| "user_id"
| "account_id"
| "identity"
| "identity_id"
| "subject"
| "sub"
)
{
out.insert(key.clone(), redact_traffic_value(value));
} else {
out.insert(key.clone(), redact_traffic_with_depth(value, depth + 1));
}
}
Value::Object(out)
}
Value::Array(values) => Value::Array(
values
.iter()
.map(|value| redact_traffic_with_depth(value, depth + 1))
.collect(),
),
_ => value.clone(),
}
}
fn redact_traffic_value(value: &Value) -> Value {
match value {
Value::String(s) => Value::String(format!("[redacted len={}]", s.len())),
Value::Object(_) | Value::Array(_) => Value::String("[redacted]".to_string()),
_ => Value::String("[redacted]".to_string()),
}
}
fn set_mode(path: &Path, mode: u32) {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
if let Ok(meta) = fs::metadata(path) {
let mut perm = meta.permissions();
perm.set_mode(mode);
let _ = fs::set_permissions(path, perm);
}
}
}
#[allow(dead_code)]
fn _provider_alias(_provider: &str) -> Option<AliasProvider> {
None
}