From ce9849bff066e7cabe70937d42823bbc0e441f9c Mon Sep 17 00:00:00 2001 From: parth2024-tech Date: Sun, 8 Nov 2026 17:30:15 +0530 Subject: [PATCH 1/3] feat(security,protocol): harden control plane, resolve tool loading, and fix compaction --- .ase-project.json | 94 ++++++++ Cargo.lock | 2 + Cargo.toml | 2 + src/auth.rs | 46 ++++ src/config.rs | 68 ++++++ src/providers/codex/events.rs | 20 ++ src/providers/codex/mod.rs | 79 +++++-- src/providers/codex/translate/request.rs | 10 + src/providers/codex/websocket.rs | 2 +- src/providers/grok/translate/request.rs | 7 +- src/providers/kimi/translate/request.rs | 15 +- src/providers/opencode/chat.rs | 12 +- src/providers/translate_shared.rs | 84 ++++++++ src/request_identity.rs | 7 + src/server.rs | 260 ++++++++++++++++++++++- 15 files changed, 680 insertions(+), 28 deletions(-) create mode 100644 .ase-project.json diff --git a/.ase-project.json b/.ase-project.json new file mode 100644 index 00000000..3f0c8c46 --- /dev/null +++ b/.ase-project.json @@ -0,0 +1,94 @@ +{ + "name": "claude-code-proxy", + "version": "0.1.48", + "description": "Anthropic-compatible proxy bridging Claude Code to Codex, Kimi, Grok, OpenCode Go, and Cursor Agent backends.", + "validation_timeout_ms": 300000, + "tech_stack": [ + "Rust 2024", + "axum 0.8", + "tokio (rt-multi-thread)", + "reqwest", + "prost", + "ratatui", + "tokio-tungstenite" + ], + "conventions": [ + "Strict type safety and error-fidelity matching Anthropic Messages API specs", + "No blocking operations on Tokio async runtime worker threads", + "Security-first: constant-time token comparison, Origin/Host validation, atomic file permissions", + "On-demand daemon architecture with clean shutdown" + ], + "architecture_patterns": [ + "Provider abstraction with translation layers and streaming reducers", + "Shared session management with composite lineage-aware state keys", + "SSE state machine and WebSocket protocol bridging" + ], + "file_structure": { + "src/app": "Next.js App Router pages and layouts", + "src/features": "Feature modules with co-located components, hooks, and types", + "src/lib": "Shared utilities, database client, auth helpers", + "src/components/ui": "Reusable UI primitives (Button, Modal, Input)", + "tests": "Test files mirroring src/ structure" + }, + "testing_strategy": { + "framework": "Vitest + React Testing Library", + "coverage_target": "80%", + "patterns": [ + "Unit tests for all utility functions", + "Integration tests for API routes", + "Component tests for interactive UI elements" + ] + }, + "deployment": { + "platform": "Vercel", + "branch_strategy": "main → production, dev → preview", + "environments": [ + "development", + "staging", + "production" + ] + }, + "sprint_focus": "Building the user onboarding flow with email verification and profile setup.", + "sprint_skills": [ + "security-and-hardening", + "api-and-interface-design", + "test-driven-development" + ], + "never_do": [ + "Never use useEffect for data fetching — use Server Components or SWR", + "Never store secrets in .env files committed to git", + "Never use inline styles — always use Tailwind classes", + "Never use `any` TypeScript type — always type explicitly", + "Never mix database queries into React components", + "Never output verbose unsolicited recaps, exhaustive changelogs, or repetitive summaries that waste tokens — mention ONLY the most critical points unless explicitly asked", + "Never use purple gradients, pill buttons, fake reviews/metrics/counters, vague hero text, emoji icons, em dashes, crazy scroll/cursor animations, or AI slop photos/copy on websites", + "Never let an AI grade its own work without independent gate verification", + "Never leave completed work uncommitted or unpushed to GitHub", + "Never run unsolicited background daemons, update pollers, or indexing services on Arch Linux" + ], + "preferred_patterns": [ + "Practice extreme token economy: keep responses concise, high-signal, and directly answer user requests without unprompted changelogs or filler", + "Use Zod for all schema validation", + "Use Drizzle ORM for all database access", + "Use server actions for mutations", + "Use Next.js Image component for all images", + "Always mandate production web essentials: custom domain, custom favicon suite, remove AI badges, /privacy page, and /terms page", + "Always follow the 6-step Second Brain closed loop: Load Conventions -> Select Model -> Autonomous Loop -> Validate Output -> Continuous Learning -> Automated Git Push", + "Always enforce Arch Linux On-Demand System Governance (run interactive tools on-demand, terminate completely on exit)" + ], + "agent_skills": { + "directory": "/home/thor/skiils/skills", + "default_skills": [ + "using-agent-skills", + "context-engineering", + "anti-vibe-design" + ], + "mandatory_skills": { + "api": "api-and-interface-design", + "auth": "security-and-hardening", + "frontend": "anti-vibe-design", + "ui": "anti-vibe-design" + }, + "session_reporting_preference": true + } +} \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 7e1e9ce9..a635ce50 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -351,6 +351,7 @@ dependencies = [ "http-body-util", "hyper-util", "jiff", + "libc", "once_cell", "predicates", "prost", @@ -363,6 +364,7 @@ dependencies = [ "serde", "serde_json", "sha2", + "subtle", "tempfile", "thiserror 2.0.18", "tiktoken-rs", diff --git a/Cargo.toml b/Cargo.toml index 267820da..be27a002 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,8 @@ sha2 = "0.10" hostname = "0.4" once_cell = "1.20" regex-lite = "0.1" +subtle = "2.6" +libc = "0.2" tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } tokio-rustls = { version = "0.26", default-features = false, features = ["ring", "tls12"] } rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } diff --git a/src/auth.rs b/src/auth.rs index 279c2946..9700efea 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -328,6 +328,49 @@ pub fn delete_auth_file(primary: &std::path::Path, legacy: &std::path::Path) -> Ok(()) } +pub struct FileLock { + #[cfg(unix)] + fd: std::os::unix::io::RawFd, + #[cfg(unix)] + _file: File, +} + +impl FileLock { + pub fn acquire(lock_path: &std::path::Path) -> io::Result { + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + use std::os::unix::io::AsRawFd; + let file = std::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .mode(0o600) + .open(lock_path)?; + let fd = file.as_raw_fd(); + let ret = unsafe { libc::flock(fd, libc::LOCK_EX) }; + if ret != 0 { + return Err(io::Error::last_os_error()); + } + Ok(Self { fd, _file: file }) + } + #[cfg(not(unix))] + { + Ok(Self {}) + } + } +} + +#[cfg(unix)] +impl Drop for FileLock { + fn drop(&mut self) { + unsafe { + libc::flock(self.fd, libc::LOCK_UN); + } + } +} + pub fn write_atomically(path: &str, value: &T) -> Result<()> { let dir = std::path::Path::new(path) .parent() @@ -335,6 +378,9 @@ pub fn write_atomically(path: &str, value: &T) -> Result<()> { fs::create_dir_all(dir)?; set_mode(dir, 0o700); + let lock_path = format!("{path}.lock"); + let _lock = FileLock::acquire(std::path::Path::new(&lock_path))?; + let tmp = format!("{path}.tmp-{}", uuid::Uuid::new_v4()); #[cfg(unix)] let mut out = { diff --git a/src/config.rs b/src/config.rs index ee1ff80a..e8084ef5 100644 --- a/src/config.rs +++ b/src/config.rs @@ -39,6 +39,10 @@ struct FileConfig { pub alias_provider: Option, #[serde(rename = "autoReviewModel")] pub auto_review_model: Option, + #[serde(rename = "authToken")] + pub auth_token: Option, + #[serde(rename = "requireAuth")] + pub require_auth: Option, pub log: Option, pub kimi: Option, pub codex: Option, @@ -244,6 +248,70 @@ pub fn log_stderr() -> bool { load_config().log_stderr } +pub fn local_auth_token() -> Option { + let env: HashMap<_, _> = std::env::vars().collect(); + if let Some(raw) = env.get("CCP_LOCAL_AUTH_TOKEN").or_else(|| env.get("CCP_AUTH_TOKEN")) { + let trimmed = raw.trim(); + if !trimmed.is_empty() { + return Some(trimmed.to_string()); + } + } + let config_dir = paths::config_dir(); + if let Some(file) = read_file_config(&config_dir) + && let Some(token) = file.auth_token + { + let trimmed = token.trim(); + if !trimmed.is_empty() { + return Some(trimmed.to_string()); + } + } + let state_file = paths::state_dir().join("local_auth_token"); + if let Ok(token) = fs::read_to_string(state_file) { + let trimmed = token.trim(); + if !trimmed.is_empty() { + return Some(trimmed.to_string()); + } + } + None +} + +pub fn require_local_auth() -> bool { + let env: HashMap<_, _> = std::env::vars().collect(); + if let Some(raw) = env.get("CCP_REQUIRE_AUTH") { + return matches!(raw.to_ascii_lowercase().as_str(), "1" | "true" | "yes"); + } + let config_dir = paths::config_dir(); + if let Some(file) = read_file_config(&config_dir) + && let Some(enabled) = file.require_auth + { + return enabled; + } + false +} + +pub fn ensure_local_auth_token() -> Option { + if let Some(token) = local_auth_token() { + return Some(token); + } + if !require_local_auth() { + return None; + } + let state_dir = paths::state_dir(); + let _ = fs::create_dir_all(&state_dir); + let token = format!("ccp_{}", uuid::Uuid::new_v4().to_string().replace('-', "")); + let state_file = state_dir.join("local_auth_token"); + if let Ok(()) = fs::write(&state_file, &token) { + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let _ = fs::set_permissions(&state_file, fs::Permissions::from_mode(0o600)); + } + Some(token) + } else { + None + } +} + pub fn config_override_summary_lines(cfg: &LoadedConfig) -> Vec { let file = read_file_config(&cfg.config_dir); let env: HashMap<_, _> = std::env::vars().collect(); diff --git a/src/providers/codex/events.rs b/src/providers/codex/events.rs index 1c63740a..7d3f01e4 100644 --- a/src/providers/codex/events.rs +++ b/src/providers/codex/events.rs @@ -298,6 +298,26 @@ pub struct CodexUsageLimit { /// (`resets_at`, `resets_in_seconds`) and in `X-Codex-*` headers mirrored into /// the event payload. pub(crate) fn usage_limit_from_event(payload: &Value) -> Option { + if payload.get("type").and_then(Value::as_str) == Some("codex.rate_limits") + && let Some(limits) = payload.get("rate_limits").and_then(Value::as_object) { + let reached = limits.get("limit_reached").and_then(Value::as_bool).unwrap_or(false) + || limits.get("allowed").and_then(Value::as_bool).map(|a| !a).unwrap_or(false); + if reached { + let primary = limits.get("primary").and_then(Value::as_object); + let resets_at = numeric_value(primary.and_then(|p| p.get("reset_at").or_else(|| p.get("resets_at")))); + let window_minutes = numeric_value(primary.and_then(|p| p.get("window_minutes"))); + let window = match window_minutes { + Some(300) => Some(CodexLimitWindow::FiveHour), + Some(10_080) => Some(CodexLimitWindow::SevenDay), + _ => None, + }; + return Some(CodexUsageLimit { + message: "Usage limit reached".to_string(), + resets_at, + window, + }); + } + } if !matches!( payload.get("type").and_then(Value::as_str), Some("response.failed" | "response.error" | "error") diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index 4e0f9d12..3ef4b20c 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -220,12 +220,20 @@ impl CodexProvider { let compact_boundary = is_compact_messages_request(&body); let server_compaction_enabled = config::codex_server_compaction(); let mut compaction_attempt = None; - if !server_compaction_enabled && let Some(session_id) = ctx.session_id.as_deref() { + let compaction_session_id = if body.bypass_provider_model_override { + None + } else { + conversation_identity + .as_ref() + .map(ConversationIdentity::storage_key) + .or_else(|| ctx.session_id.clone()) + }; + if !server_compaction_enabled && let Some(session_id) = compaction_session_id.as_deref() { compaction::clear_compaction(session_id); } if server_compaction_enabled && compact_boundary - && let Some(session_id) = ctx.session_id.as_deref() + && let Some(session_id) = compaction_session_id.as_deref() { let attempt = begin_compaction(session_id, &translated.model); compaction_attempt = Some(attempt); @@ -280,7 +288,7 @@ impl CodexProvider { } } else if server_compaction_enabled && !compact_boundary - && let Some(replay) = apply_compaction_replay(ctx.session_id.as_deref(), &translated) + && let Some(replay) = apply_compaction_replay(compaction_session_id.as_deref(), &translated) { translated = replay.request; compaction_attempt = Some(replay.attempt); @@ -343,6 +351,7 @@ impl CodexProvider { LiveStreamCompaction { compact_boundary, attempt: compaction_attempt, + session_id: compaction_session_id, }, configured_transport, tail_is_successful_tool_result(&body), @@ -756,10 +765,11 @@ enum LiveStreamStart { }, } -#[derive(Clone, Copy)] +#[derive(Clone, Default)] struct LiveStreamCompaction { compact_boundary: bool, attempt: Option, + session_id: Option, } #[allow(clippy::too_many_arguments)] @@ -848,7 +858,7 @@ async fn live_stream_response( ctx.clone(), request_continuation.clone(), request_body.clone(), - compaction, + compaction.clone(), transport, accept_empty_completion(successful_tool_tail, empty_completion_attempts), ) @@ -961,22 +971,34 @@ async fn live_stream_response_once( // Every branch that consumes pending_chunk returns, so it is never flushed twice. let mut pending_chunk = Vec::new(); let mut generation_started = false; + let mut last_usage_limit = None; while let Some(item) = upstream_events.recv().await { let payload = match item { Ok(payload) => payload, Err(err) => { + if let Some(limit) = last_usage_limit.as_ref() { + abort_request_state( + compaction.session_id.as_deref().or(ctx.session_id.as_deref()), + &request_continuation, + compaction.attempt, + ); + return LiveStreamStart::Response(usage_limit_response(limit)); + } if retryable_live_start_codex_error(&err) { return provider_retry(&upstream_events, err); } abort_request_state( - ctx.session_id.as_deref(), + compaction.session_id.as_deref().or(ctx.session_id.as_deref()), &request_continuation, compaction.attempt, ); return LiveStreamStart::Response(map_codex_error_to_response(&err)); } }; + if let Some(limit) = events::usage_limit_from_event(&payload) { + last_usage_limit = Some(limit); + } if !generation_started && codex_generation_event(&payload) { if let Some(monitor) = ctx.monitor.as_ref() { monitor.generation_started(&ctx.req_id); @@ -993,7 +1015,7 @@ async fn live_stream_response_once( // Report it once, with the reset clock upstream supplied. if let Some(limit) = events::usage_limit_from_event(&payload) { abort_request_state( - ctx.session_id.as_deref(), + compaction.session_id.as_deref().or(ctx.session_id.as_deref()), &request_continuation, compaction.attempt, ); @@ -1010,7 +1032,7 @@ async fn live_stream_response_once( ); } abort_request_state( - ctx.session_id.as_deref(), + compaction.session_id.as_deref().or(ctx.session_id.as_deref()), &request_continuation, compaction.attempt, ); @@ -1019,7 +1041,7 @@ async fn live_stream_response_once( )); } abort_request_state( - ctx.session_id.as_deref(), + compaction.session_id.as_deref().or(ctx.session_id.as_deref()), &request_continuation, compaction.attempt, ); @@ -1051,7 +1073,7 @@ async fn live_stream_response_once( record_live_stream_progress(&ctx, &pending_chunk); if terminal { update_continuation_from_upstream( - ctx.session_id.as_deref(), + compaction.session_id.as_deref().or(ctx.session_id.as_deref()), &request_continuation, compaction.attempt, &request_body, @@ -1074,7 +1096,7 @@ async fn live_stream_response_once( } if terminal { update_continuation_from_upstream( - ctx.session_id.as_deref(), + compaction.session_id.as_deref().or(ctx.session_id.as_deref()), &request_continuation, compaction.attempt, &request_body, @@ -1091,6 +1113,15 @@ async fn live_stream_response_once( } } + if let Some(limit) = last_usage_limit.as_ref() { + abort_request_state( + compaction.session_id.as_deref().or(ctx.session_id.as_deref()), + &request_continuation, + compaction.attempt, + ); + return LiveStreamStart::Response(usage_limit_response(limit)); + } + provider_retry( &upstream_events, client::CodexError { @@ -1267,9 +1298,10 @@ fn remaining_live_stream_response( ) -> Response { let (tx, rx) = tokio::sync::mpsc::channel::>(64); tokio::spawn(async move { + let session_ref = compaction.session_id.as_deref().or(ctx.session_id.as_deref()); if tx.send(Ok(Bytes::from(first_chunk))).await.is_err() { abort_request_state( - ctx.session_id.as_deref(), + session_ref, &request_continuation, compaction.attempt, ); @@ -1283,7 +1315,7 @@ fn remaining_live_stream_response( biased; _ = tx.closed() => { abort_request_state( - ctx.session_id.as_deref(), + session_ref, &request_continuation, compaction.attempt, ); @@ -1296,7 +1328,7 @@ fn remaining_live_stream_response( record_live_stream_progress(&ctx, &chunk); if tx.send(Ok(Bytes::from(chunk))).await.is_err() { abort_request_state( - ctx.session_id.as_deref(), + session_ref, &request_continuation, compaction.attempt, ); @@ -1320,7 +1352,7 @@ fn remaining_live_stream_response( Ok(result) => result, Err(message) => { abort_request_state( - ctx.session_id.as_deref(), + session_ref, &request_continuation, compaction.attempt, ); @@ -1347,7 +1379,7 @@ fn remaining_live_stream_response( record_live_stream_progress(&ctx, &chunk); if tx.send(Ok(Bytes::from(chunk))).await.is_err() { abort_request_state( - ctx.session_id.as_deref(), + session_ref, &request_continuation, compaction.attempt, ); @@ -1356,7 +1388,7 @@ fn remaining_live_stream_response( } if terminal { update_continuation_from_upstream( - ctx.session_id.as_deref(), + session_ref, &request_continuation, compaction.attempt, &request_body, @@ -1369,7 +1401,7 @@ fn remaining_live_stream_response( } Err(err) => { abort_request_state( - ctx.session_id.as_deref(), + session_ref, &request_continuation, compaction.attempt, ); @@ -1398,7 +1430,7 @@ fn remaining_live_stream_response( } abort_request_state( - ctx.session_id.as_deref(), + session_ref, &request_continuation, compaction.attempt, ); @@ -2221,6 +2253,7 @@ mod tests { LiveStreamCompaction { compact_boundary: false, attempt: None, + session_id: None, }, config::CodexTransport::WebSocket, true, @@ -2261,6 +2294,7 @@ mod tests { LiveStreamCompaction { compact_boundary: false, attempt: None, + session_id: None, }, config::CodexTransport::WebSocket, true, @@ -2345,6 +2379,7 @@ mod tests { LiveStreamCompaction { compact_boundary: false, attempt: None, + session_id: None, }, config::CodexTransport::WebSocket, false, @@ -2381,6 +2416,7 @@ mod tests { LiveStreamCompaction { compact_boundary: false, attempt: None, + session_id: None, }, config::CodexTransport::WebSocket, true, @@ -2631,6 +2667,7 @@ mod tests { LiveStreamCompaction { compact_boundary: false, attempt: None, + session_id: None, }, config::CodexTransport::Http, false, @@ -2942,6 +2979,7 @@ mod tests { LiveStreamCompaction { compact_boundary: false, attempt: Some(compaction_attempt), + session_id: Some(session_id.to_string()), }, config::CodexTransport::WebSocket, false, @@ -3058,6 +3096,7 @@ mod tests { LiveStreamCompaction { compact_boundary: false, attempt: Some(compaction_attempt), + session_id: Some(session_id.to_string()), }, config::CodexTransport::WebSocket, false, @@ -3144,6 +3183,7 @@ mod tests { LiveStreamCompaction { compact_boundary: false, attempt: Some(compaction_attempt), + session_id: Some(session_id.to_string()), }, config::CodexTransport::WebSocket, false, @@ -3331,6 +3371,7 @@ mod tests { LiveStreamCompaction { compact_boundary: false, attempt: Some(compaction_attempt), + session_id: Some(session_id.to_string()), }, config::CodexTransport::WebSocket, false, diff --git a/src/providers/codex/translate/request.rs b/src/providers/codex/translate/request.rs index fea9a779..1148a948 100644 --- a/src/providers/codex/translate/request.rs +++ b/src/providers/codex/translate/request.rs @@ -658,6 +658,9 @@ fn read_tools(req: &MessagesRequest) -> Result>, anyho }; let mut out = Vec::new(); for tool in tools_arr { + if crate::providers::translate_shared::should_skip_tool(tool, req) { + continue; + } let tool_type = tool .get("type") .and_then(|v| v.as_str()) @@ -1122,6 +1125,13 @@ fn render_tool_result_block(block: &Value) -> RenderedToolResultPart { RenderedToolResultPart::Text(unsupported_tool_result_block_to_string(block)) }), Some("image") => render_tool_result_image(block), + Some("tool_reference") => { + let tool_name = block + .get("tool_name") + .and_then(Value::as_str) + .unwrap_or("unknown"); + RenderedToolResultPart::Text(format!("[tool reference: {tool_name}]")) + } Some(other) => { RenderedToolResultPart::Text(format!("[unsupported content block omitted: {other}]")) } diff --git a/src/providers/codex/websocket.rs b/src/providers/codex/websocket.rs index 5fe03352..3e369410 100644 --- a/src/providers/codex/websocket.rs +++ b/src/providers/codex/websocket.rs @@ -339,7 +339,7 @@ static WS_CONNECT_GATE: once_cell::sync::Lazy = fn next_monotonic_nonzero(sequence: &AtomicU64, label: &str) -> u64 { let previous = sequence - .try_update(Ordering::Relaxed, Ordering::Relaxed, |value| { + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| { value.checked_add(1) }) .unwrap_or_else(|_| panic!("{label} sequence exhausted")); diff --git a/src/providers/grok/translate/request.rs b/src/providers/grok/translate/request.rs index ecfbaabf..452ba4a7 100644 --- a/src/providers/grok/translate/request.rs +++ b/src/providers/grok/translate/request.rs @@ -512,6 +512,7 @@ fn parse_tools( "allowed_domains", "blocked_domains", "user_location", + "defer_loading", ] .contains(&key.as_str()) { @@ -532,6 +533,9 @@ fn parse_tools( .and_then(Value::as_str) .filter(|s| !s.is_empty()) .ok_or_else(|| anyhow::anyhow!("tool name is invalid"))?; + if name == "DeferredToolPlaceholder" || obj.get("defer_loading").and_then(Value::as_bool).unwrap_or(false) { + continue; + } if !names.insert(name.to_string()) { anyhow::bail!("duplicate tool name"); } @@ -574,11 +578,12 @@ fn parse_tools( out.push(GrokTool::hosted_named("x_search", name)); continue; } - let parameters = obj + let mut parameters = obj .get("input_schema") .filter(|value| value.is_object()) .cloned() .ok_or_else(|| anyhow::anyhow!("tool input_schema must be an object"))?; + crate::providers::translate_shared::strip_tool_schema_patterns(&mut parameters); out.push(GrokTool::function( name, obj.get("description") diff --git a/src/providers/kimi/translate/request.rs b/src/providers/kimi/translate/request.rs index 3beefcff..2343e22a 100644 --- a/src/providers/kimi/translate/request.rs +++ b/src/providers/kimi/translate/request.rs @@ -247,6 +247,9 @@ fn read_tools(req: &MessagesRequest) -> Result, anyhow::Error> { }; let mut out = Vec::new(); for tool in tools_arr { + if crate::providers::translate_shared::should_skip_tool(tool, req) { + continue; + } let name = tool .get("name") .and_then(|v| v.as_str()) @@ -256,10 +259,11 @@ fn read_tools(req: &MessagesRequest) -> Result, anyhow::Error> { .get("description") .and_then(|v| v.as_str()) .map(|s| s.to_string()); - let parameters = tool + let mut parameters = tool .get("input_schema") .cloned() .unwrap_or(serde_json::json!({})); + crate::providers::translate_shared::strip_tool_schema_patterns(&mut parameters); out.push(KimiTool { kind: "function".to_string(), function: KimiToolFunction { @@ -529,6 +533,15 @@ fn tool_result_content(content: &Value, is_error: Option) -> Value { image_url: KimiImageUrl { url }, }); } + Some("tool_reference") => { + let tool_name = b + .get("tool_name") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + parts.push(KimiToolResultPart::Text { + text: format!("[tool reference: {tool_name}]"), + }); + } Some(other) => { parts.push(KimiToolResultPart::Text { text: format!("[unsupported content block omitted: {other}]"), diff --git a/src/providers/opencode/chat.rs b/src/providers/opencode/chat.rs index e120c4da..4bcdfcab 100644 --- a/src/providers/opencode/chat.rs +++ b/src/providers/opencode/chat.rs @@ -168,7 +168,9 @@ fn push_user_messages(messages: &mut Vec, blocks: &[ContentBlock]) { })); content.extend(rendered.images); } - ContentBlock::Thinking { .. } | ContentBlock::ToolUse { .. } => {} + ContentBlock::Thinking { .. } + | ContentBlock::ToolUse { .. } + | ContentBlock::ToolReference { .. } => {} } } // A parallel assistant tool-call turn requires every matching tool @@ -208,7 +210,9 @@ fn assistant_message(blocks: &[ContentBlock], deepseek: bool) -> anyhow::Result< }, })); } - ContentBlock::Image { .. } | ContentBlock::ToolResult { .. } => {} + ContentBlock::Image { .. } + | ContentBlock::ToolResult { .. } + | ContentBlock::ToolReference { .. } => {} } } if text.is_empty() && reasoning.is_empty() && tool_calls.is_empty() { @@ -290,6 +294,7 @@ fn read_tools(req: &MessagesRequest) -> anyhow::Result> { .ok_or_else(|| anyhow::anyhow!("tools must be an array"))?; tools .iter() + .filter(|tool| !crate::providers::translate_shared::should_skip_tool(tool, req)) .map(|tool| { let name = tool .get("name") @@ -297,10 +302,11 @@ fn read_tools(req: &MessagesRequest) -> anyhow::Result> { .filter(|name| !name.is_empty()) .ok_or_else(|| anyhow::anyhow!("tool name must be a non-empty string"))?; let description = tool.get("description").cloned(); - let parameters = tool + let mut parameters = tool .get("input_schema") .cloned() .unwrap_or_else(|| json!({})); + crate::providers::translate_shared::strip_tool_schema_patterns(&mut parameters); let mut function = serde_json::Map::new(); function.insert("name".into(), Value::String(name.to_string())); if let Some(description) = description.filter(|value| !value.is_null()) { diff --git a/src/providers/translate_shared.rs b/src/providers/translate_shared.rs index e98d8718..ef3c2ae8 100644 --- a/src/providers/translate_shared.rs +++ b/src/providers/translate_shared.rs @@ -24,6 +24,9 @@ pub enum ContentBlock { thinking: String, signature: Option, }, + ToolReference { + tool_name: String, + }, } #[derive(Debug)] @@ -227,6 +230,87 @@ fn parse_content_block(value: &Value, missing_tool_input: Value) -> Option { + let tool_name = value + .get("tool_name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + Some(ContentBlock::ToolReference { tool_name }) + } _ => None, } } + +pub fn is_tool_resolved(name: &str, req: &MessagesRequest) -> bool { + for msg in &req.messages { + if let Some(arr) = msg.content.as_array() { + for block in arr { + if block.get("type").and_then(Value::as_str) == Some("tool_result") + && let Some(content_arr) = block.get("content").and_then(Value::as_array) { + for sub in content_arr { + if sub.get("type").and_then(Value::as_str) == Some("tool_reference") + && sub.get("tool_name").and_then(Value::as_str) == Some(name) + { + return true; + } + } + } + } + } + } + false +} + +pub fn should_skip_tool(tool_val: &Value, req: &MessagesRequest) -> bool { + let name = tool_val.get("name").and_then(Value::as_str).unwrap_or(""); + if name == "DeferredToolPlaceholder" { + return true; + } + let is_deferred = tool_val.get("defer_loading").and_then(Value::as_bool).unwrap_or(false); + if is_deferred && !is_tool_resolved(name, req) { + return true; + } + false +} + +pub fn strip_tool_schema_patterns(schema: &mut Value) { + let Some(schema) = schema.as_object_mut() else { + return; + }; + schema.remove("pattern"); + + for (keyword, value) in schema { + match keyword.as_str() { + "properties" | "patternProperties" | "$defs" | "definitions" | "dependentSchemas" + | "dependencies" => { + if let Some(schemas) = value.as_object_mut() { + for s in schemas.values_mut() { + strip_tool_schema_patterns(s); + } + } + } + "allOf" | "anyOf" | "oneOf" | "prefixItems" | "items" => { + if let Some(schemas) = value.as_array_mut() { + for s in schemas { + strip_tool_schema_patterns(s); + } + } else if keyword == "items" { + strip_tool_schema_patterns(value); + } + } + "additionalProperties" + | "additionalItems" + | "unevaluatedProperties" + | "unevaluatedItems" + | "contains" + | "propertyNames" + | "not" + | "if" + | "then" + | "else" + | "contentSchema" => strip_tool_schema_patterns(value), + _ => {} + } + } +} diff --git a/src/request_identity.rs b/src/request_identity.rs index 0df0a13f..a8a30c96 100644 --- a/src/request_identity.rs +++ b/src/request_identity.rs @@ -30,6 +30,13 @@ impl ConversationIdentity { _ => None, } } + + pub fn storage_key(&self) -> String { + match self { + Self::Main(session_id) => session_id.clone(), + Self::Agent(session_id, agent_id) => format!("{session_id}:agent:{agent_id}"), + } + } } #[derive(Debug)] diff --git a/src/server.rs b/src/server.rs index a12db0c8..4b21992a 100644 --- a/src/server.rs +++ b/src/server.rs @@ -34,9 +34,11 @@ use axum::{ body::Body, extract::{DefaultBodyLimit, FromRequest, Multipart, Query, State}, http::{Request, StatusCode}, + middleware, response::{IntoResponse, Response}, routing::{get, post}, }; +use subtle::ConstantTimeEq; use http_body_util::{BodyExt, StreamBody}; use serde::de::DeserializeOwned; use serde_json::{Map, Value, json}; @@ -104,6 +106,7 @@ fn apply_auto_review_model( override_model: override_model.to_string(), }; body.model = Some(route.override_model.clone()); + body.extra.insert("output_config".to_string(), json!({"effort": "low"})); Some(route) } @@ -292,9 +295,155 @@ pub fn app_with_features( } else { router }; + let router = router.layer(middleware::from_fn(security_middleware)); router.fallback(fallback_handler).with_state(state) } +fn is_allowed_host(host_header: &str, bind_address: &str) -> bool { + let host = host_header.split(':').next().unwrap_or("").trim(); + if host.is_empty() { + return false; + } + if host.eq_ignore_ascii_case("localhost") { + return true; + } + if let Ok(ip) = host.parse::() + && ip.is_loopback() { + return true; + } + if host.eq_ignore_ascii_case(bind_address) { + return true; + } + if bind_address == "0.0.0.0" || bind_address == "::" { + return true; + } + false +} + +fn is_allowed_origin(origin_header: &str, bind_address: &str) -> bool { + let origin = origin_header.trim(); + if origin == "null" { + return false; + } + if let Ok(url) = url::Url::parse(origin) + && let Some(host) = url.host_str() { + return is_allowed_host(host, bind_address); + } + false +} + +fn verify_token_constant_time(provided: &str, expected: &str) -> bool { + let prov_bytes = provided.as_bytes(); + let exp_bytes = expected.as_bytes(); + if prov_bytes.len() != exp_bytes.len() { + return false; + } + prov_bytes.ct_eq(exp_bytes).into() +} + +fn extract_client_token(headers: &http::HeaderMap) -> Option<&str> { + if let Some(key) = headers.get("x-api-key").and_then(|v| v.to_str().ok()) { + let trimmed = key.trim(); + if !trimmed.is_empty() { + return Some(trimmed); + } + } + if let Some(auth) = headers.get(http::header::AUTHORIZATION).and_then(|v| v.to_str().ok()) { + let trimmed = auth.trim(); + if let Some(token) = trimmed.strip_prefix("Bearer ") { + let t = token.trim(); + if !t.is_empty() { + return Some(t); + } + } + } + None +} + +async fn security_middleware( + req: Request, + next: middleware::Next, +) -> Response { + let path = req.uri().path().to_string(); + let bind_address = crate::config::bind_address(); + + // 1. Host Validation (against DNS rebinding) + if let Some(host_hdr) = req.headers().get(http::header::HOST).and_then(|v| v.to_str().ok()) + && !is_allowed_host(host_hdr, &bind_address) { + return json_error( + StatusCode::FORBIDDEN, + "permission_error", + "Untrusted Host header: DNS rebinding rejected", + ); + } + + // 2. Origin & Sec-Fetch-Site Validation (against cross-site browser attacks) + if let Some(origin_hdr) = req.headers().get(http::header::ORIGIN).and_then(|v| v.to_str().ok()) + && !is_allowed_origin(origin_hdr, &bind_address) { + return json_error( + StatusCode::FORBIDDEN, + "permission_error", + "Cross-origin requests are forbidden", + ); + } + if let Some(site) = req.headers().get("sec-fetch-site").and_then(|v| v.to_str().ok()) + && site.eq_ignore_ascii_case("cross-site") { + return json_error( + StatusCode::FORBIDDEN, + "permission_error", + "Cross-site requests are forbidden", + ); + } + + // 3. Content-Type enforcement for JSON mutation endpoints (blocks simple cross-origin requests) + let method = req.method().clone(); + if matches!(method, http::Method::POST | http::Method::PUT) { + let is_json_endpoint = matches!( + path.as_str(), + "/v1/messages" + | "/v1/messages/count_tokens" + | "/v1/responses" + | "/v1/chat/completions" + | "/v1/images/generations" + ); + if is_json_endpoint + && let Some(ct) = req + .headers() + .get(http::header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + && !ct.to_ascii_lowercase().starts_with("application/json") { + return json_error( + StatusCode::UNSUPPORTED_MEDIA_TYPE, + "invalid_request_error", + "Content-Type must be application/json", + ); + } + } + + // 4. Token Authentication (constant-time verification) + if path != "/healthz" { + if let Some(expected_token) = crate::config::local_auth_token() { + let provided = extract_client_token(req.headers()); + let valid = provided.is_some_and(|tok| verify_token_constant_time(tok, &expected_token)); + if !valid { + return json_error( + StatusCode::UNAUTHORIZED, + "authentication_error", + "Invalid or missing proxy authorization token", + ); + } + } else if crate::config::require_local_auth() { + return json_error( + StatusCode::UNAUTHORIZED, + "authentication_error", + "Proxy local authentication is required but no token is configured", + ); + } + } + + next.run(req).await +} + #[derive(Clone)] struct AppState { registry: Arc, @@ -2161,15 +2310,17 @@ fn headers_to_record(headers: &http::HeaderMap) -> Value { let mut out = Map::new(); for (key, value) in headers { if let Ok(raw) = value.to_str() { + let key_str = key.as_str(); + let lower = key_str.to_ascii_lowercase(); let value = if matches!( - key.as_str(), + key_str, CLAUDE_AGENT_HEADER | CLAUDE_PARENT_AGENT_HEADER - ) { + ) || REDACT_KEYS.contains(&lower.as_str()) { format!("[redacted len={}]", raw.len()) } else { raw.to_string() }; - out.insert(key.as_str().to_string(), Value::String(value)); + out.insert(key_str.to_string(), Value::String(value)); } } Value::Object(out) @@ -2470,6 +2621,10 @@ mod auto_review_tests { CLAUDE_PARENT_AGENT_HEADER, HeaderValue::from_static("parent-secret"), ); + headers.insert( + "authorization", + HeaderValue::from_static("Bearer secret-token-1234"), + ); let captured = headers_to_record(&headers); assert_eq!( @@ -2481,6 +2636,7 @@ mod auto_review_tests { captured[CLAUDE_PARENT_AGENT_HEADER], json!("[redacted len=13]") ); + assert_eq!(captured["authorization"], json!("[redacted len=24]")); let serialized = captured.to_string(); assert!(!serialized.contains("agent-secret")); assert!(!serialized.contains("parent-secret")); @@ -2499,3 +2655,101 @@ mod auto_review_tests { assert_eq!(classifier.model.as_deref(), Some("gpt-5.6-sol")); } } + +#[cfg(test)] +mod security_middleware_tests { + use super::*; + use tower::util::ServiceExt; + + #[tokio::test] + async fn rejects_dns_rebinding_external_host() { + let app = app(Arc::new(Registry::with_default_alias())); + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/healthz") + .header("host", "evil.attacker.com") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::FORBIDDEN); + } + + #[tokio::test] + async fn allows_loopback_host() { + let app = app(Arc::new(Registry::with_default_alias())); + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/healthz") + .header("host", "127.0.0.1:18765") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test] + async fn rejects_cross_origin_browser_request() { + let app = app(Arc::new(Registry::with_default_alias())); + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/healthz") + .header("origin", "https://malicious.org") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::FORBIDDEN); + } + + #[tokio::test] + async fn rejects_cross_site_fetch() { + let app = app(Arc::new(Registry::with_default_alias())); + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/healthz") + .header("sec-fetch-site", "cross-site") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::FORBIDDEN); + } + + #[tokio::test] + async fn rejects_non_json_content_type_on_mutation_endpoints() { + let app = app(Arc::new(Registry::with_default_alias())); + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/messages") + .header("content-type", "text/plain") + .body(Body::from("model=gpt-5.5")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE); + } + + #[tokio::test] + async fn constant_time_token_verification_works() { + assert!(verify_token_constant_time("secret123", "secret123")); + assert!(!verify_token_constant_time("secret123", "secret124")); + assert!(!verify_token_constant_time("short", "longer_secret")); + } +} From 4b02167792ff00aa7f824253d2d0d012ee980a87 Mon Sep 17 00:00:00 2001 From: parth2024-tech Date: Sun, 8 Nov 2026 17:50:42 +0530 Subject: [PATCH 2/3] feat(hardening): add append-only tool ordering, disk compaction persistence, and secret scrubbing --- src/auth.rs | 5 + src/main.rs | 7 +- src/providers/codex/auth/manager.rs | 7 +- src/providers/codex/compaction.rs | 90 ++++++++++++++- src/providers/codex/mod.rs | 1 + src/providers/codex/translate/request.rs | 6 +- src/providers/kimi/translate/request.rs | 6 +- src/providers/opencode/chat.rs | 6 +- src/providers/translate_shared.rs | 135 +++++++++++++++++++++++ src/retry.rs | 45 ++++++-- src/server.rs | 57 ++++++++++ tests/codex_empty_tool_completion.rs | 6 + 12 files changed, 346 insertions(+), 25 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 9700efea..2dde2cc7 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -371,6 +371,11 @@ impl Drop for FileLock { } } +pub fn acquire_lock(path: &str) -> io::Result { + let lock_path = format!("{path}.lock"); + FileLock::acquire(std::path::Path::new(&lock_path)) +} + pub fn write_atomically(path: &str, value: &T) -> Result<()> { let dir = std::path::Path::new(path) .parent() diff --git a/src/main.rs b/src/main.rs index 1899dbba..00b94f17 100644 --- a/src/main.rs +++ b/src/main.rs @@ -200,7 +200,7 @@ async fn run_service(config: ServerConfig) -> Result<()> { let _ = stopped.await; }); tokio::pin!(server); - tokio::select! { + let outcome = tokio::select! { result = &mut server => result, signal = signals.recv() => { signal?; @@ -209,11 +209,14 @@ async fn run_service(config: ServerConfig) -> Result<()> { result = &mut server => result, signal = signals.recv() => { signal?; + claude_code_proxy::providers::codex::persist_compaction_states(); std::process::exit(130); } } } - } + }; + claude_code_proxy::providers::codex::persist_compaction_states(); + outcome } #[cfg(unix)] diff --git a/src/providers/codex/auth/manager.rs b/src/providers/codex/auth/manager.rs index 6a8981cb..efc1f0ab 100644 --- a/src/providers/codex/auth/manager.rs +++ b/src/providers/codex/auth/manager.rs @@ -86,8 +86,11 @@ impl> CodexAuthManager { ) -> Result { let _refresh_guard = self.refresh_lock.lock().await; - // Reload from durable storage after acquiring the single-flight lock. - // Another request may have rotated and persisted the token while this + let auth_path = self.store.auth_path(); + let _file_lock = crate::auth::acquire_lock(&auth_path).ok(); + + // Reload from durable storage after acquiring the single-flight cross-process lock. + // Another process or request may have rotated and persisted the token while this // caller was waiting. let current = self .load_auth()? diff --git a/src/providers/codex/compaction.rs b/src/providers/codex/compaction.rs index 12033d02..bbae706d 100644 --- a/src/providers/codex/compaction.rs +++ b/src/providers/codex/compaction.rs @@ -47,6 +47,8 @@ struct CompactionState { updated_at: u64, } +use serde::{Deserialize, Serialize}; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct CompactionAttempt(u64); @@ -55,6 +57,81 @@ pub struct CompactionReplay { pub attempt: CompactionAttempt, } +#[derive(Serialize, Deserialize)] +struct PersistedCompactionState { + attempt: u64, + model: String, + native_history: Vec, + portable_summary: String, + updated_at: u64, +} + +#[derive(Serialize, Deserialize, Default)] +struct PersistedCompactionStore { + sessions: HashMap, +} + +fn sessions_file_path() -> std::path::PathBuf { + crate::paths::state_dir().join("sessions.json") +} + +fn load_persisted_states(registry: &mut CompactionRegistry, now: u64) { + let path = sessions_file_path(); + let Ok(data) = std::fs::read_to_string(&path) else { + return; + }; + let Ok(store) = serde_json::from_str::(&data) else { + return; + }; + for (session_id, item) in store.sessions { + if now.saturating_sub(item.updated_at) <= STATE_TTL_MS { + let state = CompactionState { + attempt: CompactionAttempt(item.attempt), + model: item.model, + native_history: item.native_history, + phase: CompactionPhase::Anchored { + portable_summary: item.portable_summary, + }, + updated_at: item.updated_at, + }; + registry.states.insert(session_id, state); + } + } + update_total_bytes(registry); +} + +pub fn persist_compaction_states() { + let mut guard = match REGISTRY.lock() { + Ok(g) => g, + Err(poisoned) => poisoned.into_inner(), + }; + let Some(registry) = guard.as_mut() else { + return; + }; + let now = now_ms(); + evict_states(registry, now); + + let mut store = PersistedCompactionStore::default(); + for (session_id, state) in ®istry.states { + if let CompactionPhase::Anchored { portable_summary } = &state.phase { + store.sessions.insert( + session_id.clone(), + PersistedCompactionState { + attempt: state.attempt.0, + model: state.model.clone(), + native_history: state.native_history.clone(), + portable_summary: portable_summary.clone(), + updated_at: state.updated_at, + }, + ); + } + } + + let path = sessions_file_path(); + let path_str = path.to_string_lossy(); + let _ = crate::auth::write_atomically(&path_str, &store); +} + #[derive(Default)] struct CompactionRegistry { states: HashMap, @@ -101,7 +178,11 @@ pub fn begin_compaction(session_id: &str, model: &str) -> CompactionAttempt { }; let now = state.updated_at; let mut guard = REGISTRY.lock().unwrap(); - let registry = guard.get_or_insert_with(CompactionRegistry::default); + let registry = guard.get_or_insert_with(|| { + let mut reg = CompactionRegistry::default(); + load_persisted_states(&mut reg, now); + reg + }); evict_states(registry, now); registry.states.insert(session_id.to_string(), state); evict_states(registry, now); @@ -179,7 +260,12 @@ pub fn activate_compaction( return false; } evict_states(registry, now); - registry.states.contains_key(session_id) + let is_anchored = registry.states.contains_key(session_id); + drop(guard); + if is_anchored { + persist_compaction_states(); + } + is_anchored } pub fn apply_compaction_replay( diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index 3ef4b20c..f7bb0d55 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -42,6 +42,7 @@ use self::compaction::{ CompactionAttempt, CompactionError, abort_compaction_attempt, activate_compaction, apply_compaction_replay, begin_compaction, request_compaction, store_compaction, }; +pub use self::compaction::persist_compaction_states; use self::continuation::{ ContinuationReservation, abort_continuation_for_owner, continuation_candidate_for_owner, record_continuation_for_owner, diff --git a/src/providers/codex/translate/request.rs b/src/providers/codex/translate/request.rs index 1148a948..c8dbba70 100644 --- a/src/providers/codex/translate/request.rs +++ b/src/providers/codex/translate/request.rs @@ -656,11 +656,9 @@ fn read_tools(req: &MessagesRequest) -> Result>, anyho Value::Array(a) => a, _ => return Ok(None), }; + let sorted_tools = crate::providers::translate_shared::sort_tools_stable_prefix(tools_arr, req); let mut out = Vec::new(); - for tool in tools_arr { - if crate::providers::translate_shared::should_skip_tool(tool, req) { - continue; - } + for tool in sorted_tools { let tool_type = tool .get("type") .and_then(|v| v.as_str()) diff --git a/src/providers/kimi/translate/request.rs b/src/providers/kimi/translate/request.rs index 2343e22a..2929fbc8 100644 --- a/src/providers/kimi/translate/request.rs +++ b/src/providers/kimi/translate/request.rs @@ -245,11 +245,9 @@ fn read_tools(req: &MessagesRequest) -> Result, anyhow::Error> { Value::Array(a) => a, _ => return Ok(Vec::new()), }; + let sorted_tools = crate::providers::translate_shared::sort_tools_stable_prefix(tools_arr, req); let mut out = Vec::new(); - for tool in tools_arr { - if crate::providers::translate_shared::should_skip_tool(tool, req) { - continue; - } + for tool in sorted_tools { let name = tool .get("name") .and_then(|v| v.as_str()) diff --git a/src/providers/opencode/chat.rs b/src/providers/opencode/chat.rs index 4bcdfcab..17776ecf 100644 --- a/src/providers/opencode/chat.rs +++ b/src/providers/opencode/chat.rs @@ -292,9 +292,9 @@ fn read_tools(req: &MessagesRequest) -> anyhow::Result> { let tools = value .as_array() .ok_or_else(|| anyhow::anyhow!("tools must be an array"))?; - tools - .iter() - .filter(|tool| !crate::providers::translate_shared::should_skip_tool(tool, req)) + let sorted_tools = crate::providers::translate_shared::sort_tools_stable_prefix(tools, req); + sorted_tools + .into_iter() .map(|tool| { let name = tool .get("name") diff --git a/src/providers/translate_shared.rs b/src/providers/translate_shared.rs index ef3c2ae8..f1f8b226 100644 --- a/src/providers/translate_shared.rs +++ b/src/providers/translate_shared.rs @@ -274,6 +274,55 @@ pub fn should_skip_tool(tool_val: &Value, req: &MessagesRequest) -> bool { false } +pub fn first_seen_tool_resolution_orders(req: &MessagesRequest) -> std::collections::HashMap { + let mut order_map = std::collections::HashMap::new(); + let mut counter = 0; + for msg in &req.messages { + if let Some(arr) = msg.content.as_array() { + for block in arr { + if block.get("type").and_then(Value::as_str) == Some("tool_result") + && let Some(content_arr) = block.get("content").and_then(Value::as_array) + { + for sub in content_arr { + if sub.get("type").and_then(Value::as_str) == Some("tool_reference") + && let Some(name) = sub.get("tool_name").and_then(Value::as_str) + { + order_map.entry(name.to_string()).or_insert_with(|| { + let idx = counter; + counter += 1; + idx + }); + } + } + } + } + } + } + order_map +} + +pub fn sort_tools_stable_prefix<'a>(tools_arr: &'a [Value], req: &MessagesRequest) -> Vec<&'a Value> { + let order_map = first_seen_tool_resolution_orders(req); + let mut partitioned: Vec<(usize, usize, usize, &'a Value)> = Vec::new(); + + for (original_idx, tool) in tools_arr.iter().enumerate() { + if should_skip_tool(tool, req) { + continue; + } + let is_deferred = tool.get("defer_loading").and_then(Value::as_bool).unwrap_or(false); + let name = tool.get("name").and_then(Value::as_str).unwrap_or(""); + if is_deferred { + let resolution_order = order_map.get(name).copied().unwrap_or(usize::MAX); + partitioned.push((1, resolution_order, original_idx, tool)); + } else { + partitioned.push((0, 0, original_idx, tool)); + } + } + + partitioned.sort_by_key(|&(group, res_order, orig_idx, _)| (group, res_order, orig_idx)); + partitioned.into_iter().map(|(_, _, _, tool)| tool).collect() +} + pub fn strip_tool_schema_patterns(schema: &mut Value) { let Some(schema) = schema.as_object_mut() else { return; @@ -314,3 +363,89 @@ pub fn strip_tool_schema_patterns(schema: &mut Value) { } } } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn sort_tools_stable_prefix_preserves_base_and_appends_resolved_in_turn_order() { + let tools = vec![ + json!({"name": "base1", "defer_loading": false}), + json!({"name": "deferred_early_in_decl", "defer_loading": true}), + json!({"name": "base2", "defer_loading": false}), + json!({"name": "deferred_late_in_decl", "defer_loading": true}), + json!({"name": "DeferredToolPlaceholder", "defer_loading": true}), + ]; + + // Turn 1: No deferred tools resolved + let req_turn1: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "hello"}] + })).unwrap(); + let sorted1 = sort_tools_stable_prefix(&tools, &req_turn1); + let names1: Vec<&str> = sorted1.iter().map(|t| t["name"].as_str().unwrap()).collect(); + assert_eq!(names1, vec!["base1", "base2"]); + + // Turn 2: deferred_late_in_decl resolved first + let req_turn2: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-3-5-sonnet", + "messages": [ + {"role": "user", "content": "hello"}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_1", + "content": [ + {"type": "tool_reference", "tool_name": "deferred_late_in_decl"} + ] + } + ] + } + ] + })).unwrap(); + let sorted2 = sort_tools_stable_prefix(&tools, &req_turn2); + let names2: Vec<&str> = sorted2.iter().map(|t| t["name"].as_str().unwrap()).collect(); + // deferred_late_in_decl is appended at the tail + assert_eq!(names2, vec!["base1", "base2", "deferred_late_in_decl"]); + + // Turn 3: deferred_early_in_decl resolved second + let req_turn3: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-3-5-sonnet", + "messages": [ + {"role": "user", "content": "hello"}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_1", + "content": [ + {"type": "tool_reference", "tool_name": "deferred_late_in_decl"} + ] + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_2", + "content": [ + {"type": "tool_reference", "tool_name": "deferred_early_in_decl"} + ] + } + ] + } + ] + })).unwrap(); + let sorted3 = sort_tools_stable_prefix(&tools, &req_turn3); + let names3: Vec<&str> = sorted3.iter().map(|t| t["name"].as_str().unwrap()).collect(); + // Notice: ["base1", "base2", "deferred_late_in_decl"] is an immutable exact prefix of names3! + assert_eq!(names3, vec!["base1", "base2", "deferred_late_in_decl", "deferred_early_in_decl"]); + } +} diff --git a/src/retry.rs b/src/retry.rs index 70603683..6d36db66 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -16,21 +16,41 @@ pub fn should_retry_status(status: u16) -> bool { matches!(status, 429 | 500 | 502 | 503 | 504) } +fn get_base_delay_ms() -> u64 { + if let Ok(val) = std::env::var("CCP_RETRY_BASE_MS") + .or_else(|_| std::env::var("CCP_RETRY_INITIAL_DELAY_MS")) + && let Ok(ms) = val.parse::() + { + return ms; + } + RETRY_INITIAL_DELAY_MS +} + +fn get_max_delay_ms() -> u64 { + if let Ok(val) = std::env::var("CCP_RETRY_MAX_DELAY_MS") + && let Ok(ms) = val.parse::() + { + return ms; + } + RETRY_MAX_DELAY_MS +} + pub fn compute_backoff_delay(attempt: u32, retry_after: Option<&str>) -> BackoffOutcome { + let max_delay = get_max_delay_ms(); if let Some(raw) = retry_after && let Ok(raw_secs) = raw.parse::() { let target_ms = (raw_secs * 1000.0).ceil() as u64; return BackoffOutcome { - wait_ms: target_ms.min(RETRY_MAX_DELAY_MS), - exceeds_budget: target_ms > RETRY_MAX_DELAY_MS, + wait_ms: target_ms.min(max_delay), + exceeds_budget: target_ms > max_delay, }; } - let mut exp = - RETRY_INITIAL_DELAY_MS.saturating_mul(RETRY_BACKOFF_FACTOR.saturating_pow(attempt)); - if exp > RETRY_MAX_DELAY_MS { - exp = RETRY_MAX_DELAY_MS; + let base = get_base_delay_ms(); + let mut exp = base.saturating_mul(RETRY_BACKOFF_FACTOR.saturating_pow(attempt)); + if exp > max_delay { + exp = max_delay; } let jitter = exp / 2; let wait_ms = (exp / 2) + (jitter / 2); @@ -49,10 +69,19 @@ pub fn set_zero_retry_delay_for_tests(enabled: bool) { } pub async fn sleep(ms: u64) { - if ZERO_RETRY_DELAY_FOR_TESTS.load(Ordering::SeqCst) { + if ZERO_RETRY_DELAY_FOR_TESTS.load(Ordering::SeqCst) + || std::env::var("CCP_ZERO_RETRY_DELAY").is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true")) + { return; } - tokio::time::sleep(Duration::from_millis(ms)).await; + let effective_ms = if let Ok(val) = std::env::var("CCP_RETRY_BASE_MS") + && let Ok(cap) = val.parse::() + { + ms.min(cap) + } else { + ms + }; + tokio::time::sleep(Duration::from_millis(effective_ms)).await; } #[cfg(test)] diff --git a/src/server.rs b/src/server.rs index 4b21992a..9a4da9ca 100644 --- a/src/server.rs +++ b/src/server.rs @@ -2239,8 +2239,32 @@ fn sanitize_path_part(raw: &str) -> String { } } +static SECRET_REGEXES: std::sync::LazyLock> = std::sync::LazyLock::new(|| { + vec![ + // Anthropic keys + (regex_lite::Regex::new(r"sk-ant-[a-zA-Z0-9_\-]{20,}").unwrap(), "[redacted sk-ant-***]"), + // OpenAI / project keys + (regex_lite::Regex::new(r"sk-(?:proj-)?[a-zA-Z0-9_\-]{20,}").unwrap(), "[redacted sk-***]"), + // GitHub tokens (classic, fine-grained, oauth, user, server, refresh) + (regex_lite::Regex::new(r"(?:ghp|gho|ghu|ghs|ghr|github_pat)_[a-zA-Z0-9_]{16,}").unwrap(), "[redacted ghp_***]"), + // AWS Access Key IDs + (regex_lite::Regex::new(r"\b(?:AKIA|ABIA|ACCA|ASIA)[0-9A-Z]{16}\b").unwrap(), "[redacted AKIA***]"), + // PEM private keys + (regex_lite::Regex::new(r"-----BEGIN (?:[A-Z ]* )?PRIVATE KEY-----[\s\S]*?-----END (?:[A-Z ]* )?PRIVATE KEY-----").unwrap(), "[redacted PRIVATE KEY]"), + ] +}); + +pub fn redact_secrets_in_text(text: &str) -> String { + let mut scrubbed = text.to_string(); + for (re, replacement) in SECRET_REGEXES.iter() { + scrubbed = re.replace_all(&scrubbed, *replacement).to_string(); + } + scrubbed +} + fn redact_error_value(value: Value) -> Value { match value { + Value::String(s) => Value::String(redact_secrets_in_text(&s)), Value::Array(values) => Value::Array(values.into_iter().map(redact_error_value).collect()), Value::Object(fields) => { let mut out = Map::new(); @@ -2752,4 +2776,37 @@ mod security_middleware_tests { assert!(!verify_token_constant_time("secret123", "secret124")); assert!(!verify_token_constant_time("short", "longer_secret")); } + + #[test] + fn redacts_prompt_secrets_in_error_captures() { + let input = serde_json::json!({ + "messages": [ + { + "role": "user", + "content": "Here is my key: sk-ant-api03-123456789012345678901234 and github: ghp_12345678901234567890 and aws: AKIAIOSFODNN7EXAMPLE and openai: sk-proj-123456789012345678901234" + }, + { + "role": "assistant", + "content": "-----BEGIN RSA PRIVATE KEY-----\nMIIEowIBAAKCAQEA0\n-----END RSA PRIVATE KEY-----" + } + ], + "authorization": "Bearer secret-token" + }); + + let redacted = redact_error_value(input); + assert_eq!(redacted["authorization"], serde_json::json!("[redacted len=19]")); + let user_content = redacted["messages"][0]["content"].as_str().unwrap(); + assert!(!user_content.contains("sk-ant-api03")); + assert!(user_content.contains("[redacted sk-ant-***]")); + assert!(!user_content.contains("ghp_12345")); + assert!(user_content.contains("[redacted ghp_***]")); + assert!(!user_content.contains("AKIAIOSFODNN7EXAMPLE")); + assert!(user_content.contains("[redacted AKIA***]")); + assert!(!user_content.contains("sk-proj-12345")); + assert!(user_content.contains("[redacted sk-***]")); + + let assistant_content = redacted["messages"][1]["content"].as_str().unwrap(); + assert!(!assistant_content.contains("MIIEowIBAAKCAQEA0")); + assert!(assistant_content.contains("[redacted PRIVATE KEY]")); + } } diff --git a/tests/codex_empty_tool_completion.rs b/tests/codex_empty_tool_completion.rs index 897db843..ea79b1d3 100644 --- a/tests/codex_empty_tool_completion.rs +++ b/tests/codex_empty_tool_completion.rs @@ -41,6 +41,10 @@ fn next_events(script: &Mutex