diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 5a217ad..19d9553 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "opensecret" -version = "3.3.0" +version = "3.4.0" edition = "2021" authors = ["OpenSecret"] description = "Rust SDK for OpenSecret - secure AI API interactions with nitro attestation" @@ -14,6 +14,7 @@ categories = ["cryptography", "api-bindings", "web-programming"] [dependencies] # HTTP and async runtime reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls", "http2", "charset", "system-proxy"] } +http = "1.1" tokio = { version = "1.41", features = ["full"] } async-trait = "0.1" @@ -47,6 +48,7 @@ chrono = { version = "0.4", features = ["serde"] } # SSE support for streaming eventsource-stream = "0.2" +async-stream = "0.3" futures = "0.3" pin-project = "1.1" bytes = "1.8" diff --git a/rust/README.md b/rust/README.md index 53471c3..beed83c 100644 --- a/rust/README.md +++ b/rust/README.md @@ -8,6 +8,7 @@ Rust SDK for OpenSecret - secure AI API interactions with nitro attestation. - 🔑 **End-to-End Encryption**: All API calls encrypted with session keys - 👤 **Authentication**: Support for both email-based and guest users - 🔄 **Token Management**: Automatic token refresh and session management +- ↔️ **Lossless Inference Transport**: Provider-specific request and response fields pass through as raw bytes - 🛡️ **Secure by Default**: No plaintext data transmission ## Installation @@ -16,7 +17,10 @@ Add to your `Cargo.toml`: ```toml [dependencies] -opensecret = "0.1.0" +opensecret = "3.4.0" +bytes = "1" +futures = "0.3" +http = "1" ``` ## Quick Start @@ -47,6 +51,47 @@ async fn main() -> Result<()> { } ``` +## Inference APIs + +`send_inference_request` is the lossless inference API. The caller owns the +HTTP method, route, query, headers, body bytes, and response parsing; the SDK +owns attestation, authentication, encryption, retrying an expired session, and +response decryption. It does not parse or rewrite inference parameters such as +`stream`: + +```rust +use bytes::Bytes; +use futures::StreamExt; +use http::Request; + +let request = Request::post("/v1/chat/completions") + .header("x-provider-option", "value") + .body(Bytes::from_static(br#"{ + "model": "provider-model", + "messages": [{"role": "user", "content": "Hello"}], + "provider_option": {"enabled": true} + }"#)) + .expect("valid inference request"); + +let response = client.send_inference_request(request).await?; +let status = response.status(); +let mut body = response.into_body(); +while let Some(chunk) = body.next().await { + let decrypted_bytes = chunk?; + // Parse or forward the bytes in the calling application. +} +``` + +The transport accepts only the SDK's explicitly allowed inference routes; +it cannot be used to bypass JWT-only account or conversation APIs. The existing +typed model, embedding, and chat-completion helpers remain available as +compatibility wrappers over the same transport. + +The SDK manages transport credentials and framing. Caller-provided `Host`, +`Authorization`, `x-session-id`, `Content-Length`, `Content-Type`, +`Content-Encoding`, `Accept-Encoding`, `Content-MD5`, `Digest`, hop-by-hop, and +`Connection`-listed headers are not forwarded; other headers are preserved. + ## Authentication ### User Registration @@ -179,4 +224,4 @@ cargo run --example auth_example ## License -MIT \ No newline at end of file +MIT diff --git a/rust/src/client.rs b/rust/src/client.rs index beca80a..193d021 100644 --- a/rust/src/client.rs +++ b/rust/src/client.rs @@ -7,15 +7,42 @@ use crate::{ types::*, }; use base64::{engine::general_purpose::STANDARD as BASE64, Engine}; +use bytes::{Bytes, BytesMut}; +use futures::{Stream, StreamExt}; +use http::{header, HeaderMap as HttpHeaderMap, Request as HttpRequest, Response as HttpResponse}; use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC}; use reqwest::{ header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE}, Client, }; -use serde::{de::DeserializeOwned, Serialize}; -use std::sync::{Arc, RwLock}; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; +use std::{ + pin::Pin, + sync::{Arc, RwLock}, +}; use uuid::Uuid; +/// A decrypted response body returned by [`OpenSecretClient::send_inference_request`]. +/// +/// Ordinary responses contain one chunk. Server-sent event responses remain a +/// stream, with each encrypted `data:` field decrypted without interpreting its +/// payload as JSON. +pub type OpenSecretResponseBody = Pin> + Send + 'static>>; + +/// A caller-owned request to an allowed OpenSecret inference endpoint. +pub type InferenceRequest = HttpRequest; + +/// A decrypted HTTP response from an OpenSecret inference endpoint. +pub type InferenceResponse = HttpResponse; + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct EncryptedBody { + encrypted: String, +} + +const MAX_INFERENCE_SSE_LINE_BYTES: usize = 16 * 1024 * 1024; + pub struct OpenSecretClient { client: Client, base_url: String, @@ -149,6 +176,199 @@ enum AuthHeaderMode { ApiKeyOrJwt, } +fn is_allowed_inference_endpoint(method: &http::Method, path: &str) -> bool { + matches!( + (method.as_str(), path), + ("GET", "/v1/models") + | ("GET", "/v1/models/catalog") + | ("POST", "/v1/chat/completions") + | ("POST", "/v1/embeddings") + | ("POST", "/v1/audio/speech") + | ("POST", "/v1/audio/transcriptions") + ) +} + +fn is_hop_by_hop_header(name: &http::HeaderName) -> bool { + matches!( + name.as_str(), + "connection" + | "keep-alive" + | "proxy-authenticate" + | "proxy-authorization" + | "te" + | "trailer" + | "transfer-encoding" + | "upgrade" + ) +} + +fn sanitize_inference_request_headers(headers: &HttpHeaderMap) -> HttpHeaderMap { + let connection_headers = headers + .get_all(header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .map(str::trim) + .filter(|value| !value.is_empty()) + .filter_map(|value| http::HeaderName::from_bytes(value.as_bytes()).ok()) + .collect::>(); + + headers + .iter() + .filter(|(name, _)| { + !is_hop_by_hop_header(name) + && !connection_headers.contains(name) + && *name != header::HOST + && *name != header::AUTHORIZATION + && name.as_str() != "x-session-id" + && *name != header::CONTENT_LENGTH + && *name != header::CONTENT_TYPE + && *name != header::CONTENT_ENCODING + && *name != header::ACCEPT_ENCODING + && name.as_str() != "content-md5" + && name.as_str() != "digest" + }) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +fn sanitize_inference_response_headers(headers: &HttpHeaderMap) -> HttpHeaderMap { + let connection_headers = headers + .get_all(header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .map(str::trim) + .filter(|value| !value.is_empty()) + .filter_map(|value| http::HeaderName::from_bytes(value.as_bytes()).ok()) + .collect::>(); + + headers + .iter() + .filter(|(name, _)| { + !is_hop_by_hop_header(name) + && !connection_headers.contains(name) + && *name != header::CONTENT_LENGTH + && *name != header::CONTENT_ENCODING + && name.as_str() != "content-md5" + && name.as_str() != "digest" + && *name != header::ETAG + }) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +fn is_event_stream(headers: &HttpHeaderMap) -> bool { + headers + .get(header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.split(';').next()) + .is_some_and(|media_type| media_type.trim().eq_ignore_ascii_case("text/event-stream")) +} + +fn transform_sse_line(line: Bytes, session_key: &[u8; 32]) -> Result { + let (content, line_ending) = if line.ends_with(b"\r\n") { + (&line[..line.len() - 2], &b"\r\n"[..]) + } else if line.ends_with(b"\n") { + (&line[..line.len() - 1], &b"\n"[..]) + } else { + (&line[..], &b""[..]) + }; + + let Some(mut payload) = content.strip_prefix(b"data:") else { + return Ok(line); + }; + let prefix_len = if payload.starts_with(b" ") { + payload = &payload[1..]; + 6 + } else { + 5 + }; + + if payload == b"[DONE]" { + return Ok(line); + } + + if payload.is_empty() { + return Ok(line); + } + + // OpenSecret encrypts every normal data event. A base64-shaped payload + // that fails authentication is corrupt transport data, not plaintext. + let encrypted = match BASE64.decode(payload) { + Ok(encrypted) => encrypted, + Err(_) => return Ok(line), + }; + if encrypted.len() < 28 { + return Ok(line); + } + let decrypted = crypto::decrypt_data(session_key, &encrypted) + .map_err(|error| Error::Decryption(format!("Failed to decrypt SSE data: {error}")))?; + + let mut transformed = BytesMut::with_capacity(prefix_len + decrypted.len() + line_ending.len()); + transformed.extend_from_slice(&content[..prefix_len]); + transformed.extend_from_slice(&decrypted); + transformed.extend_from_slice(line_ending); + Ok(transformed.freeze()) +} + +fn decrypt_sse_stream( + source: OpenSecretResponseBody, + session_key: [u8; 32], +) -> OpenSecretResponseBody { + decrypt_sse_stream_with_line_limit(source, session_key, MAX_INFERENCE_SSE_LINE_BYTES) +} + +fn decrypt_sse_stream_with_line_limit( + mut source: OpenSecretResponseBody, + session_key: [u8; 32], + max_line_bytes: usize, +) -> OpenSecretResponseBody { + let stream = async_stream::try_stream! { + let mut buffered = BytesMut::new(); + + while let Some(chunk) = source.next().await { + buffered.extend_from_slice(&chunk?); + while let Some(line_end) = buffered.iter().position(|byte| *byte == b'\n') { + let line_len = line_end + 1; + if line_len > max_line_bytes { + Err(Error::InvalidResponse(format!( + "Inference SSE line exceeds {max_line_bytes}-byte limit" + )))?; + } + let line = buffered.split_to(line_len).freeze(); + yield transform_sse_line(line, &session_key)?; + } + if buffered.len() > max_line_bytes { + Err(Error::InvalidResponse(format!( + "Inference SSE line exceeds {max_line_bytes}-byte limit" + )))?; + } + } + + if !buffered.is_empty() { + yield transform_sse_line(buffered.freeze(), &session_key)?; + } + }; + + Box::pin(stream) +} + +fn decrypt_sse_body(response: reqwest::Response, session_key: [u8; 32]) -> OpenSecretResponseBody { + let source = response + .bytes_stream() + .map(|chunk| chunk.map_err(Error::Http)); + decrypt_sse_stream(Box::pin(source), session_key) +} + +async fn collect_response_body(mut body: OpenSecretResponseBody) -> Result { + let mut collected = BytesMut::new(); + while let Some(chunk) = body.next().await { + collected.extend_from_slice(&chunk?); + } + Ok(collected.freeze()) +} + impl OpenSecretClient { pub fn new(base_url: impl Into) -> Result { let base_url = base_url.into(); @@ -508,16 +728,210 @@ impl OpenSecretClient { Ok(result) } - /// Encrypted API call specifically for OpenAI endpoints (/v1/*) - /// This supports both API key and JWT authentication, with API key taking priority - async fn encrypted_openai_call( + /// Sends a lossless encrypted request to an OpenSecret inference endpoint. + /// + /// The request URI must be relative and target one of the SDK's explicitly + /// allowed inference routes. Its method, query string, headers, and body + /// bytes are otherwise caller-owned. The SDK does not parse the request as + /// JSON and never adds or changes inference parameters such as `stream`. + /// + /// OpenSecret authentication, attestation sessions, and the encrypted + /// envelope remain SDK-owned. Caller-provided `Host`, `Authorization`, + /// `x-session-id`, `Content-Length`, `Content-Type`, `Content-Encoding`, + /// `Accept-Encoding`, `Content-MD5`, `Digest`, and hop-by-hop headers + /// (including fields named by `Connection`) are therefore not forwarded. + /// Other headers are preserved. + /// + /// The returned HTTP response preserves the final OpenSecret status and + /// safe response headers. Its body is decrypted raw bytes. SSE framing is + /// preserved while encrypted `data:` fields are decrypted without parsing + /// their contents as completion JSON. Individual SSE lines are limited to + /// 16 MiB. + pub async fn send_inference_request( + &self, + request: InferenceRequest, + ) -> Result { + let (parts, body) = request.into_parts(); + if parts.uri.scheme().is_some() || parts.uri.authority().is_some() { + return Err(Error::Configuration( + "Inference request URI must be relative".to_string(), + )); + } + if !is_allowed_inference_endpoint(&parts.method, parts.uri.path()) { + return Err(Error::Configuration(format!( + "Inference endpoint is not allowed: {} {}", + parts.method, + parts.uri.path() + ))); + } + + let path_and_query = parts + .uri + .path_and_query() + .ok_or_else(|| Error::Configuration("Inference request URI has no path".to_string()))? + .as_str() + .to_string(); + let headers = sanitize_inference_request_headers(&parts.headers); + let mut retried_attestation = false; + let mut retried_refresh = false; + + loop { + let result = self + .send_inference_request_once(&parts.method, &path_and_query, &headers, body.clone()) + .await; + + match result { + // OpenSecret currently uses the same 400 for stale sessions and + // pre-provider validation. Preserve the legacy one-retry behavior + // until the backend exposes an explicit re-attestation signal. + Ok((response, _session_key)) + if response.status() == reqwest::StatusCode::BAD_REQUEST + && !retried_attestation => + { + self.perform_attestation_handshake().await?; + retried_attestation = true; + } + Ok((response, session_key)) + if response.status() == reqwest::StatusCode::UNAUTHORIZED + && !retried_refresh + && !self.using_api_key(AuthHeaderMode::ApiKeyOrJwt)? => + { + if self.refresh_token().await.is_ok() { + retried_refresh = true; + } else { + return self.finish_inference_response(response, session_key).await; + } + } + Ok((response, session_key)) => { + return self.finish_inference_response(response, session_key).await + } + Err(error) if !retried_attestation && Self::is_attestation_retryable(&error) => { + self.perform_attestation_handshake().await?; + retried_attestation = true; + } + Err(error) => return Err(error), + } + } + } + + async fn send_inference_request_once( + &self, + method: &http::Method, + path_and_query: &str, + caller_headers: &HttpHeaderMap, + body: Bytes, + ) -> Result<(reqwest::Response, [u8; 32])> { + let session = self.session_manager.get_session()?.ok_or_else(|| { + Error::Session( + "No active session. Call perform_attestation_handshake first".to_string(), + ) + })?; + let url = format!("{}{}", self.base_url, path_and_query); + let mut headers = caller_headers.clone(); + headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + headers.insert( + "x-session-id", + HeaderValue::from_str(&session.session_id.to_string()) + .map_err(|error| Error::Session(format!("Invalid session ID: {error}")))?, + ); + if let Some(token) = self.resolve_auth_token(AuthHeaderMode::ApiKeyOrJwt)? { + headers.insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {token}")).map_err(|error| { + Error::Authentication(format!( + "Invalid authorization credential format: {error}" + )) + })?, + ); + } + + let encrypted_body = if body.is_empty() { + None + } else { + let encrypted = crypto::encrypt_data(&session.session_key, &body)?; + Some(EncryptedRequest { + encrypted: BASE64.encode(encrypted), + }) + }; + let request = self.client.request(method.clone(), url).headers(headers); + let response = match encrypted_body { + Some(encrypted_body) => request.json(&encrypted_body).send().await?, + None => request.send().await?, + }; + + Ok((response, session.session_key)) + } + + async fn finish_inference_response( + &self, + response: reqwest::Response, + session_key: [u8; 32], + ) -> Result> { + let status = response.status(); + let headers = sanitize_inference_response_headers(response.headers()); + let body: OpenSecretResponseBody = if is_event_stream(&headers) { + decrypt_sse_body(response, session_key) + } else { + let raw_body = response.bytes().await?; + let decrypted_body = match serde_json::from_slice::(&raw_body) { + Ok(encrypted_body) => { + let encrypted = BASE64.decode(encrypted_body.encrypted)?; + Bytes::from(crypto::decrypt_data(&session_key, &encrypted)?) + } + Err(_) if status.is_success() && !raw_body.is_empty() => { + return Err(Error::InvalidResponse( + "Successful inference response did not contain an encrypted body" + .to_string(), + )); + } + Err(_) => raw_body, + }; + Box::pin(futures::stream::once(async move { Ok(decrypted_body) })) + }; + + let mut result = HttpResponse::builder() + .status(status) + .body(body) + .map_err(|error| { + Error::InvalidResponse(format!("Failed to construct inference response: {error}")) + })?; + *result.headers_mut() = headers; + Ok(result) + } + + /// Typed compatibility wrapper over the lossless inference transport. + async fn encrypted_openai_call( &self, endpoint: &str, method: &str, data: Option, ) -> Result { - self.retry_encrypted_json_call(endpoint, method, data, AuthHeaderMode::ApiKeyOrJwt, true) - .await + let method = http::Method::from_bytes(method.as_bytes()).map_err(|error| { + Error::Configuration(format!("Invalid inference HTTP method: {error}")) + })?; + let body = match data { + Some(data) => Bytes::from(serde_json::to_vec(&data)?), + None => Bytes::new(), + }; + let request = HttpRequest::builder() + .method(method) + .uri(endpoint) + .body(body) + .map_err(|error| { + Error::Configuration(format!("Failed to build inference request: {error}")) + })?; + let response = self.send_inference_request(request).await?; + let status = response.status(); + let body = collect_response_body(response.into_body()).await?; + + if !status.is_success() { + return Err(Error::Api { + status: status.as_u16(), + message: String::from_utf8_lossy(&body).into_owned(), + }); + } + + Ok(serde_json::from_slice(&body)?) } async fn retry_encrypted_stream_call( @@ -1528,45 +1942,49 @@ impl OpenSecretClient { /// Creates a chat completion (non-streaming) pub async fn create_chat_completion( &self, - request: ChatCompletionRequest, + mut request: ChatCompletionRequest, ) -> Result { - let mut modified_request = request; - modified_request.stream = Some(false); - self.encrypted_openai_call("/v1/chat/completions", "POST", Some(modified_request)) + request.stream = Some(false); + self.encrypted_openai_call("/v1/chat/completions", "POST", Some(request)) .await } /// Creates a streaming chat completion pub async fn create_chat_completion_stream( &self, - request: ChatCompletionRequest, + mut request: ChatCompletionRequest, ) -> Result> + Send>>> { - use eventsource_stream::Eventsource; - use futures::StreamExt; - - let mut modified_request = request; - modified_request.stream = Some(true); - modified_request.stream_options = Some(StreamOptions { + request.stream = Some(true); + request.stream_options = Some(StreamOptions { include_usage: true, }); + use eventsource_stream::Eventsource; + use futures::StreamExt; - let (response, session_key) = self - .retry_encrypted_stream_call( - "/v1/chat/completions", - "POST", - Some(modified_request), - AuthHeaderMode::ApiKeyOrJwt, - true, - ) - .await?; + let request = HttpRequest::builder() + .method(http::Method::POST) + .uri("/v1/chat/completions") + .header(header::ACCEPT, "text/event-stream") + .body(Bytes::from(serde_json::to_vec(&request)?)) + .map_err(|error| { + Error::Configuration(format!("Failed to build inference request: {error}")) + })?; + let response = self.send_inference_request(request).await?; + let status = response.status(); + if !status.is_success() { + let body = collect_response_body(response.into_body()).await?; + return Err(Error::Api { + status: status.as_u16(), + message: String::from_utf8_lossy(&body).into_owned(), + }); + } let stream = response - .bytes_stream() + .into_body() .map(|result| result.map_err(std::io::Error::other)); let event_stream = stream.eventsource().filter_map(move |event| { - let session_key = session_key; async move { match event { Ok(event) => { @@ -1575,37 +1993,17 @@ impl OpenSecretClient { return None; } - // Decrypt the event data - server sends base64 encrypted chunks. - // Skip non-base64 events (heartbeats, retries, etc.) to match TS SDK. - let encrypted_bytes = match BASE64.decode(&event.data) { - Ok(bytes) => bytes, - Err(_) => return None, - }; - match crypto::decrypt_data(&session_key, &encrypted_bytes) { - Ok(decrypted) => match String::from_utf8(decrypted) { - Ok(json_str) => { - match serde_json::from_str::(&json_str) { - Ok(chunk) => Some(Ok(chunk)), - Err(e) => Some(Err(Error::Api { - status: 0, - message: format!("Failed to parse chunk: {}", e), - })), - } - } - Err(e) => Some(Err(Error::Api { - status: 0, - message: format!("Invalid UTF-8 in decrypted data: {}", e), - })), - }, - Err(e) => Some(Err(Error::Decryption(format!( - "Failed to decrypt chunk: {}", - e - )))), + match serde_json::from_str::(&event.data) { + Ok(chunk) => Some(Ok(chunk)), + Err(error) => Some(Err(Error::Api { + status: 0, + message: format!("Failed to parse chunk: {error}"), + })), } } - Err(e) => Some(Err(Error::Api { + Err(error) => Some(Err(Error::Api { status: 0, - message: format!("SSE error: {}", e), + message: format!("SSE error: {error}"), })), } } @@ -1934,7 +2332,7 @@ mod tests { use futures::StreamExt; use serde_json::json; use wiremock::{ - matchers::{header, method, path}, + matchers::{header, method, path, query_param}, Match, Mock, MockServer, Request, Respond, ResponseTemplate, }; @@ -1954,6 +2352,53 @@ mod tests { } } + #[derive(Debug)] + struct EncryptedJsonBodyMatcher { + session_key: [u8; 32], + expected: serde_json::Value, + } + + impl Match for EncryptedJsonBodyMatcher { + fn matches(&self, request: &Request) -> bool { + let Ok(body) = serde_json::from_slice::(request.body.as_ref()) else { + return false; + }; + let Ok(encrypted) = BASE64.decode(body.encrypted.as_bytes()) else { + return false; + }; + let Ok(plaintext) = crypto::decrypt_data(&self.session_key, &encrypted) else { + return false; + }; + let Ok(actual) = serde_json::from_slice::(&plaintext) else { + return false; + }; + + actual == self.expected + } + } + + #[derive(Debug)] + struct EncryptedBytesBodyMatcher { + session_key: [u8; 32], + expected: Bytes, + } + + impl Match for EncryptedBytesBodyMatcher { + fn matches(&self, request: &Request) -> bool { + let Ok(body) = serde_json::from_slice::(request.body.as_ref()) else { + return false; + }; + let Ok(encrypted) = BASE64.decode(body.encrypted.as_bytes()) else { + return false; + }; + let Ok(actual) = crypto::decrypt_data(&self.session_key, &encrypted) else { + return false; + }; + + actual == self.expected + } + } + struct AttestationResponder { server_public_key: [u8; 32], } @@ -2024,12 +2469,22 @@ mod tests { json!({ "encrypted": BASE64.encode(encrypted) }) } + fn encrypted_response_bytes(session_key: &[u8; 32], payload: &[u8]) -> serde_json::Value { + let encrypted = crypto::encrypt_data(session_key, payload).unwrap(); + json!({ "encrypted": BASE64.encode(encrypted) }) + } + fn encrypted_sse_data(session_key: &[u8; 32], payload: &T) -> String { let plaintext = serde_json::to_vec(payload).unwrap(); let encrypted = crypto::encrypt_data(session_key, &plaintext).unwrap(); format!("data: {}\n\n", BASE64.encode(encrypted)) } + fn encrypted_sse_bytes(session_key: &[u8; 32], payload: &[u8]) -> String { + let encrypted = crypto::encrypt_data(session_key, payload).unwrap(); + BASE64.encode(encrypted) + } + fn decrypt_request_body( request: &Request, session_key: &[u8; 32], @@ -2667,11 +3122,18 @@ mod tests { .and(path("/v1/chat/completions")) .and(header("authorization", "Bearer access_token")) .and(header("x-session-id", session_id.to_string())) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ + "model": "kimi-k2-5", + "messages": [{"role": "user", "content": "What is 2+2?"}], + "temperature": 0.0, + "max_tokens": 100, + "stream": true, + "stream_options": {"include_usage": true} + }), + }) + .respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream")) .expect(1) .mount(&mock_server) .await; @@ -2702,6 +3164,531 @@ mod tests { assert!(stream.next().await.is_none()); } + #[tokio::test] + async fn inference_transport_preserves_raw_request_and_response_bytes() { + let mock_server = MockServer::start().await; + let client = + OpenSecretClient::new_with_api_key(mock_server.uri(), "real_api_key".to_string()) + .unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [27u8; 32]; + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + + let request_body = Bytes::from_static( + br#"{ "z": 184467440737095516160000000000000, "stream":false, "a":[1, 2] } +"#, + ); + let response_body = Bytes::from_static( + br#"{ "provider": {"huge":184467440737095516160000000000001}, "a": 1 } +"#, + ); + + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(query_param("trace", "one two")) + .and(header("authorization", "Bearer real_api_key")) + .and(header("x-session-id", session_id.to_string())) + .and(header("content-type", "application/json")) + .and(header("x-provider-beta", "raw-v2")) + .and(MissingHeaderMatcher("x-remove")) + .and(MissingHeaderMatcher("content-encoding")) + .and(MissingHeaderMatcher("accept-encoding")) + .and(MissingHeaderMatcher("content-md5")) + .and(EncryptedBytesBodyMatcher { + session_key, + expected: request_body.clone(), + }) + .respond_with( + ResponseTemplate::new(201) + .insert_header("content-type", "application/json") + .insert_header("x-provider-result", "kept") + .insert_header("connection", "x-remove-response") + .insert_header("x-remove-response", "gone") + .insert_header("content-encoding", "identity") + .set_body_json(encrypted_response_bytes(&session_key, &response_body)), + ) + .expect(1) + .mount(&mock_server) + .await; + + let request = HttpRequest::builder() + .method(http::Method::POST) + .uri("/v1/chat/completions?trace=one%20two") + .header(header::AUTHORIZATION, "Bearer caller_must_not_control_this") + .header("x-session-id", "caller_must_not_control_this") + .header(header::CONTENT_TYPE, "application/custom") + .header(header::CONTENT_ENCODING, "gzip") + .header(header::ACCEPT_ENCODING, "gzip, br") + .header("content-md5", "caller-body-digest") + .header(header::CONNECTION, "x-remove") + .header("x-remove", "gone") + .header("x-provider-beta", "raw-v2") + .body(request_body) + .unwrap(); + let response = client.send_inference_request(request).await.unwrap(); + + assert_eq!(response.status(), http::StatusCode::CREATED); + assert_eq!(response.headers().get("x-provider-result").unwrap(), "kept"); + assert!(!response.headers().contains_key(header::CONTENT_LENGTH)); + assert!(!response.headers().contains_key(header::CONTENT_ENCODING)); + assert!(!response.headers().contains_key(header::CONNECTION)); + assert!(!response.headers().contains_key("x-remove-response")); + assert_eq!( + collect_response_body(response.into_body()).await.unwrap(), + response_body + ); + } + + #[tokio::test] + async fn inference_transport_rejects_absolute_and_non_inference_routes() { + let client = OpenSecretClient::new("http://localhost:3000").unwrap(); + for (method, uri) in [ + ( + http::Method::POST, + "https://example.test/v1/chat/completions", + ), + (http::Method::POST, "/v1/responses"), + (http::Method::GET, "/v1/chat/completions"), + (http::Method::POST, "/protected/user"), + ] { + let request = HttpRequest::builder() + .method(method) + .uri(uri) + .body(Bytes::new()) + .unwrap(); + assert!(matches!( + client.send_inference_request(request).await, + Err(Error::Configuration(_)) + )); + } + } + + #[tokio::test] + async fn inference_transport_preserves_final_plaintext_error_after_refresh() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [28u8; 32]; + let request_body = Bytes::from_static(br#"{"model":"test","messages":[]}"#); + let error_body = + Bytes::from_static(br#"{"error":{"type":"rate_limit","n":999999999999999999999}}"#); + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + client + .session_manager + .set_tokens( + "expired_access".to_string(), + Some("refresh_token".to_string()), + ) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(header("authorization", "Bearer expired_access")) + .and(EncryptedBytesBodyMatcher { + session_key, + expected: request_body.clone(), + }) + .respond_with(ResponseTemplate::new(401).set_body_string("jwt expired")) + .expect(1) + .mount(&mock_server) + .await; + Mock::given(method("POST")) + .and(path("/refresh")) + .and(MissingHeaderMatcher("authorization")) + .respond_with(ResponseTemplate::new(200).set_body_json(encrypted_response( + &session_key, + &json!({ + "access_token": "fresh_access", + "refresh_token": "fresh_refresh" + }), + ))) + .expect(1) + .mount(&mock_server) + .await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(header("authorization", "Bearer fresh_access")) + .and(EncryptedBytesBodyMatcher { + session_key, + expected: request_body.clone(), + }) + .respond_with( + ResponseTemplate::new(429) + .insert_header("x-request-id", "provider-429") + .set_body_bytes(error_body.clone()), + ) + .expect(1) + .mount(&mock_server) + .await; + + let request = HttpRequest::builder() + .method(http::Method::POST) + .uri("/v1/chat/completions") + .body(request_body) + .unwrap(); + let response = client.send_inference_request(request).await.unwrap(); + + assert_eq!(response.status(), http::StatusCode::TOO_MANY_REQUESTS); + assert_eq!( + response.headers().get("x-request-id").unwrap(), + "provider-429" + ); + assert_eq!( + collect_response_body(response.into_body()).await.unwrap(), + error_body + ); + } + + #[tokio::test] + async fn inference_transport_decrypts_non_success_encrypted_body() { + let mock_server = MockServer::start().await; + let client = + OpenSecretClient::new_with_api_key(mock_server.uri(), "api_key".to_string()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [29u8; 32]; + let error_body = Bytes::from_static(br#"{ "error": "provider rejected", "extra": 7 }"#); + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with( + ResponseTemplate::new(422) + .set_body_json(encrypted_response_bytes(&session_key, &error_body)), + ) + .expect(1) + .mount(&mock_server) + .await; + + let request = HttpRequest::builder() + .method(http::Method::POST) + .uri("/v1/embeddings") + .body(Bytes::from_static(br#"{"model":"x","input":"y"}"#)) + .unwrap(); + let response = client.send_inference_request(request).await.unwrap(); + + assert_eq!(response.status(), http::StatusCode::UNPROCESSABLE_ENTITY); + assert_eq!( + collect_response_body(response.into_body()).await.unwrap(), + error_body + ); + } + + #[tokio::test] + async fn inference_transport_preserves_non_envelope_error_with_encrypted_field() { + let mock_server = MockServer::start().await; + let client = + OpenSecretClient::new_with_api_key(mock_server.uri(), "api_key".to_string()).unwrap(); + let session_key = [36u8; 32]; + let error_body = Bytes::from( + serde_json::to_vec(&json!({ + "encrypted": BASE64.encode([0u8; 28]), + "message": "plain backend error" + })) + .unwrap(), + ); + client + .session_manager + .set_session(Uuid::new_v4(), session_key) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(ResponseTemplate::new(422).set_body_bytes(error_body.clone())) + .expect(1) + .mount(&mock_server) + .await; + + let request = HttpRequest::builder() + .method(http::Method::POST) + .uri("/v1/embeddings") + .body(Bytes::from_static(br#"{"model":"x","input":"y"}"#)) + .unwrap(); + let response = client.send_inference_request(request).await.unwrap(); + + assert_eq!(response.status(), http::StatusCode::UNPROCESSABLE_ENTITY); + assert_eq!( + collect_response_body(response.into_body()).await.unwrap(), + error_body + ); + } + + #[tokio::test] + async fn inference_transport_establishes_attestation_and_replays_exact_bytes() { + let mock_server = MockServer::start().await; + let client = + OpenSecretClient::new_with_api_key(mock_server.uri(), "api_key".to_string()).unwrap(); + let server_secret_key = [34u8; 32]; + let server_public_key = + x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from(server_secret_key)); + let session_key = [35u8; 32]; + let session_id = Uuid::new_v4(); + let request_body = Bytes::from_static(br#"{ "model":"x", "input":"exact bytes" }"#); + let response_body = Bytes::from_static(br#"{ "ok": true }"#); + + Mock::given(method("GET")) + .and(PathPrefixMatcher("/attestation/")) + .respond_with(AttestationResponder { + server_public_key: server_public_key.to_bytes(), + }) + .expect(1) + .mount(&mock_server) + .await; + Mock::given(method("POST")) + .and(path("/key_exchange")) + .and(MissingHeaderMatcher("authorization")) + .respond_with(KeyExchangeResponder { + server_secret_key, + session_key, + session_id: session_id.to_string(), + }) + .expect(1) + .mount(&mock_server) + .await; + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .and(header("authorization", "Bearer api_key")) + .and(header("x-session-id", session_id.to_string())) + .and(EncryptedBytesBodyMatcher { + session_key, + expected: request_body.clone(), + }) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(encrypted_response_bytes(&session_key, &response_body)), + ) + .expect(1) + .mount(&mock_server) + .await; + + let request = HttpRequest::builder() + .method(http::Method::POST) + .uri("/v1/embeddings") + .body(request_body) + .unwrap(); + let response = client.send_inference_request(request).await.unwrap(); + + assert_eq!( + collect_response_body(response.into_body()).await.unwrap(), + response_body + ); + } + + #[tokio::test] + async fn inference_sse_transport_preserves_framing_across_arbitrary_chunks() { + let session_key = [30u8; 32]; + let decrypted_payload = + br#"{ "delta": {"huge":184467440737095516160000000000000}, "text":"hi" }"#; + let encrypted_payload = encrypted_sse_bytes(&session_key, decrypted_payload); + let encrypted_sse = format!( + ": heartbeat\r\nevent: chunk\r\nid: provider-7\r\nretry: 1500\r\ndata: {encrypted_payload}\r\n\r\ndata: provider-heartbeat\n\ndata: [DONE]\n\n" + ); + let expected = format!( + ": heartbeat\r\nevent: chunk\r\nid: provider-7\r\nretry: 1500\r\ndata: {}\r\n\r\ndata: provider-heartbeat\n\ndata: [DONE]\n\n", + String::from_utf8_lossy(decrypted_payload) + ); + let chunks = encrypted_sse + .as_bytes() + .chunks(3) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect::>>(); + let source: OpenSecretResponseBody = Box::pin(futures::stream::iter(chunks)); + + let actual = collect_response_body(decrypt_sse_stream(source, session_key)) + .await + .unwrap(); + + assert_eq!(actual, expected.as_bytes()); + assert_eq!( + actual + .windows(6) + .filter(|window| *window == b"[DONE]") + .count(), + 1 + ); + } + + #[tokio::test] + async fn inference_sse_transport_reports_corrupt_ciphertext() { + let session_key = [31u8; 32]; + let mut encrypted = crypto::encrypt_data(&session_key, br#"{"delta":"x"}"#).unwrap(); + *encrypted.last_mut().unwrap() ^= 0xff; + let source: OpenSecretResponseBody = Box::pin(futures::stream::iter([Ok(Bytes::from( + format!("data: {}\n\n", BASE64.encode(encrypted)), + ))])); + + let error = collect_response_body(decrypt_sse_stream(source, session_key)) + .await + .unwrap_err(); + + assert!(matches!(error, Error::Decryption(message) if message.contains("SSE data"))); + } + + #[tokio::test] + async fn inference_sse_transport_bounds_each_line_not_each_network_chunk() { + let session_key = [37u8; 32]; + let within_limit = Bytes::from_static(b":1\n:2\n:3\n"); + let source: OpenSecretResponseBody = + Box::pin(futures::stream::iter([Ok(within_limit.clone())])); + let actual = + collect_response_body(decrypt_sse_stream_with_line_limit(source, session_key, 8)) + .await + .unwrap(); + assert_eq!(actual, within_limit); + + let exact_limit = Bytes::from_static(b":1234567"); + let source: OpenSecretResponseBody = + Box::pin(futures::stream::iter([Ok(exact_limit.clone())])); + let actual = + collect_response_body(decrypt_sse_stream_with_line_limit(source, session_key, 8)) + .await + .unwrap(); + assert_eq!(actual, exact_limit); + + for chunks in [ + vec![ + Ok(Bytes::from_static(b":1234")), + Ok(Bytes::from_static(b"5678")), + ], + vec![Ok(Bytes::from_static(b":12345678\n"))], + ] { + let source: OpenSecretResponseBody = Box::pin(futures::stream::iter(chunks)); + let error = + collect_response_body(decrypt_sse_stream_with_line_limit(source, session_key, 8)) + .await + .unwrap_err(); + assert!(matches!( + error, + Error::InvalidResponse(message) if message.contains("8-byte limit") + )); + } + } + + #[tokio::test] + async fn typed_chat_completion_keeps_legacy_stream_false_and_error_mapping() { + let mock_server = MockServer::start().await; + let client = + OpenSecretClient::new_with_api_key(mock_server.uri(), "api_key".to_string()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [32u8; 32]; + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ + "model": "typed-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": false + }), + }) + .respond_with(ResponseTemplate::new(409).set_body_string("typed conflict")) + .expect(1) + .mount(&mock_server) + .await; + + let error = client + .create_chat_completion(ChatCompletionRequest { + model: "typed-model".to_string(), + messages: vec![ChatMessage { + role: "user".to_string(), + content: json!("hi"), + tool_calls: None, + reasoning_content: None, + }], + temperature: None, + max_tokens: None, + stream: Some(true), + stream_options: None, + tools: None, + tool_choice: None, + }) + .await + .unwrap_err(); + + assert!(matches!( + error, + Error::Api { status: 409, message } if message == "typed conflict" + )); + } + + #[tokio::test] + async fn typed_models_and_embeddings_remain_compatible() { + let mock_server = MockServer::start().await; + let client = + OpenSecretClient::new_with_api_key(mock_server.uri(), "api_key".to_string()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [33u8; 32]; + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + + Mock::given(method("GET")) + .and(path("/v1/models")) + .respond_with(ResponseTemplate::new(200).set_body_json(encrypted_response( + &session_key, + &json!({ + "object": "list", + "data": [{"id": "model-a", "object": "model"}] + }), + ))) + .expect(1) + .mount(&mock_server) + .await; + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ + "input": "hello", + "model": "embedding-model" + }), + }) + .respond_with(ResponseTemplate::new(200).set_body_json(encrypted_response( + &session_key, + &json!({ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}], + "model": "embedding-model", + "usage": {"prompt_tokens": 1, "total_tokens": 1} + }), + ))) + .expect(1) + .mount(&mock_server) + .await; + + let models = client.get_models().await.unwrap(); + let embeddings = client + .create_embeddings(EmbeddingRequest { + input: "hello".into(), + model: "embedding-model".to_string(), + encoding_format: None, + dimensions: None, + user: None, + }) + .await + .unwrap(); + + assert_eq!(models.data[0].id, "model-a"); + assert_eq!( + embeddings.data[0].embedding.as_floats(), + Some(&[0.25, 0.5][..]) + ); + } + #[tokio::test] async fn test_refresh_reestablishes_attestation_without_sending_auth_headers() { let mock_server = MockServer::start().await; diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 7a4b6dd..dd1859c 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -7,7 +7,7 @@ pub mod push; pub mod session; pub mod types; -pub use client::OpenSecretClient; +pub use client::{InferenceRequest, InferenceResponse, OpenSecretClient, OpenSecretResponseBody}; pub use error::{Error, Result}; pub use push::*; pub use types::*; diff --git a/website/docs/maple-ai/index.md b/website/docs/maple-ai/index.md index 50ae6ca..60d8f20 100644 --- a/website/docs/maple-ai/index.md +++ b/website/docs/maple-ai/index.md @@ -306,16 +306,20 @@ Add to your `Cargo.toml`: ```toml [dependencies] -opensecret = "0.2" -tokio = { version = "1", features = ["full"] } +opensecret = "3.4.0" +bytes = "1" futures = "0.3" +http = "1" +tokio = { version = "1", features = ["full"] } ``` ### Quick Start ```rust -use opensecret::{OpenSecretClient, ChatCompletionRequest, ChatMessage, Result}; +use bytes::Bytes; use futures::StreamExt; +use http::Request; +use opensecret::{OpenSecretClient, Result}; #[tokio::main] async fn main() -> Result<()> { @@ -330,37 +334,24 @@ async fn main() -> Result<()> { // Perform attestation handshake (required) client.perform_attestation_handshake().await?; - // Create streaming chat completion - let request = ChatCompletionRequest { - model: "gpt-oss-120b".to_string(), - messages: vec![ChatMessage { - role: "user".to_string(), - content: serde_json::json!("Hello, world!"), - tool_calls: None, - }], - temperature: Some(0.7), - max_tokens: Some(1000), - stream: Some(true), - stream_options: None, - tools: None, - tool_choice: None, - }; - - let mut stream = client.create_chat_completion_stream(request).await?; - - while let Some(chunk_result) = stream.next().await { - match chunk_result { - Ok(chunk) => { - if !chunk.choices.is_empty() { - if let Some(serde_json::Value::String(content)) = - &chunk.choices[0].delta.content - { - print!("{}", content); - } - } - } - Err(e) => eprintln!("Error: {}", e), - } + // The lossless transport preserves provider-specific parameters and + // returns decrypted response bytes for the caller to parse or forward. + let request = Request::post("/v1/chat/completions") + .body(Bytes::from_static(br#"{ + "model": "gemma4-31b", + "messages": [{"role": "user", "content": "Hello, world!"}], + "stream": true, + "include_reasoning": false, + "chat_template_kwargs": {"enable_thinking": false} + }"#)) + .expect("valid request"); + + let response = client.send_inference_request(request).await?; + let mut body = response.into_body(); + while let Some(chunk) = body.next().await { + let decrypted_bytes = chunk?; + // Parse or forward the SSE bytes in the calling application. + print!("{}", String::from_utf8_lossy(&decrypted_bytes)); } Ok(()) @@ -417,7 +408,7 @@ match client.get_models().await { | Client creation | `createCustomFetch({ apiKey })` | `OpenSecretClient::new_with_api_key(url, key)` | | Attestation | Automatic (in fetch) | Manual: `client.perform_attestation_handshake().await?` | | Streaming | `for await (const chunk of stream)` | `while let Some(chunk) = stream.next().await` | -| OpenAI compat | Uses `openai` npm package | Native implementation | +| Inference request | Uses `openai` npm package | Lossless encrypted HTTP transport; typed compatibility helpers also available | ---