diff --git a/Cargo.lock b/Cargo.lock index 2c3d31a..4eddf78 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -327,12 +327,6 @@ version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" -[[package]] -name = "byteorder" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" - [[package]] name = "bytes" version = "1.11.1" @@ -1063,16 +1057,6 @@ version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" -[[package]] -name = "hdrhistogram" -version = "7.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "765c9198f173dd59ce26ff9f95ef0aafd0a0fe01fb9d72841bc5066a4c06511d" -dependencies = [ - "byteorder", - "num-traits", -] - [[package]] name = "heck" version = "0.5.0" @@ -1461,7 +1445,7 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "maple-proxy" -version = "0.1.12" +version = "0.2.0" dependencies = [ "anyhow", "async-stream", @@ -1476,7 +1460,6 @@ dependencies = [ "serde", "serde_json", "tokio", - "tokio-test", "tower", "tower-http", "tracing", @@ -1654,12 +1637,13 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" [[package]] name = "opensecret" -version = "3.3.0" +version = "3.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92ce94955ead29ea33002f32d19b94420582e15e0b2c8b7c8e66e64dae6ff48a" +checksum = "a7652ddb1661444ea0cbacacbc62c13f5073f1687a2e765e0195ebf29d3e394d" dependencies = [ "aes-gcm", "anyhow", + "async-stream", "async-trait", "base64", "bytes", @@ -1670,6 +1654,7 @@ dependencies = [ "futures", "hex", "hkdf", + "http", "p256", "percent-encoding", "pin-project", @@ -2556,28 +2541,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "tokio-stream" -version = "0.1.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" -dependencies = [ - "futures-core", - "pin-project-lite", - "tokio", -] - -[[package]] -name = "tokio-test" -version = "0.4.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f6d24790a10a7af737693a3e8f1d03faef7e6ca0cc99aae5066f533766de545" -dependencies = [ - "futures-core", - "tokio", - "tokio-stream", -] - [[package]] name = "tokio-util" version = "0.7.18" @@ -2599,13 +2562,9 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", - "hdrhistogram", - "indexmap", "pin-project-lite", - "slab", "sync_wrapper", "tokio", - "tokio-util", "tower-layer", "tower-service", "tracing", diff --git a/Cargo.toml b/Cargo.toml index 1daf1e7..b282dcd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,14 +1,23 @@ [package] name = "maple-proxy" -version = "0.1.12" +version = "0.2.0" edition = "2021" authors = ["OpenSecret"] description = "Lightweight OpenAI-compatible proxy server for Maple/OpenSecret TEE infrastructure" license = "MIT" -repository = "https://github.com/OpenSecret/maple-proxy" -homepage = "https://github.com/OpenSecret/maple-proxy" +repository = "https://github.com/OpenSecretCloud/maple-proxy" +homepage = "https://github.com/OpenSecretCloud/maple-proxy" keywords = ["openai", "proxy", "tee", "opensecret", "maple"] categories = ["web-programming::http-server", "api-bindings"] +include = [ + "/src/**", + "/examples/**", + "/tests/**", + "/Cargo.toml", + "/Cargo.lock", + "/README.md", + "/LICENSE", +] [lib] name = "maple_proxy" @@ -19,13 +28,13 @@ name = "maple-proxy" path = "src/main.rs" [dependencies] -# OpenSecret SDK -opensecret = "3.3.0" +# OpenSecret SDK +opensecret = "3.4.0" # Web server axum = { version = "0.8.4", features = ["http2", "macros"] } tokio = { version = "1.47", features = ["net", "rt-multi-thread", "macros", "sync", "time"] } -tower = { version = "0.5.2", features = ["full"] } +tower = { version = "0.5.2", features = ["util"] } tower-http = { version = "0.6.6", features = ["cors", "trace"] } # Serialization @@ -51,6 +60,3 @@ http = "1.0" [dev-dependencies] axum-test = "18.0.1" -tokio-test = "0.4" -tower = { version = "0.5.2", features = ["util"] } -serde_json = "1.0" diff --git a/README.md b/README.md index 6082b63..3042a1c 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,17 @@ # 🍁 Maple Proxy -A lightweight OpenAI-compatible proxy server for Maple/OpenSecret's TEE infrastructure. Works with **any** OpenAI client library while providing the security and privacy benefits of Trusted Execution Environment (TEE) processing. +A lightweight proxy for Maple/OpenSecret's OpenAI-compatible inference +endpoints, with the security and privacy benefits of Trusted Execution +Environment (TEE) processing. ## 🚀 Features -- **100% OpenAI Compatible** - Drop-in replacement for OpenAI API +- **OpenAI-Compatible Surface** - Models, chat completions, and embeddings endpoints - **Secure TEE Processing** - All requests processed in secure enclaves -- **Streaming Support** - Full Server-Sent Events streaming for chat completions +- **Lossless Chat Parameters** - Provider-specific request fields pass through unchanged +- **Streaming and Non-Streaming** - Supports both chat completion response modes - **Flexible Authentication** - Environment variables or per-request API keys -- **Zero Client Changes** - Works with existing OpenAI client code +- **Familiar Clients** - Point compatible OpenAI clients at the proxy base URL - **Lightweight** - Minimal overhead, maximum performance - **CORS Support** - Ready for web applications @@ -30,7 +33,7 @@ Add to your `Cargo.toml`: [dependencies] maple-proxy = { git = "https://github.com/opensecretcloud/maple-proxy" } # Or if published to crates.io: -# maple-proxy = "0.1.0" +# maple-proxy = "0.2.0" ``` ## ⚙️ Configuration @@ -40,7 +43,7 @@ Set environment variables or use command-line arguments: ```bash # Environment Variables export MAPLE_HOST=127.0.0.1 # Server host (default: 127.0.0.1) -export MAPLE_PORT=3000 # Server port (default: 3000) +export MAPLE_PORT=8080 # Server port (default: 8080) export MAPLE_BACKEND_URL=http://localhost:3000 # Maple backend URL (prod: https://enclave.trymaple.ai) export MAPLE_API_KEY=your-maple-api-key # Default API key (optional) export MAPLE_DEBUG=true # Enable debug logging @@ -70,7 +73,7 @@ You should see: 📋 Available endpoints: GET /health - Health check GET /v1/models - List available models - POST /v1/chat/completions - Create chat completions (streaming) + POST /v1/chat/completions - Create chat completions (streaming & non-streaming) POST /v1/embeddings - Create embeddings ``` @@ -82,7 +85,7 @@ curl http://localhost:8080/v1/models \ -H "Authorization: Bearer YOUR_MAPLE_API_KEY" ``` -#### Chat Completions (Streaming) +#### Chat Completions ```bash curl -N http://localhost:8080/v1/chat/completions \ -H "Authorization: Bearer YOUR_MAPLE_API_KEY" \ @@ -96,7 +99,9 @@ curl -N http://localhost:8080/v1/chat/completions \ }' ``` -**Note:** Maple currently only supports streaming responses. +Set `stream` to `true` for Server-Sent Events or `false` for one JSON response. +Additional provider-specific JSON fields are forwarded without being parsed or +rewritten by the proxy or Rust SDK. #### Embeddings ```bash diff --git a/src/config.rs b/src/config.rs index 49c1900..f0024d5 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,5 +1,5 @@ use clap::Parser; -use serde::{Deserialize, Serialize}; +use serde::Serialize; use std::{net::SocketAddr, time::Duration}; pub const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 300; @@ -123,22 +123,22 @@ impl Config { } } -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct OpenAIError { - pub error: OpenAIErrorDetails, +#[derive(Debug, Serialize)] +pub(crate) struct OpenAIError { + error: OpenAIErrorDetails, } -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct OpenAIErrorDetails { - pub message: String, +#[derive(Debug, Serialize)] +struct OpenAIErrorDetails { + message: String, #[serde(rename = "type")] - pub error_type: String, - pub param: Option, - pub code: Option, + error_type: String, + param: Option, + code: Option, } impl OpenAIError { - pub fn new(message: impl Into, error_type: impl Into) -> Self { + fn new(message: impl Into, error_type: impl Into) -> Self { Self { error: OpenAIErrorDetails { message: message.into(), @@ -149,11 +149,11 @@ impl OpenAIError { } } - pub fn authentication_error(message: impl Into) -> Self { + pub(crate) fn authentication_error(message: impl Into) -> Self { Self::new(message, "invalid_request_error") } - pub fn server_error(message: impl Into) -> Self { + pub(crate) fn server_error(message: impl Into) -> Self { Self::new(message, "server_error") } } diff --git a/src/lib.rs b/src/lib.rs index d767f1c..41d3f9e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,8 +1,8 @@ -pub mod config; -pub mod proxy; +mod config; +mod proxy; -pub use config::{Config, OpenAIError, OpenAIErrorDetails}; -use proxy::{create_chat_completion, create_embeddings, health_check, list_models, ProxyState}; +pub use config::Config; +use proxy::{health_check, proxy_openai_request, ProxyState}; use axum::{ extract::DefaultBodyLimit, @@ -18,20 +18,23 @@ use tower_http::{ }; use tracing::Level; -pub const MAX_PROXY_REQUEST_BODY_BYTES: usize = 50 * 1024 * 1024; +const MAX_PROXY_REQUEST_BODY_BYTES: usize = 50 * 1024 * 1024; /// Create the Axum application with the given configuration pub fn create_app(config: Config) -> Router { let state = Arc::new(ProxyState::new(config.clone())); + create_app_with_state(config, state) +} +pub(crate) fn create_app_with_state(config: Config, state: Arc) -> Router { let mut app = Router::new() // Health check endpoints .route("/health", get(health_check)) .route("/", get(health_check)) // OpenAI-compatible endpoints - .route("/v1/models", get(list_models)) - .route("/v1/chat/completions", post(create_chat_completion)) - .route("/v1/embeddings", post(create_embeddings)) + .route("/v1/models", get(proxy_openai_request)) + .route("/v1/chat/completions", post(proxy_openai_request)) + .route("/v1/embeddings", post(proxy_openai_request)) .with_state(state) .layer( ServiceBuilder::new() diff --git a/src/main.rs b/src/main.rs index 99d9541..42d624a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -49,7 +49,7 @@ async fn main() -> anyhow::Result<()> { info!( " Set MAPLE_API_KEY environment variable or provide Authorization: Bearer header" ); - info!(" Compatible with any OpenAI client library!"); + info!(" OpenAI-compatible clients can use this proxy as their base URL"); info!(""); info!("🔗 Example curl:"); info!(" curl http://{} \\", config.socket_addr()?); diff --git a/src/proxy.rs b/src/proxy.rs index be38a8a..95378d4 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -1,18 +1,18 @@ use crate::config::{Config, OpenAIError}; use axum::{ - extract::State, - http::{HeaderMap, StatusCode}, - response::{IntoResponse, Response, Sse}, + body::{Body, Bytes}, + extract::{OriginalUri, State}, + http::{header, HeaderMap, HeaderName, Method, Request, StatusCode, Uri}, + response::{IntoResponse, Response}, Json, }; use dashmap::DashMap; -use futures::Stream; -use opensecret::{ - ChatCompletionChunk, ChatCompletionRequest, EmbeddingRequest, EmbeddingResponse, - ModelsResponse, OpenSecretClient, Result as OpenSecretResult, -}; +use futures::{future::BoxFuture, Stream, StreamExt}; +use opensecret::{client::OpenSecretResponseBody, OpenSecretClient, Result as OpenSecretResult}; use std::{ - convert::Infallible, + collections::HashSet, + io, + pin::Pin, sync::Arc, time::{Duration, Instant}, }; @@ -24,6 +24,22 @@ const CLIENT_CACHE_ENTRY_TTL: Duration = Duration::from_secs(60 * 60); type ProxyError = (StatusCode, Json); +trait InferenceTransport: Send + Sync { + fn send_inference_request( + &self, + request: Request, + ) -> BoxFuture<'_, OpenSecretResult>>; +} + +impl InferenceTransport for OpenSecretClient { + fn send_inference_request( + &self, + request: Request, + ) -> BoxFuture<'_, OpenSecretResult>> { + Box::pin(OpenSecretClient::send_inference_request(self, request)) + } +} + struct CachedClientEntry { cell: OnceCell>, created_at: Instant, @@ -43,16 +59,27 @@ impl CachedClientEntry { } #[derive(Clone)] -pub struct ProxyState { - pub config: Config, +pub(crate) struct ProxyState { + config: Config, clients: DashMap>, + transport_override: Option>, } impl ProxyState { - pub fn new(config: Config) -> Self { + pub(crate) fn new(config: Config) -> Self { + Self { + config, + clients: DashMap::new(), + transport_override: None, + } + } + + #[cfg(test)] + fn with_transport(config: Config, transport: Arc) -> Self { Self { config, clients: DashMap::new(), + transport_override: Some(transport), } } @@ -105,6 +132,18 @@ impl ProxyState { } } + async fn transport_for_api_key( + &self, + api_key: &str, + ) -> Result, ProxyError> { + if let Some(transport) = &self.transport_override { + return Ok(Arc::clone(transport)); + } + + let client = self.client_for_api_key(api_key).await?; + Ok(client) + } + fn remove_client_entry_if_same(&self, api_key: &str, client_entry: &Arc) { self.clients .remove_if(api_key, |_, entry| Arc::ptr_eq(entry, client_entry)); @@ -131,7 +170,7 @@ impl ProxyState { } } -pub async fn health_check() -> impl IntoResponse { +pub(crate) async fn health_check() -> impl IntoResponse { Json(serde_json::json!({ "status": "ok", "service": "maple-proxy", @@ -166,30 +205,14 @@ async fn create_client_with_auth( api_key: &str, request_timeout: Duration, ) -> Result { - let client = - OpenSecretClient::new_with_api_key(backend_url, api_key.to_string()).map_err(|e| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(OpenAIError::server_error(format!( - "Failed to create client: {}", - e - ))), - ) - })?; + let client = OpenSecretClient::new_with_api_key(backend_url, api_key.to_string()) + .map_err(|e| transport_error_response("OpenSecret client creation", &e))?; // Perform attestation handshake tokio::time::timeout(request_timeout, client.perform_attestation_handshake()) .await .map_err(|_| timeout_response("Attestation handshake", request_timeout))? - .map_err(|e| { - error!("Attestation handshake failed: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(OpenAIError::server_error( - "Failed to establish secure connection with Maple backend", - )), - ) - })?; + .map_err(|e| transport_error_response("OpenSecret attestation handshake", &e))?; Ok(client) } @@ -210,205 +233,191 @@ fn timeout_response(operation: &str, timeout: Duration) -> ProxyError { ) } -pub async fn list_models( +/// Transparently forwards an OpenAI-compatible inference request through the +/// encrypted OpenSecret transport without interpreting either body. +pub(crate) async fn proxy_openai_request( State(state): State>, + OriginalUri(uri): OriginalUri, + method: Method, headers: HeaderMap, -) -> Result, (StatusCode, Json)> { + body: Bytes, +) -> Result { let api_key = extract_api_key(&headers, &state.config.default_api_key) .map_err(|e| (StatusCode::UNAUTHORIZED, Json(e)))?; debug!( - "Listing models for API key: {}...", + "Proxying {} {} for API key: {}...", + method, + uri, &api_key[..8.min(api_key.len())] ); - let client = state.client_for_api_key(&api_key).await?; - + let transport = state.transport_for_api_key(&api_key).await?; + let request = build_upstream_request(method, uri, &headers, body); let request_timeout = state.config.request_timeout(); - let models = tokio::time::timeout(request_timeout, client.get_models()) + let response = tokio::time::timeout(request_timeout, transport.send_inference_request(request)) .await - .map_err(|_| timeout_response("List models request", request_timeout))? - .map_err(|e| { - error!("Failed to get models: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(OpenAIError::server_error(format!( - "Failed to retrieve models: {}", - e - ))), - ) - })?; + .map_err(|_| timeout_response("OpenAI-compatible request", request_timeout))? + .map_err(|error| transport_error_response("OpenSecret inference request", &error))?; - debug!("Successfully retrieved {} models", models.data.len()); - Ok(Json(models)) + Ok(build_downstream_response( + response, + state.config.stream_idle_timeout(), + )) } -pub async fn create_chat_completion( - State(state): State>, - headers: HeaderMap, - Json(mut request): Json, -) -> Result)> { - let api_key = extract_api_key(&headers, &state.config.default_api_key) - .map_err(|e| (StatusCode::UNAUTHORIZED, Json(e)))?; +fn build_upstream_request( + method: Method, + uri: Uri, + headers: &HeaderMap, + body: Bytes, +) -> Request { + let mut request = Request::new(body); + *request.method_mut() = method; + *request.uri_mut() = uri; + copy_safe_request_headers(headers, request.headers_mut()); + request +} - debug!( - "Chat completion request for model: {}, stream: {:?}", - request.model, - request.stream.unwrap_or(false) - ); +fn copy_safe_request_headers(source: &HeaderMap, destination: &mut HeaderMap) { + let connection_headers = connection_header_names(source); - let client = state.client_for_api_key(&api_key).await?; + for name in source.keys() { + if is_unsafe_request_header(name) || connection_headers.contains(name) { + continue; + } + for value in source.get_all(name) { + destination.append(name.clone(), value.clone()); + } + } +} - // Check if streaming is requested - if request.stream.unwrap_or(false) { - // Handle streaming response - let request_timeout = state.config.request_timeout(); - let stream = tokio::time::timeout( - request_timeout, - client.create_chat_completion_stream(request), - ) - .await - .map_err(|_| timeout_response("Streaming chat completion request", request_timeout))? - .map_err(|e| { - error!("Failed to create streaming chat completion: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(OpenAIError::server_error(format!( - "Failed to create streaming completion: {}", - e - ))), - ) - })?; +fn connection_header_names(headers: &HeaderMap) -> HashSet { + headers + .get_all(header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .filter_map(|name| HeaderName::from_bytes(name.trim().as_bytes()).ok()) + .collect() +} - let sse_stream = create_sse_stream(stream, state.config.stream_idle_timeout()); - Ok(Sse::new(sse_stream).into_response()) - } else { - // Handle non-streaming response - request.stream = Some(false); // Ensure it's explicitly false - - let request_timeout = state.config.request_timeout(); - let response = - tokio::time::timeout(request_timeout, client.create_chat_completion(request)) - .await - .map_err(|_| timeout_response("Chat completion request", request_timeout))? - .map_err(|e| { - error!("Failed to create chat completion: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(OpenAIError::server_error(format!( - "Failed to create completion: {}", - e - ))), - ) - })?; - - debug!("Successfully created chat completion: {}", response.id); - Ok(Json(response).into_response()) - } +fn is_unsafe_request_header(name: &HeaderName) -> bool { + matches!( + name.as_str(), + "authorization" + | "cookie" + | "accept-encoding" + | "proxy-authorization" + | "host" + | "content-length" + | "transfer-encoding" + | "connection" + | "keep-alive" + | "proxy-authenticate" + | "te" + | "trailer" + | "upgrade" + | "x-session-id" + ) +} + +fn transport_error_response(operation: &str, error: &impl std::fmt::Display) -> ProxyError { + error!("{} failed: {}", operation, error); + ( + StatusCode::BAD_GATEWAY, + Json(OpenAIError::server_error( + "Failed to communicate securely with the Maple backend", + )), + ) } -fn create_sse_stream( - mut stream: std::pin::Pin> + Send>>, +fn build_downstream_response( + response: http::Response, stream_idle_timeout: Duration, -) -> impl Stream> { - async_stream::stream! { - use futures::StreamExt; - let mut completed_normally = false; +) -> Response { + let (parts, body) = response.into_parts(); + let mut response = Response::new(Body::from_stream(stream_with_idle_timeout( + body, + stream_idle_timeout, + ))); + *response.status_mut() = parts.status; + copy_safe_response_headers(&parts.headers, response.headers_mut()); + response +} + +fn copy_safe_response_headers(source: &HeaderMap, destination: &mut HeaderMap) { + let connection_headers = connection_header_names(source); + + for name in source.keys() { + if is_unsafe_response_header(name) || connection_headers.contains(name) { + continue; + } + for value in source.get_all(name) { + destination.append(name.clone(), value.clone()); + } + } +} + +fn is_unsafe_response_header(name: &HeaderName) -> bool { + matches!( + name.as_str(), + "set-cookie" + | "content-length" + | "transfer-encoding" + | "connection" + | "keep-alive" + | "proxy-authenticate" + | "proxy-authorization" + | "te" + | "trailer" + | "upgrade" + ) +} +fn stream_with_idle_timeout( + mut stream: OpenSecretResponseBody, + stream_idle_timeout: Duration, +) -> Pin> + Send>> { + Box::pin(async_stream::stream! { loop { let chunk_result = match tokio::time::timeout(stream_idle_timeout, stream.next()).await { Ok(Some(chunk_result)) => chunk_result, - Ok(None) => { - completed_normally = true; - break; - } + Ok(None) => break, Err(_) => { error!( - "Streaming chat completion idle timed out after {} seconds", + "OpenSecret response stream idle timed out after {} seconds", stream_idle_timeout.as_secs() ); - let error_event = axum::response::sse::Event::default() - .data(format!( - r#"{{"error": "Stream idle timeout after {} seconds"}}"#, - stream_idle_timeout.as_secs() - )); - yield Ok(error_event); + yield Err(io::Error::new( + io::ErrorKind::TimedOut, + "Maple backend response stream timed out", + )); break; } }; match chunk_result { - Ok(chunk) => { - match serde_json::to_string(&chunk) { - Ok(json) => { - let event = axum::response::sse::Event::default() - .data(json); - yield Ok(event); - } - Err(e) => { - error!("Failed to serialize chunk: {}", e); - let error_event = axum::response::sse::Event::default() - .data(format!(r#"{{"error": "Failed to serialize chunk: {}"}}"#, e)); - yield Ok(error_event); - break; - } - } - } - Err(e) => { - error!("Stream error: {}", e); - let error_event = axum::response::sse::Event::default() - .data(format!(r#"{{"error": "Stream error: {}"}}"#, e)); - yield Ok(error_event); + Ok(bytes) => yield Ok(bytes), + Err(error) => { + error!("OpenSecret response stream failed: {}", error); + yield Err(io::Error::other("Maple backend response stream failed")); break; } } } - - if completed_normally { - let done_event = axum::response::sse::Event::default() - .data("[DONE]"); - yield Ok(done_event); - } - } -} - -pub async fn create_embeddings( - State(state): State>, - headers: HeaderMap, - Json(request): Json, -) -> Result, (StatusCode, Json)> { - let api_key = extract_api_key(&headers, &state.config.default_api_key) - .map_err(|e| (StatusCode::UNAUTHORIZED, Json(e)))?; - - debug!("Embeddings request for model: {}", request.model); - - let client = state.client_for_api_key(&api_key).await?; - - let request_timeout = state.config.request_timeout(); - let response = tokio::time::timeout(request_timeout, client.create_embeddings(request)) - .await - .map_err(|_| timeout_response("Embeddings request", request_timeout))? - .map_err(|e| { - error!("Failed to create embeddings: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(OpenAIError::server_error(format!( - "Failed to create embeddings: {}", - e - ))), - ) - })?; - - debug!( - "Successfully created embeddings with {} vectors", - response.data.len() - ); - Ok(Json(response)) + }) } #[cfg(test)] mod tests { use super::*; + use axum::{ + body::to_bytes, + http::{HeaderValue, Request as AxumRequest}, + }; + use std::{collections::VecDeque, sync::Mutex}; + use tower::ServiceExt; fn test_config() -> Config { Config { @@ -423,6 +432,77 @@ mod tests { } } + struct MockTransport { + requests: Mutex>>, + responses: Mutex>>>, + } + + impl MockTransport { + fn new(responses: Vec>>) -> Self { + Self { + requests: Mutex::new(Vec::new()), + responses: Mutex::new(responses.into()), + } + } + + fn take_requests(&self) -> Vec> { + std::mem::take(&mut *self.requests.lock().unwrap()) + } + } + + impl InferenceTransport for MockTransport { + fn send_inference_request( + &self, + request: Request, + ) -> BoxFuture<'_, OpenSecretResult>> { + self.requests.lock().unwrap().push(request); + let response = self + .responses + .lock() + .unwrap() + .pop_front() + .expect("a mock response for every request"); + Box::pin(async move { response }) + } + } + + struct PendingTransport; + + impl InferenceTransport for PendingTransport { + fn send_inference_request( + &self, + _request: Request, + ) -> BoxFuture<'_, OpenSecretResult>> { + Box::pin(std::future::pending()) + } + } + + fn raw_response( + status: StatusCode, + headers: &[(&str, &str)], + chunks: Vec, + ) -> http::Response { + let body: OpenSecretResponseBody = Box::pin(futures::stream::iter( + chunks.into_iter().map(Ok::<_, opensecret::Error>), + )); + let mut response = http::Response::new(body); + *response.status_mut() = status; + for (name, value) in headers { + response.headers_mut().append( + HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_str(value).unwrap(), + ); + } + response + } + + fn mock_app(transport: Arc) -> axum::Router { + let mut config = test_config(); + config.default_api_key = Some("default-key".to_string()); + let state = Arc::new(ProxyState::with_transport(config.clone(), transport)); + crate::create_app_with_state(config, state) + } + #[test] fn reuses_client_cell_for_same_api_key() { let state = ProxyState::new(test_config()); @@ -488,24 +568,311 @@ mod tests { } #[tokio::test] - async fn sse_stream_emits_error_when_idle_timeout_expires() { - use futures::{stream, StreamExt}; + async fn all_explicit_inference_routes_forward_method_uri_headers_and_exact_body() { + let responses = (0..3) + .map(|_| { + Ok(raw_response( + StatusCode::OK, + &[], + vec![Bytes::from_static(b"ok")], + )) + }) + .collect(); + let transport = Arc::new(MockTransport::new(responses)); + let app = mock_app(Arc::clone(&transport)); + let chat_body = Bytes::from_static( + br#"{"model":"gemma4-31b","messages":[],"include_reasoning":false,"chat_template_kwargs":{"enable_thinking":false,"future":{"keep":true}}}"#, + ); + let embedding_body = Bytes::from_static(b"\0\xffraw-provider-body"); + + for request in [ + AxumRequest::builder() + .method(Method::GET) + .uri("/v1/models?provider=tinfoil") + .header(header::AUTHORIZATION, "Bearer request-key") + .header("x-provider-beta", "models-v2") + .body(Body::empty()) + .unwrap(), + AxumRequest::builder() + .method(Method::POST) + .uri("/v1/chat/completions?preview=1") + .header(header::AUTHORIZATION, "Bearer request-key") + .header(header::CONTENT_TYPE, "application/json") + .header("x-provider-beta", "thinking-controls") + .body(Body::from(chat_body.clone())) + .unwrap(), + AxumRequest::builder() + .method(Method::POST) + .uri("/v1/embeddings") + .header(header::AUTHORIZATION, "Bearer request-key") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(embedding_body.clone())) + .unwrap(), + ] { + let response = app.clone().oneshot(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(to_bytes(response.into_body(), 16).await.unwrap(), "ok"); + } + + let requests = transport.take_requests(); + assert_eq!(requests.len(), 3); + assert_eq!(requests[0].method(), Method::GET); + assert_eq!(requests[0].uri(), "/v1/models?provider=tinfoil"); + assert!(requests[0].body().is_empty()); + assert_eq!(requests[0].headers()["x-provider-beta"], "models-v2"); + assert_eq!(requests[1].method(), Method::POST); + assert_eq!(requests[1].uri(), "/v1/chat/completions?preview=1"); + assert_eq!(requests[1].body(), &chat_body); + assert_eq!( + requests[1].headers()["x-provider-beta"], + "thinking-controls" + ); + assert_eq!(requests[2].uri(), "/v1/embeddings"); + assert_eq!(requests[2].body(), &embedding_body); + assert!(requests + .iter() + .all(|request| request.headers().get(header::AUTHORIZATION).is_none())); + } + + #[tokio::test] + async fn sse_response_is_forwarded_byte_for_byte_with_one_done_marker() { + let first = Bytes::from_static(b"data: {\"id\":\"one\",\"future\":true}\n\n"); + let done = Bytes::from_static(b"data: [DONE]\n\n"); + let transport = Arc::new(MockTransport::new(vec![Ok(raw_response( + StatusCode::OK, + &[ + ("content-type", "text/event-stream"), + ("x-request-id", "req-sse"), + ], + vec![first.clone(), done.clone()], + ))])); + let response = mock_app(transport) + .oneshot( + AxumRequest::builder() + .method(Method::POST) + .uri("/v1/chat/completions") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(r#"{"stream":true}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.headers()[header::CONTENT_TYPE], + "text/event-stream" + ); + assert_eq!(response.headers()["x-request-id"], "req-sse"); + let body = to_bytes(response.into_body(), 1024).await.unwrap(); + let mut expected = first.to_vec(); + expected.extend_from_slice(&done); + assert_eq!(body.as_ref(), expected); + assert_eq!( + body.windows(b"[DONE]".len()) + .filter(|w| *w == b"[DONE]") + .count(), + 1 + ); + } - let backend_stream = Box::pin(stream::pending::>()); - let sse_stream = create_sse_stream(backend_stream, Duration::from_millis(1)); - futures::pin_mut!(sse_stream); + #[tokio::test] + async fn upstream_error_status_safe_headers_and_body_are_preserved() { + let error_body = Bytes::from_static( + br#"{"error":{"message":"rate limited","type":"rate_limit_error"}}"#, + ); + let transport = Arc::new(MockTransport::new(vec![Ok(raw_response( + StatusCode::TOO_MANY_REQUESTS, + &[ + ("content-type", "application/json"), + ("retry-after", "7"), + ("x-request-id", "req-429"), + ("content-length", "999"), + ("connection", "x-remove"), + ("x-remove", "not-forwarded"), + ], + vec![error_body.clone()], + ))])); + let response = mock_app(transport) + .oneshot( + AxumRequest::builder() + .method(Method::POST) + .uri("/v1/embeddings") + .body(Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); - assert!( - tokio::time::timeout(Duration::from_secs(1), sse_stream.next()) - .await - .unwrap() - .is_some() + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(response.headers()[header::CONTENT_TYPE], "application/json"); + assert_eq!(response.headers()[header::RETRY_AFTER], "7"); + assert_eq!(response.headers()["x-request-id"], "req-429"); + assert!(response.headers().get(header::CONTENT_LENGTH).is_none()); + assert!(response.headers().get("x-remove").is_none()); + assert_eq!( + to_bytes(response.into_body(), 1024).await.unwrap(), + error_body ); - assert!( - tokio::time::timeout(Duration::from_secs(1), sse_stream.next()) - .await - .unwrap() - .is_none() + } + + #[tokio::test] + async fn upstream_server_error_status_and_body_are_preserved() { + let error_body = Bytes::from_static(b"provider unavailable"); + let transport = Arc::new(MockTransport::new(vec![Ok(raw_response( + StatusCode::SERVICE_UNAVAILABLE, + &[("content-type", "text/plain"), ("retry-after", "2")], + vec![error_body.clone()], + ))])); + let response = mock_app(transport) + .oneshot( + AxumRequest::builder() + .method(Method::GET) + .uri("/v1/models") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(response.headers()[header::RETRY_AFTER], "2"); + assert_eq!( + to_bytes(response.into_body(), 1024).await.unwrap(), + error_body + ); + } + + #[tokio::test] + async fn response_start_timeout_is_gateway_timeout() { + let mut config = test_config(); + config.default_api_key = Some("default-key".to_string()); + config.request_timeout_secs = 1; + let state = Arc::new(ProxyState::with_transport( + config.clone(), + Arc::new(PendingTransport), + )); + let response = crate::create_app_with_state(config, state) + .oneshot( + AxumRequest::builder() + .method(Method::GET) + .uri("/v1/models") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::GATEWAY_TIMEOUT); + } + + #[tokio::test] + async fn routes_outside_the_explicit_proxy_surface_are_not_forwarded() { + let transport = Arc::new(MockTransport::new(Vec::new())); + let response = mock_app(Arc::clone(&transport)) + .oneshot( + AxumRequest::builder() + .method(Method::POST) + .uri("/v1/audio/speech") + .body(Body::from("opaque")) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert!(transport.take_requests().is_empty()); + } + + #[tokio::test] + async fn transport_errors_are_safe_bad_gateway_responses() { + let transport = Arc::new(MockTransport::new(vec![Err(opensecret::Error::Other( + "sensitive transport detail".to_string(), + ))])); + let response = mock_app(transport) + .oneshot( + AxumRequest::builder() + .method(Method::GET) + .uri("/v1/models") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + let body = to_bytes(response.into_body(), 1024).await.unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + assert!(body.contains("Failed to communicate securely")); + assert!(!body.contains("sensitive transport detail")); + } + + #[test] + fn request_headers_strip_credentials_framing_and_connection_options() { + let mut source = HeaderMap::new(); + source.insert(header::AUTHORIZATION, "Bearer secret".parse().unwrap()); + source.insert(header::COOKIE, "session=local-secret".parse().unwrap()); + source.insert(header::ACCEPT_ENCODING, "gzip, br".parse().unwrap()); + source.insert(header::HOST, "localhost".parse().unwrap()); + source.insert(header::CONTENT_LENGTH, "10".parse().unwrap()); + source.insert(header::CONNECTION, "x-remove".parse().unwrap()); + source.insert("x-session-id", "forged".parse().unwrap()); + source.insert("x-remove", "also-forbidden".parse().unwrap()); + source.insert("x-provider-beta", "keep-me".parse().unwrap()); + + let mut destination = HeaderMap::new(); + copy_safe_request_headers(&source, &mut destination); + + assert_eq!(destination.get("x-provider-beta").unwrap(), "keep-me"); + assert!(destination.get(header::AUTHORIZATION).is_none()); + assert!(destination.get(header::COOKIE).is_none()); + assert!(destination.get(header::ACCEPT_ENCODING).is_none()); + assert!(destination.get(header::CONTENT_LENGTH).is_none()); + assert!(destination.get("x-session-id").is_none()); + assert!(destination.get("x-remove").is_none()); + } + + #[test] + fn response_headers_strip_cookie_credentials() { + let mut source = HeaderMap::new(); + source.insert( + header::SET_COOKIE, + "session=backend-secret".parse().unwrap(), + ); + source.insert("x-request-id", "req-1".parse().unwrap()); + + let mut destination = HeaderMap::new(); + copy_safe_response_headers(&source, &mut destination); + + assert!(destination.get(header::SET_COOKIE).is_none()); + assert_eq!(destination.get("x-request-id").unwrap(), "req-1"); + } + + #[tokio::test] + async fn response_stream_reports_idle_timeout() { + let backend_stream: OpenSecretResponseBody = + Box::pin(futures::stream::pending::>()); + let mut stream = stream_with_idle_timeout(backend_stream, Duration::from_millis(1)); + + let error = tokio::time::timeout(Duration::from_secs(1), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::TimedOut); + } + + #[test] + fn auth_header_overrides_default_and_default_remains_supported() { + let mut headers = HeaderMap::new(); + headers.insert(header::AUTHORIZATION, "Bearer request-key".parse().unwrap()); + assert_eq!( + extract_api_key(&headers, &Some("default-key".to_string())).unwrap(), + "request-key" + ); + assert_eq!( + extract_api_key(&HeaderMap::new(), &Some("default-key".to_string())).unwrap(), + "default-key" ); } } diff --git a/tests/health_test.rs b/tests/health_test.rs index 819c64c..c1b5f2d 100644 --- a/tests/health_test.rs +++ b/tests/health_test.rs @@ -1,6 +1,6 @@ use axum::http::StatusCode; use axum_test::TestServer; -use maple_proxy::{create_app, Config, MAX_PROXY_REQUEST_BODY_BYTES}; +use maple_proxy::{create_app, Config}; use serde_json::{json, Value}; #[tokio::test] @@ -47,8 +47,6 @@ async fn test_root_health_check() { #[tokio::test] async fn chat_completion_accepts_large_payloads_above_axum_default() { - assert_eq!(MAX_PROXY_REQUEST_BODY_BYTES, 50 * 1024 * 1024); - let config = Config::new( "127.0.0.1".to_string(), 0,