diff --git a/CLAUDE.md b/CLAUDE.md index 07028f2..38f42a7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -12,7 +12,7 @@ See PLAN.md for the full specification and TASKS.md for phased implementation st **Working (verified):** - Build: `cargo build` clean, `cargo clippy -- -D warnings` clean -- Tests: ~367 tests passing, 4 ignored (live API) +- Tests: ~390 tests passing, 4 ignored (live API) - Full Anthropic Messages API translation: non-streaming, streaming SSE, tool calling, file/document blocks - Proxy middleware: health, auth, request ID, size limits, concurrency limits, retry with backoff - Compatibility endpoints: /v1/models, count_tokens (approximate via tiktoken), batches (stub) @@ -27,7 +27,7 @@ See PLAN.md for the full specification and TASKS.md for phased implementation st ```bash cargo build # build everything -cargo test # run all tests (~367 tests, 4 ignored) +cargo test # run all tests (~390 tests, 4 ignored) cargo test -p anthropic_openai_translate # translator crate only cargo test -p anthropic_openai_proxy # proxy crate only cargo test health_endpoint # single test by name @@ -85,14 +85,17 @@ Pure translation logic, no IO. Key modules: ### `crates/proxy` (bin: `anthropic_openai_proxy`) HTTP proxy built on axum + reqwest: -- **`config.rs`**: Env-based configuration +- **`config/`**: Env-based configuration (`mod.rs`), TLS client cert setup (`tls.rs`), URL validation (`url_validation.rs`) - **`server/routes.rs`**: Axum router (POST /v1/messages, GET /health, GET /metrics, GET /v1/models, stubs for count_tokens and batches) - **`server/middleware.rs`**: Auth validation (x-api-key), request ID injection, 32MB size limit, concurrency limit, logging - **`server/sse.rs`**: SSE response helpers for Anthropic-format streaming +- **`server/streaming.rs`**: SSE streaming handler with pre-stream error propagation and backpressure +- **`server/passthrough.rs`**: Anthropic passthrough handler (no translation, forwards as-is) +- **`server/token_counting.rs`**: Approximate token counting via tiktoken - **`backend/mod.rs`**: `BackendClient` enum (OpenAI/OpenAIResponses/Vertex/GeminiOpenAI/Anthropic), `BackendError`, shared retry helpers - **`backend/openai_client.rs`**: reqwest client calling OpenAI-compatible Chat Completions with retry/backoff on 429/5xx (used for OpenAI, Vertex, and Gemini backends) - **`backend/anthropic_client.rs`**: Passthrough client forwarding Anthropic requests as-is to upstream Anthropic API (no translation) -- **`admin/`**: Admin server (localhost-only) with config management, WebSocket live updates, token auth (`auth.rs`, `db.rs`, `mod.rs`, `routes.rs`, `state.rs`) +- **`admin/`**: Admin server (localhost-only) with config management, WebSocket live updates (`ws.rs`), token auth (`auth.rs`, `db.rs`, `mod.rs`, `routes.rs`, `state.rs`) - **`admin-ui/`**: Static admin UI served by the admin server (`index.html`) - **`metrics/`**: Request count, success/error tracking, exposed via GET /metrics @@ -115,13 +118,15 @@ Client (Anthropic format) -> proxy (axum) - Retry logic: 3 retries with exponential backoff + 25% jitter, respects retry-after header. - Backoff jitter is deterministic (upper bound, not random) to keep tests predictable. - `ChatCompletionRequest` uses `#[serde(flatten)] pub extra: serde_json::Map` to capture unknown OpenAI fields (e.g., `seed`, `logprobs`, `logit_bias`, `n`, `reasoning_effort`). These pass through to OpenAI without typed handling. Only fields that require translation logic (not just forwarding) need explicit struct fields. +- DeepSeek/Qwen thinking model support: `reasoning_content` on `ChatMessage` and `ChunkDelta` maps bidirectionally to Anthropic thinking blocks. Request direction: Anthropic `Thinking` content blocks become `reasoning_content` on the assistant message. Response direction: `reasoning_content` becomes an Anthropic `Thinking` block preceding the text content. Streaming: `reasoning_content` deltas open a thinking content block, which is closed when regular `content` deltas begin. The `thinking` config (`budget_tokens`) is stripped with a warning since it has no standard OpenAI equivalent. +- Local LLM compatibility: streaming tool calls handle missing/empty IDs by generating synthetic `toolu_` IDs. `FinishReason::Unknown` (serde catch-all) maps to `end_turn` for providers like DeepSeek that use non-standard finish reasons (e.g., `insufficient_system_resource`). ## Conventions - Most source files reference their PLAN.md line ranges in a comment at the top. - Test files live alongside source (`#[cfg(test)]` modules) and in `crates/proxy/tests/` for integration tests. - Error types use `thiserror` derive macros. -- Test distribution: translator (~242 tests), proxy (~125 tests including integration/compatibility). +- Test distribution: translator (~240 tests), proxy (~150 tests including integration/compatibility). Counts shift as features are added. ## References diff --git a/crates/proxy/Cargo.toml b/crates/proxy/Cargo.toml index 6b8db3d..0684938 100644 --- a/crates/proxy/Cargo.toml +++ b/crates/proxy/Cargo.toml @@ -25,6 +25,7 @@ toml = "1.0.7" indexmap = { version = "2.13.0", features = ["serde"] } bytes = "1.11.1" subtle = "2" +sha2 = "0.10" rusqlite = { version = "0.32", features = ["bundled"] } httpdate = "1" diff --git a/crates/proxy/src/admin/db.rs b/crates/proxy/src/admin/db.rs index eb16ffe..166760e 100644 --- a/crates/proxy/src/admin/db.rs +++ b/crates/proxy/src/admin/db.rs @@ -187,17 +187,12 @@ pub fn delete_config_override(conn: &Connection, key: &str) -> rusqlite::Result< /// Delete request log entries older than the given number of days. pub fn purge_old_logs(conn: &Connection, retention_days: u32) -> rusqlite::Result { // SQLite datetime comparison: delete rows where timestamp < cutoff - let cutoff = format!( - "{}", - // Approximate: subtract seconds - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs() - .saturating_sub(retention_days as u64 * 86400) - ); - // Convert epoch to ISO 8601 for comparison - let cutoff_iso = epoch_to_iso8601(cutoff.parse::().unwrap_or(0)); + let cutoff = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + .saturating_sub(retention_days as u64 * 86400); + let cutoff_iso = epoch_to_iso8601(cutoff); let changed = conn.execute( "DELETE FROM request_log WHERE timestamp < ?1", params![cutoff_iso], diff --git a/crates/proxy/src/admin/routes.rs b/crates/proxy/src/admin/routes.rs index c18db06..dfb97bd 100644 --- a/crates/proxy/src/admin/routes.rs +++ b/crates/proxy/src/admin/routes.rs @@ -134,6 +134,7 @@ async fn serve_spa() -> impl IntoResponse { connect-src 'self'; frame-ancestors 'none'", ), ("x-frame-options", "DENY"), + ("referrer-policy", "no-referrer"), ], SPA_HTML, ) @@ -237,6 +238,11 @@ async fn put_config( } } + // Serialize config writes so concurrent requests cannot interleave + // Phase 1 (SQLite) and Phase 2 (in-memory), which would leave them + // inconsistent. + let _config_guard = shared.config_write_lock.lock().await; + // Phase 1: Persist to SQLite first. If the process crashes between // phases, the database is the source of truth and config is restored // on restart. Reversing the order would lose updates on crash. @@ -284,6 +290,8 @@ async fn put_config( } } + drop(_config_guard); + // Broadcast config changes. for (key, value) in &db_writes { let _ = shared @@ -345,14 +353,17 @@ async fn delete_config_override( Json(serde_json::json!({"error": "override not found"})), ) .into_response(), - Some(Err(e)) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": e.to_string()})), - ) - .into_response(), + Some(Err(e)) => { + tracing::error!(error = %e, "delete_config_override failed"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "internal database error"})), + ) + .into_response() + } None => ( StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": "task panicked"})), + Json(serde_json::json!({"error": "internal error"})), ) .into_response(), } @@ -470,10 +481,13 @@ async fn get_requests( "limit": limit, "offset": offset, })), - Some(Err(e)) => Json(serde_json::json!({ - "error": e.to_string(), - "requests": [], - })), + Some(Err(e)) => { + tracing::error!(error = %e, "query_request_log failed"); + Json(serde_json::json!({ + "error": "internal database error", + "requests": [], + })) + } None => Json(serde_json::json!({ "error": "task panicked", "requests": [], @@ -499,14 +513,17 @@ async fn get_request_by_id( Json(serde_json::json!({"error": "request not found"})), ) .into_response(), - Some(Err(e)) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": e.to_string()})), - ) - .into_response(), + Some(Err(e)) => { + tracing::error!(error = %e, "get_request_by_id failed"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "internal database error"})), + ) + .into_response() + } None => ( StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": "task panicked"})), + Json(serde_json::json!({"error": "internal error"})), ) .into_response(), } diff --git a/crates/proxy/src/admin/state.rs b/crates/proxy/src/admin/state.rs index 164369c..0d20b26 100644 --- a/crates/proxy/src/admin/state.rs +++ b/crates/proxy/src/admin/state.rs @@ -33,6 +33,9 @@ pub struct SharedState { pub log_tx: tokio::sync::mpsc::Sender, /// Closure to reload tracing filter at runtime. None in tests. pub log_reload: Option, + /// Serializes config write operations (Phase 1: SQLite + Phase 2: in-memory) + /// so concurrent PUT /admin/api/config requests cannot interleave. + pub config_write_lock: Arc>, } /// Run a synchronous closure against the SQLite connection on the blocking @@ -115,6 +118,7 @@ impl SharedState { backend_metrics: Arc::new(HashMap::new()), log_tx, log_reload: None, + config_write_lock: Arc::new(tokio::sync::Mutex::new(())), } } } diff --git a/crates/proxy/src/config/mod.rs b/crates/proxy/src/config/mod.rs index cd6c06b..63e2a3d 100644 --- a/crates/proxy/src/config/mod.rs +++ b/crates/proxy/src/config/mod.rs @@ -63,6 +63,20 @@ pub struct Config { pub openai_api_format: OpenAIApiFormat, } +/// Validate that a GCP identifier (project ID, region) contains only safe characters. +/// Prevents URL injection when these values are interpolated into Vertex AI endpoint URLs. +fn validate_gcp_identifier(name: &str, value: &str) { + if value.is_empty() + || !value + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_' || b == b'.') + { + panic!( + "{name} contains invalid characters: only alphanumeric, '-', '_', '.' are allowed, got: {value}" + ); + } +} + impl Config { pub fn from_env() -> Self { let backend_str = std::env::var("BACKEND").unwrap_or_else(|_| "openai".into()); @@ -122,6 +136,8 @@ impl Config { .unwrap_or_else(|_| panic!("VERTEX_PROJECT is required when BACKEND=vertex")); let region = std::env::var("VERTEX_REGION") .unwrap_or_else(|_| panic!("VERTEX_REGION is required when BACKEND=vertex")); + validate_gcp_identifier("VERTEX_PROJECT", &project); + validate_gcp_identifier("VERTEX_REGION", ®ion); let backend_auth = if let Ok(api_key) = std::env::var("VERTEX_API_KEY") { BackendAuth::GoogleApiKey(api_key) @@ -476,6 +492,8 @@ impl MultiConfig { .region .as_deref() .unwrap_or_else(|| panic!("backend '{name}': 'region' is required for vertex")); + validate_gcp_identifier("project", project); + validate_gcp_identifier("region", region); let base_url = tb.base_url.clone().unwrap_or_else(|| { format!( diff --git a/crates/proxy/src/config/url_validation.rs b/crates/proxy/src/config/url_validation.rs index 0628b44..5c9104a 100644 --- a/crates/proxy/src/config/url_validation.rs +++ b/crates/proxy/src/config/url_validation.rs @@ -49,18 +49,28 @@ pub fn validate_base_url(raw: &str) -> Result<(), String> { .port() .unwrap_or(if parsed.scheme() == "https" { 443 } else { 80 }); let lookup = format!("{domain}:{port}"); - if let Ok(addrs) = std::net::ToSocketAddrs::to_socket_addrs(&lookup) { - for addr in addrs { - if is_private_ip(addr.ip()) { - return Err(format!( - "hostname '{domain}' resolves to private/loopback IP {}, not allowed", - addr.ip() - )); + match std::net::ToSocketAddrs::to_socket_addrs(&lookup) { + Ok(addrs) => { + for addr in addrs { + if is_private_ip(addr.ip()) { + return Err(format!( + "hostname '{domain}' resolves to private/loopback IP {}, not allowed", + addr.ip() + )); + } } } + Err(e) => { + // Allow through: the domain may not be resolvable in the + // build/test environment but will work at runtime. The + // runtime SsrfSafeDnsResolver provides connection-time protection. + tracing::warn!( + domain = %domain, + error = %e, + "DNS resolution failed at startup; domain will be validated at connection time" + ); + } } - // If DNS resolution fails, allow it (the domain may not be resolvable - // in the build/test environment but will work at runtime). } } diff --git a/crates/proxy/src/main.rs b/crates/proxy/src/main.rs index 9e5ad2a..5fea814 100644 --- a/crates/proxy/src/main.rs +++ b/crates/proxy/src/main.rs @@ -53,7 +53,21 @@ async fn main() { if let Ok(overrides) = admin::db::get_config_overrides(&conn) { for (key, value, _) in &overrides { match key.as_str() { - "log_level" => runtime_config.log_level = value.clone(), + "log_level" => { + // Apply the same allowlist enforced by the admin API to + // prevent a tampered SQLite database from enabling trace-level + // logging, which would expose API keys in HTTP headers. + const ALLOWED_LOG_LEVELS: &[&str] = &["error", "warn", "info", "debug"]; + let normalized = value.trim().to_lowercase(); + if ALLOWED_LOG_LEVELS.contains(&normalized.as_str()) { + runtime_config.log_level = normalized; + } else { + tracing::warn!( + value = %value, + "ignoring invalid log_level override from database" + ); + } + } "log_bodies" => runtime_config.log_bodies = value == "true", k if k.ends_with(".big_model") => { let backend = k.strip_suffix(".big_model").unwrap(); @@ -114,6 +128,7 @@ async fn main() { backend_metrics: Arc::new(backend_metrics), log_tx, log_reload: Some(log_reload), + config_write_lock: Arc::new(tokio::sync::Mutex::new(())), }; // Admin token: use env var or generate random UUID written to a file. @@ -124,9 +139,12 @@ async fn main() { // Write token to file with restrictive permissions instead of stderr, // because stderr is captured by container log drivers in production. if let Err(e) = write_token_file(&token_path, &token) { - // Fall back to stderr if file write fails (e.g., read-only filesystem). - eprintln!("WARNING: could not write admin token to {token_path}: {e}"); - eprintln!("Admin token: {token}"); + // Do not print the token to stderr: container log drivers capture + // stderr and persist it in centralized logging systems. + panic!( + "Cannot write admin token to {token_path}: {e}. \ + Set ADMIN_TOKEN env var explicitly or ensure the path is writable." + ); } else { // Log the path, not the token itself. tracing::info!(path = %token_path, "generated admin token written to file (set ADMIN_TOKEN env var to avoid this)"); @@ -263,7 +281,14 @@ fn write_token_file(path: &str, token: &str) -> std::io::Result<()> { }; #[cfg(not(unix))] - let mut file = std::fs::File::create(path)?; + let mut file = { + tracing::warn!( + path = %path, + "non-Unix platform: admin token file may be world-readable. \ + Set ADMIN_TOKEN env var explicitly in production." + ); + std::fs::File::create(path)? + }; file.write_all(token.as_bytes())?; file.write_all(b"\n")?; diff --git a/crates/proxy/src/server/middleware.rs b/crates/proxy/src/server/middleware.rs index 0037eb1..4efd711 100644 --- a/crates/proxy/src/server/middleware.rs +++ b/crates/proxy/src/server/middleware.rs @@ -9,12 +9,13 @@ use axum::{ middleware::Next, response::{IntoResponse, Json, Response}, }; +use sha2::{Digest, Sha256}; use std::sync::LazyLock; use subtle::ConstantTimeEq; -/// Allowed API keys loaded from `PROXY_API_KEYS` (comma-separated). -/// When the list is empty, any non-empty key is accepted (open-relay mode). -static ALLOWED_API_KEYS: LazyLock> = LazyLock::new(|| { +/// Pre-hashed allowed API keys for constant-time comparison without +/// leaking key length via timing. Each key is SHA-256 hashed at startup. +static ALLOWED_KEY_HASHES: LazyLock> = LazyLock::new(|| { let keys: Vec = std::env::var("PROXY_API_KEYS") .unwrap_or_default() .split(',') @@ -22,12 +23,31 @@ static ALLOWED_API_KEYS: LazyLock> = LazyLock::new(|| { .filter(|s| !s.is_empty()) .collect(); if keys.is_empty() { - tracing::warn!( - "PROXY_API_KEYS is not set: proxy accepts ANY non-empty key (open-relay mode). \ - Set PROXY_API_KEYS to restrict access." - ); + let open_relay = std::env::var("PROXY_OPEN_RELAY") + .map(|v| v == "true" || v == "1") + .unwrap_or(false); + if open_relay { + tracing::warn!( + "PROXY_OPEN_RELAY=true: proxy accepts ANY non-empty key. \ + Set PROXY_API_KEYS to restrict access." + ); + } else { + tracing::error!( + "PROXY_API_KEYS is not set and PROXY_OPEN_RELAY is not enabled. \ + The proxy will reject all requests. Set PROXY_API_KEYS or \ + set PROXY_OPEN_RELAY=true to allow unauthenticated access." + ); + } } - keys + keys.iter().map(|k| Sha256::digest(k.as_bytes()).into()).collect() +}); + +/// Whether open-relay mode is explicitly enabled via PROXY_OPEN_RELAY=true. +static OPEN_RELAY: LazyLock = LazyLock::new(|| { + ALLOWED_KEY_HASHES.is_empty() + && std::env::var("PROXY_OPEN_RELAY") + .map(|v| v == "true" || v == "1") + .unwrap_or(false) }); /// Validate that the request carries a valid API key. @@ -64,14 +84,14 @@ pub async fn validate_auth( } }; - // Validate against the allowlist with constant-time comparison. Without - // this, an attacker could infer the correct key byte-by-byte by measuring - // response times (timing side-channel attack). - let is_allowed = ALLOWED_API_KEYS.iter().any(|allowed| { - allowed.len() == credential.len() - && bool::from(allowed.as_bytes().ct_eq(credential.as_bytes())) - }); - if !ALLOWED_API_KEYS.is_empty() && !is_allowed { + // Compare SHA-256 hashes of the credential against pre-hashed allowed keys. + // Hashing eliminates the timing side-channel on key length: all comparisons + // operate on fixed-size 32-byte digests regardless of original key length. + let credential_hash: [u8; 32] = Sha256::digest(credential.as_bytes()).into(); + let is_allowed = ALLOWED_KEY_HASHES + .iter() + .any(|h| bool::from(h.ct_eq(&credential_hash))); + if !ALLOWED_KEY_HASHES.is_empty() && !is_allowed { let err = create_anthropic_error( anthropic::ErrorType::AuthenticationError, "Invalid API key.".to_string(), @@ -79,6 +99,15 @@ pub async fn validate_auth( ); return Err((StatusCode::UNAUTHORIZED, Json(err)).into_response()); } + // Reject if no keys configured and open-relay not explicitly enabled. + if ALLOWED_KEY_HASHES.is_empty() && !*OPEN_RELAY { + let err = create_anthropic_error( + anthropic::ErrorType::AuthenticationError, + "Server not configured for access. Contact the administrator.".to_string(), + None, + ); + return Err((StatusCode::UNAUTHORIZED, Json(err)).into_response()); + } Ok(next.run(request).await) } diff --git a/crates/proxy/src/server/routes.rs b/crates/proxy/src/server/routes.rs index 088a89c..99af9db 100644 --- a/crates/proxy/src/server/routes.rs +++ b/crates/proxy/src/server/routes.rs @@ -159,9 +159,9 @@ pub fn app_multi_with_shared(config: MultiConfig, shared: Option) - backend_metrics: Arc::new(backend_metrics), }; - // Health and metrics are public, bypass auth and concurrency limits. - Router::new() - .route("/health", get(health)) + // Metrics requires auth (prevents unauthenticated reconnaissance of + // backend names and traffic patterns). + let metrics_route = Router::new() .route( "/metrics", get(|State(gs): State| async move { @@ -189,6 +189,12 @@ pub fn app_multi_with_shared(config: MultiConfig, shared: Option) - })) }), ) + .layer(axum::middleware::from_fn(super::middleware::validate_auth)); + + // Health is public (no auth required). + Router::new() + .route("/health", get(health)) + .merge(metrics_route) .merge(router) .fallback(fallback_not_found) .layer(axum::middleware::from_fn(super::middleware::add_request_id)) @@ -235,12 +241,14 @@ fn backend_router(state: AppState, is_anthropic: bool) -> Router { /// Reject requests when the concurrency limit is reached (429), rather than /// queueing them like Tower's ConcurrencyLimitLayer would. +/// The permit is stored in request extensions so streaming handlers can hold +/// it until the stream completes (not just until headers are sent). async fn enforce_concurrency( State(state): State, - request: axum::extract::Request, + mut request: axum::extract::Request, next: axum::middleware::Next, ) -> Response { - let Ok(_permit) = state.concurrency.try_acquire() else { + let Ok(permit) = state.concurrency.clone().try_acquire_owned() else { let err = mapping::errors_map::create_anthropic_error( anthropic::ErrorType::RateLimitError, "Proxy concurrency limit reached".to_string(), @@ -248,9 +256,16 @@ async fn enforce_concurrency( ); return (StatusCode::TOO_MANY_REQUESTS, Json(err)).into_response(); }; + request.extensions_mut().insert(ConcurrencyPermit(Arc::new(permit))); next.run(request).await } +/// Wrapper so OwnedSemaphorePermit can be stored in request extensions. +/// The field is never read directly; it exists as an RAII guard to hold +/// the permit until the struct is dropped. +#[derive(Clone)] +pub(crate) struct ConcurrencyPermit(#[allow(dead_code)] pub(crate) Arc); + static MODELS_RESPONSE: std::sync::LazyLock = std::sync::LazyLock::new(|| { serde_json::json!({ "data": [ @@ -306,8 +321,12 @@ fn backend_error_to_response(error: BackendError) -> Response { async fn messages( State(state): State, headers: axum::http::HeaderMap, + permit: Option>, AnthropicJson(body): AnthropicJson, ) -> Response { + // Hold concurrency permit for streaming: passed to the spawned task so + // the permit lives until the stream completes, not just until headers are sent. + let permit = permit.map(|axum::Extension(p)| p); let ctx = RequestCtx { request_id: headers .get("x-request-id") @@ -335,7 +354,7 @@ async fn messages( } let mapped_model = state.map_model(&body.model); // Logging deferred until stream completes (inside messages_stream tasks). - match messages_stream(state, body, ctx, mapped_model).await { + match messages_stream(state, body, ctx, mapped_model, permit).await { Ok((rate_limits, sse)) => { let mut response = sse.into_response(); rate_limits.inject_anthropic_headers(response.headers_mut()); diff --git a/crates/proxy/src/server/streaming.rs b/crates/proxy/src/server/streaming.rs index 70a498f..052b30f 100644 --- a/crates/proxy/src/server/streaming.rs +++ b/crates/proxy/src/server/streaming.rs @@ -174,6 +174,7 @@ pub(crate) async fn messages_stream( body: anthropic::MessageCreateRequest, ctx: RequestCtx, mapped_model: String, + concurrency_permit: Option, ) -> Result< ( RateLimitHeaders, @@ -198,8 +199,13 @@ pub(crate) async fn messages_stream( super::routes::inject_gemini_thinking(&body, &state.backend, &mut openai_req); openai_req.model = state.map_model(&openai_req.model); let model = body.model.clone(); + let permit = concurrency_permit.clone(); tokio::spawn(async move { + // Hold concurrency permit until the stream completes, not just + // until headers are sent, so the semaphore accurately bounds + // concurrent streaming connections. + let _permit = permit; match client.chat_completion_stream(&openai_req).await { Ok((response, rate_limits)) => { rl_tx.send(Ok(rate_limits)).ok(); @@ -271,8 +277,10 @@ pub(crate) async fn messages_stream( responses_req.model = state.map_model(&responses_req.model); responses_req.stream = Some(true); let model = body.model.clone(); + let permit = concurrency_permit; tokio::spawn(async move { + let _permit = permit; match client.responses_stream(&responses_req).await { Ok((response, rate_limits)) => { rl_tx.send(Ok(rate_limits)).ok(); diff --git a/crates/proxy/tests/compatibility.rs b/crates/proxy/tests/compatibility.rs index 0e356e0..692fbc4 100644 --- a/crates/proxy/tests/compatibility.rs +++ b/crates/proxy/tests/compatibility.rs @@ -23,6 +23,8 @@ fn test_config() -> Config { } async fn spawn_test_server() -> String { + // Enable open-relay mode for tests (no PROXY_API_KEYS configured). + std::env::set_var("PROXY_OPEN_RELAY", "true"); let app = routes::app(test_config()); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -159,12 +161,26 @@ async fn health_no_auth_required() { assert_eq!(resp.status(), 200); } +#[tokio::test] +async fn metrics_endpoint_requires_auth() { + let base = spawn_test_server().await; + let client = Client::new(); + // No auth header -- should be rejected + let resp = client.get(format!("{base}/metrics")).send().await.unwrap(); + assert_eq!(resp.status(), 401); +} + #[tokio::test] async fn metrics_endpoint_returns_counters() { let base = spawn_test_server().await; let client = Client::new(); - // No auth required for metrics - let resp = client.get(format!("{base}/metrics")).send().await.unwrap(); + // Auth required for metrics + let resp = client + .get(format!("{base}/metrics")) + .header("x-api-key", "test") + .send() + .await + .unwrap(); assert_eq!(resp.status(), 200); let body: serde_json::Value = resp.json().await.unwrap(); // Multi-backend metrics format: { "backends": {...}, "total": {...} } diff --git a/crates/proxy/tests/multi_backend.rs b/crates/proxy/tests/multi_backend.rs index 1e3a72f..e44fe5f 100644 --- a/crates/proxy/tests/multi_backend.rs +++ b/crates/proxy/tests/multi_backend.rs @@ -1,6 +1,6 @@ // Integration tests for multi-backend path-prefix routing. -use anthropic_openai_proxy::config::{self, MultiConfig}; +use anthropic_openai_proxy::config::MultiConfig; use anthropic_openai_proxy::server::routes; use reqwest::Client; @@ -21,6 +21,8 @@ fn test_multi_config() -> MultiConfig { } async fn spawn_multi_server() -> String { + // Enable open-relay mode for tests (no PROXY_API_KEYS configured). + std::env::set_var("PROXY_OPEN_RELAY", "true"); let app = routes::app_multi(test_multi_config()); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -112,7 +114,12 @@ async fn unknown_prefix_returns_404() { async fn metrics_shows_per_backend_breakdown() { let base = spawn_multi_server().await; let client = Client::new(); - let resp = client.get(format!("{base}/metrics")).send().await.unwrap(); + let resp = client + .get(format!("{base}/metrics")) + .header("x-api-key", "any-key") + .send() + .await + .unwrap(); assert_eq!(resp.status(), 200); let body: serde_json::Value = resp.json().await.unwrap(); // Should have per-backend metrics diff --git a/crates/translator/src/mapping/responses_streaming_map.rs b/crates/translator/src/mapping/responses_streaming_map.rs index 0cb396a..b17faeb 100644 --- a/crates/translator/src/mapping/responses_streaming_map.rs +++ b/crates/translator/src/mapping/responses_streaming_map.rs @@ -325,6 +325,7 @@ impl ResponsesStreamingTranslator { } fn handle_error(&mut self, event: &ResponsesStreamEvent) -> Vec { + self.finished = true; // Error is terminal; prevent finish() from emitting closure events let message = event .data .get("response") diff --git a/crates/translator/src/mapping/streaming_map.rs b/crates/translator/src/mapping/streaming_map.rs index 71ad4fa..998f409 100644 --- a/crates/translator/src/mapping/streaming_map.rs +++ b/crates/translator/src/mapping/streaming_map.rs @@ -326,7 +326,7 @@ impl StreamingTranslator { while self.active_tool_calls.len() <= idx { self.active_tool_calls.push(ToolCallAccumulator { block_index: 0, - closed: false, + closed: true, // Padding: never opened, so must not emit ContentBlockStop }); } self.active_tool_calls[idx] = ToolCallAccumulator { diff --git a/crates/translator/src/middleware/handler.rs b/crates/translator/src/middleware/handler.rs index 19e02dd..27d23f3 100644 --- a/crates/translator/src/middleware/handler.rs +++ b/crates/translator/src/middleware/handler.rs @@ -143,9 +143,9 @@ const MAX_SSE_BUFFER_SIZE: usize = 10 * 1024 * 1024; /// Find the first SSE frame boundary (`\n\n` or `\r\n\r\n`) in a byte slice. /// Returns `(position, delimiter_length)` so the caller can skip the full delimiter. -fn find_double_newline(buf: &[u8]) -> Option<(usize, usize)> { +fn find_double_newline(buf: &[u8], start: usize) -> Option<(usize, usize)> { let len = buf.len(); - let mut i = 0; + let mut i = start; while i < len.saturating_sub(1) { if buf[i] == b'\n' && buf[i + 1] == b'\n' { return Some((i, 2)); @@ -178,6 +178,7 @@ where // split across TCP chunk boundaries. let mut buffer = BytesMut::new(); let mut frame_events: Vec = Vec::new(); + let mut search_from: usize = 0; while let Some(chunk_result) = stream.next().await { let bytes = match chunk_result { @@ -197,7 +198,7 @@ where return false; } - while let Some((pos, delim_len)) = find_double_newline(&buffer) { + while let Some((pos, delim_len)) = find_double_newline(&buffer, search_from) { frame_events.clear(); match std::str::from_utf8(&buffer[..pos]) { Ok(frame_str) => { @@ -215,11 +216,16 @@ where } } let _ = buffer.split_to(pos + delim_len); + // split_to shifted the buffer; restart search at the beginning + search_from = 0; if !send_events(tx, &frame_events).await { return false; } } + // Next chunk: resume scanning 3 bytes back from the end so a + // 4-byte delimiter (\r\n\r\n) straddling a chunk boundary is found. + search_from = buffer.len().saturating_sub(3); } true