From eb9785d3df9bde36ec6d9d2f7c5817f799f243d6 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Fri, 17 Jul 2026 06:03:30 +0000 Subject: [PATCH 1/7] fix(rust): serialize concurrent token refresh --- rust/src/client.rs | 301 ++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 271 insertions(+), 30 deletions(-) diff --git a/rust/src/client.rs b/rust/src/client.rs index 193d021..04a3bb0 100644 --- a/rust/src/client.rs +++ b/rust/src/client.rs @@ -20,6 +20,7 @@ use std::{ pin::Pin, sync::{Arc, RwLock}, }; +use tokio::sync::Mutex; use uuid::Uuid; /// A decrypted response body returned by [`OpenSecretClient::send_inference_request`]. @@ -47,10 +48,16 @@ pub struct OpenSecretClient { client: Client, base_url: String, session_manager: SessionManager, + refresh_lock: Mutex<()>, use_mock_attestation: bool, server_public_key: Arc>>>, // Store server's public key from attestation } +struct ResolvedAuth { + token: Option, + using_api_key: bool, +} + fn append_query_param(query: &mut Vec, key: &str, value: impl ToString) { let encoded = utf8_percent_encode(&value.to_string(), NON_ALPHANUMERIC).to_string(); query.push(format!("{}={}", key, encoded)); @@ -381,6 +388,7 @@ impl OpenSecretClient { client: Client::new(), base_url: base_url.trim_end_matches('/').to_string(), session_manager: SessionManager::new(), + refresh_lock: Mutex::new(()), use_mock_attestation: use_mock, server_public_key: Arc::new(RwLock::new(None)), }) @@ -397,6 +405,7 @@ impl OpenSecretClient { client: Client::new(), base_url: base_url.trim_end_matches('/').to_string(), session_manager: SessionManager::new_with_api_key(api_key), + refresh_lock: Mutex::new(()), use_mock_attestation: use_mock, server_public_key: Arc::new(RwLock::new(None)), }) @@ -664,8 +673,9 @@ impl OpenSecretClient { let mut retried_attestation = false; loop { + let auth = self.resolve_auth(auth_mode)?; match self - .encrypted_json_call_inner(endpoint, method, data.clone(), auth_mode) + .encrypted_json_call_inner(endpoint, method, data.clone(), &auth) .await { Ok(result) => return Ok(result), @@ -690,8 +700,9 @@ impl OpenSecretClient { let mut retried_refresh = false; loop { + let auth = self.resolve_auth(auth_mode)?; match self - .encrypted_json_call_inner(endpoint, method, data.clone(), auth_mode) + .encrypted_json_call_inner(endpoint, method, data.clone(), &auth) .await { Ok(result) => return Ok(result), @@ -700,9 +711,10 @@ impl OpenSecretClient { retried_attestation = true; } Err(Error::Api { status: 401, .. }) - if allow_refresh && !retried_refresh && !self.using_api_key(auth_mode)? => + if allow_refresh && !retried_refresh && !auth.using_api_key => { - self.refresh_token().await?; + self.refresh_after_unauthorized(auth.token.as_deref()) + .await?; retried_refresh = true; } Err(error) => return Err(error), @@ -715,10 +727,10 @@ impl OpenSecretClient { endpoint: &str, method: &str, data: Option, - auth_mode: AuthHeaderMode, + auth: &ResolvedAuth, ) -> Result { let (response, session_key) = self - .send_encrypted_request(endpoint, method, data, auth_mode, false) + .send_encrypted_request(endpoint, method, data, auth, false) .await?; let encrypted_response: EncryptedResponse = response.json().await?; let decrypted = @@ -776,8 +788,15 @@ impl OpenSecretClient { let mut retried_refresh = false; loop { + let auth = self.resolve_auth(AuthHeaderMode::ApiKeyOrJwt)?; let result = self - .send_inference_request_once(&parts.method, &path_and_query, &headers, body.clone()) + .send_inference_request_once( + &parts.method, + &path_and_query, + &headers, + body.clone(), + &auth, + ) .await; match result { @@ -794,9 +813,13 @@ impl OpenSecretClient { Ok((response, session_key)) if response.status() == reqwest::StatusCode::UNAUTHORIZED && !retried_refresh - && !self.using_api_key(AuthHeaderMode::ApiKeyOrJwt)? => + && !auth.using_api_key => { - if self.refresh_token().await.is_ok() { + if self + .refresh_after_unauthorized(auth.token.as_deref()) + .await + .is_ok() + { retried_refresh = true; } else { return self.finish_inference_response(response, session_key).await; @@ -820,6 +843,7 @@ impl OpenSecretClient { path_and_query: &str, caller_headers: &HttpHeaderMap, body: Bytes, + auth: &ResolvedAuth, ) -> Result<(reqwest::Response, [u8; 32])> { let session = self.session_manager.get_session()?.ok_or_else(|| { Error::Session( @@ -834,7 +858,7 @@ impl OpenSecretClient { 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)? { + if let Some(token) = &auth.token { headers.insert( AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {token}")).map_err(|error| { @@ -946,8 +970,9 @@ impl OpenSecretClient { let mut retried_refresh = false; loop { + let auth = self.resolve_auth(auth_mode)?; match self - .send_encrypted_request(endpoint, method, data.clone(), auth_mode, true) + .send_encrypted_request(endpoint, method, data.clone(), &auth, true) .await { Ok(response) => return Ok(response), @@ -956,9 +981,10 @@ impl OpenSecretClient { retried_attestation = true; } Err(Error::Api { status: 401, .. }) - if allow_refresh && !retried_refresh && !self.using_api_key(auth_mode)? => + if allow_refresh && !retried_refresh && !auth.using_api_key => { - self.refresh_token().await?; + self.refresh_after_unauthorized(auth.token.as_deref()) + .await?; retried_refresh = true; } Err(error) => return Err(error), @@ -971,7 +997,7 @@ impl OpenSecretClient { endpoint: &str, method: &str, data: Option, - auth_mode: AuthHeaderMode, + auth: &ResolvedAuth, accept_sse: bool, ) -> Result<(reqwest::Response, [u8; 32])> { let session = self.session_manager.get_session()?.ok_or_else(|| { @@ -992,7 +1018,7 @@ impl OpenSecretClient { None }; - let headers = self.build_encrypted_headers(&session, auth_mode, accept_sse)?; + let headers = self.build_encrypted_headers(&session, auth, accept_sse)?; let request_builder = match method { "GET" => self.client.get(&url), "POST" => self.client.post(&url), @@ -1031,7 +1057,7 @@ impl OpenSecretClient { fn build_encrypted_headers( &self, session: &crate::types::SessionState, - auth_mode: AuthHeaderMode, + auth: &ResolvedAuth, accept_sse: bool, ) -> Result { let mut headers = HeaderMap::new(); @@ -1047,7 +1073,7 @@ impl OpenSecretClient { .map_err(|e| Error::Session(format!("Invalid session ID: {}", e)))?, ); - if let Some(token) = self.resolve_auth_token(auth_mode)? { + if let Some(token) = &auth.token { headers.insert( AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {}", token)).map_err(|e| { @@ -1059,27 +1085,32 @@ impl OpenSecretClient { Ok(headers) } - fn resolve_auth_token(&self, auth_mode: AuthHeaderMode) -> Result> { + fn resolve_auth(&self, auth_mode: AuthHeaderMode) -> Result { match auth_mode { - AuthHeaderMode::None => Ok(None), - AuthHeaderMode::Jwt => self.session_manager.get_access_token(), + AuthHeaderMode::None => Ok(ResolvedAuth { + token: None, + using_api_key: false, + }), + AuthHeaderMode::Jwt => Ok(ResolvedAuth { + token: self.session_manager.get_access_token()?, + using_api_key: false, + }), AuthHeaderMode::ApiKeyOrJwt => { if let Some(api_key) = self.session_manager.get_api_key()? { - Ok(Some(api_key)) + Ok(ResolvedAuth { + token: Some(api_key), + using_api_key: true, + }) } else { - self.session_manager.get_access_token() + Ok(ResolvedAuth { + token: self.session_manager.get_access_token()?, + using_api_key: false, + }) } } } } - fn using_api_key(&self, auth_mode: AuthHeaderMode) -> Result { - match auth_mode { - AuthHeaderMode::ApiKeyOrJwt => Ok(self.session_manager.get_api_key()?.is_some()), - _ => Ok(false), - } - } - fn is_attestation_retryable(error: &Error) -> bool { matches!( error, @@ -1339,7 +1370,7 @@ impl OpenSecretClient { Ok(response) } - pub async fn refresh_token(&self) -> Result<()> { + async fn refresh_token_inner(&self) -> Result<()> { let refresh_token = self .session_manager .get_refresh_token()? @@ -1357,6 +1388,26 @@ impl OpenSecretClient { Ok(()) } + async fn refresh_after_unauthorized(&self, failed_access_token: Option<&str>) -> Result<()> { + let _refresh_guard = self.refresh_lock.lock().await; + + // Another request may already have replaced the credential while this + // request was in flight or waiting for the refresh lock. In that case, + // retry with the replacement instead of rotating the refresh token a + // second time. + let current_access_token = self.session_manager.get_access_token()?; + if current_access_token.as_deref() != failed_access_token { + return Ok(()); + } + + self.refresh_token_inner().await + } + + pub async fn refresh_token(&self) -> Result<()> { + let _refresh_guard = self.refresh_lock.lock().await; + self.refresh_token_inner().await + } + async fn logout_inner(&self, push_device_id: Option) -> Result<()> { let refresh_token = self .session_manager @@ -2331,6 +2382,7 @@ mod tests { use crate::PushNotificationKeyPair; use futures::StreamExt; use serde_json::json; + use std::time::Duration; use wiremock::{ matchers::{header, method, path, query_param}, Match, Mock, MockServer, Request, Respond, ResponseTemplate, @@ -2936,6 +2988,90 @@ mod tests { ); } + #[tokio::test] + async fn concurrent_authenticated_401s_share_one_refresh() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [39u8; 32]; + + 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("GET")) + .and(path("/protected/user")) + .and(header("authorization", "Bearer expired_access")) + .respond_with( + ResponseTemplate::new(401) + .set_delay(Duration::from_millis(50)) + .set_body_json(json!({ "message": "jwt expired" })), + ) + .expect(2) + .mount(&mock_server) + .await; + + Mock::given(method("POST")) + .and(path("/refresh")) + .and(MissingHeaderMatcher("authorization")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(100)) + .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("GET")) + .and(path("/protected/user")) + .and(header("authorization", "Bearer fresh_access")) + .respond_with(ResponseTemplate::new(200).set_body_json(encrypted_response( + &session_key, + &json!({ + "user": { + "id": Uuid::new_v4(), + "name": null, + "email": "sdk@test.dev", + "email_verified": true, + "login_method": "email", + "created_at": "2024-01-01T00:00:00Z", + "updated_at": "2024-01-01T00:00:00Z" + } + }), + ))) + .expect(2) + .mount(&mock_server) + .await; + + let (first, second) = tokio::join!(client.get_user(), client.get_user()); + assert_eq!(first.unwrap().user.email.as_deref(), Some("sdk@test.dev")); + assert_eq!(second.unwrap().user.email.as_deref(), Some("sdk@test.dev")); + assert_eq!( + client.get_access_token().unwrap().as_deref(), + Some("fresh_access") + ); + assert_eq!( + client.get_refresh_token().unwrap().as_deref(), + Some("fresh_refresh") + ); + mock_server.verify().await; + } + #[tokio::test] async fn test_corrupted_access_token_recovers_via_refresh_on_next_call() { let mock_server = MockServer::start().await; @@ -3345,6 +3481,111 @@ mod tests { ); } + #[tokio::test] + async fn concurrent_inference_401s_share_one_refresh() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [40u8; 32]; + let request_body = Bytes::from_static(br#"{"model":"test","messages":[]}"#); + let response_body = Bytes::from_static(br#"{"id":"completion-ok"}"#); + + 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_delay(Duration::from_millis(50)) + .set_body_string("jwt expired"), + ) + .expect(2) + .mount(&mock_server) + .await; + + Mock::given(method("POST")) + .and(path("/refresh")) + .and(MissingHeaderMatcher("authorization")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(100)) + .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(200) + .set_body_json(encrypted_response_bytes(&session_key, &response_body)), + ) + .expect(2) + .mount(&mock_server) + .await; + + let make_request = || { + HttpRequest::builder() + .method(http::Method::POST) + .uri("/v1/chat/completions") + .body(request_body.clone()) + .unwrap() + }; + let (first, second) = tokio::join!( + client.send_inference_request(make_request()), + client.send_inference_request(make_request()) + ); + + let first = first.unwrap(); + let second = second.unwrap(); + assert_eq!(first.status(), http::StatusCode::OK); + assert_eq!(second.status(), http::StatusCode::OK); + assert_eq!( + collect_response_body(first.into_body()).await.unwrap(), + response_body + ); + assert_eq!( + collect_response_body(second.into_body()).await.unwrap(), + response_body + ); + assert_eq!( + client.get_access_token().unwrap().as_deref(), + Some("fresh_access") + ); + assert_eq!( + client.get_refresh_token().unwrap().as_deref(), + Some("fresh_refresh") + ); + mock_server.verify().await; + } + #[tokio::test] async fn inference_transport_decrypts_non_success_encrypted_body() { let mock_server = MockServer::start().await; From cd634373de83fa34e43e2bce3b44702afdd8495c Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Fri, 17 Jul 2026 06:28:58 +0000 Subject: [PATCH 2/7] fix(rust): harden authenticated transport races --- rust/src/client.rs | 809 +++++++++++++++++++++++++++++++++++++++----- rust/src/session.rs | 172 ++++++++-- 2 files changed, 867 insertions(+), 114 deletions(-) diff --git a/rust/src/client.rs b/rust/src/client.rs index 04a3bb0..24640c7 100644 --- a/rust/src/client.rs +++ b/rust/src/client.rs @@ -16,10 +16,7 @@ use reqwest::{ Client, }; use serde::{de::DeserializeOwned, Deserialize, Serialize}; -use std::{ - pin::Pin, - sync::{Arc, RwLock}, -}; +use std::pin::Pin; use tokio::sync::Mutex; use uuid::Uuid; @@ -50,12 +47,13 @@ pub struct OpenSecretClient { session_manager: SessionManager, refresh_lock: Mutex<()>, use_mock_attestation: bool, - server_public_key: Arc>>>, // Store server's public key from attestation } +#[derive(Debug, Clone, PartialEq, Eq)] struct ResolvedAuth { token: Option, using_api_key: bool, + generation: u64, } fn append_query_param(query: &mut Vec, key: &str, value: impl ToString) { @@ -300,14 +298,16 @@ fn transform_sse_line(line: Bytes, session_key: &[u8; 32]) -> Result { 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), - }; + // OpenSecret encrypts every normal data event. Plaintext, malformed, or + // unauthenticated data is corrupt transport data and must never be passed + // through as trusted inference output. + let encrypted = BASE64.decode(payload).map_err(|_| { + Error::InvalidResponse("Inference SSE data was not valid encrypted payload".to_string()) + })?; if encrypted.len() < 28 { - return Ok(line); + return Err(Error::InvalidResponse( + "Inference SSE data was shorter than the encrypted payload minimum".to_string(), + )); } let decrypted = crypto::decrypt_data(session_key, &encrypted) .map_err(|error| Error::Decryption(format!("Failed to decrypt SSE data: {error}")))?; @@ -390,7 +390,6 @@ impl OpenSecretClient { session_manager: SessionManager::new(), refresh_lock: Mutex::new(()), use_mock_attestation: use_mock, - server_public_key: Arc::new(RwLock::new(None)), }) } @@ -407,7 +406,6 @@ impl OpenSecretClient { session_manager: SessionManager::new_with_api_key(api_key), refresh_lock: Mutex::new(()), use_mock_attestation: use_mock, - server_public_key: Arc::new(RwLock::new(None)), }) } @@ -435,19 +433,15 @@ impl OpenSecretClient { self.parse_mock_attestation(&attestation_doc.attestation_document)? }; - // Store server's public key from attestation document - if let Some(pub_key) = doc.public_key { - *self.server_public_key.write().map_err(|e| { - Error::KeyExchange(format!("Failed to write server public key: {}", e)) - })? = Some(pub_key); - } else { - return Err(Error::AttestationVerificationFailed( + let server_public_key = doc.public_key.ok_or_else(|| { + Error::AttestationVerificationFailed( "No public key in attestation document".to_string(), - )); - } + ) + })?; // Step 3: Perform key exchange - self.perform_key_exchange(&nonce).await?; + self.perform_key_exchange(&nonce, &server_public_key) + .await?; Ok(()) } @@ -472,7 +466,7 @@ impl OpenSecretClient { response.json().await.map_err(Into::into) } - async fn perform_key_exchange(&self, nonce: &str) -> Result<()> { + async fn perform_key_exchange(&self, nonce: &str, server_public_key: &[u8]) -> Result<()> { // Generate ephemeral keypair let (secret, public_key) = crypto::generate_static_keypair(); let public_key_bytes = public_key.as_bytes(); @@ -510,18 +504,9 @@ impl OpenSecretClient { let key_exchange_response: KeyExchangeResponse = response.json().await?; - // Get server's public key from attestation - let server_public_key_bytes = self - .server_public_key - .read() - .map_err(|e| Error::KeyExchange(format!("Failed to read server public key: {}", e)))?; - let server_public_key_bytes = server_public_key_bytes - .as_ref() - .ok_or_else(|| Error::KeyExchange("Server public key not available".to_string()))?; - // Convert server's public key bytes to x25519 PublicKey let server_public_key = x25519_dalek::PublicKey::from( - <[u8; 32]>::try_from(server_public_key_bytes.as_slice()) + <[u8; 32]>::try_from(server_public_key) .map_err(|_| Error::KeyExchange("Invalid server public key length".to_string()))?, ); @@ -710,12 +695,17 @@ impl OpenSecretClient { self.perform_attestation_handshake().await?; retried_attestation = true; } - Err(Error::Api { status: 401, .. }) - if allow_refresh && !retried_refresh && !auth.using_api_key => + Err(error @ Error::Api { status: 401, .. }) + if allow_refresh && !retried_refresh => { - self.refresh_after_unauthorized(auth.token.as_deref()) - .await?; - retried_refresh = true; + if self + .recover_auth_after_unauthorized(auth_mode, &auth) + .await? + { + retried_refresh = true; + } else { + return Err(error); + } } Err(error) => return Err(error), } @@ -812,14 +802,13 @@ impl OpenSecretClient { } Ok((response, session_key)) if response.status() == reqwest::StatusCode::UNAUTHORIZED - && !retried_refresh - && !auth.using_api_key => + && !retried_refresh => { - if self - .refresh_after_unauthorized(auth.token.as_deref()) - .await - .is_ok() - { + if matches!( + self.recover_auth_after_unauthorized(AuthHeaderMode::ApiKeyOrJwt, &auth,) + .await, + Ok(true) + ) { retried_refresh = true; } else { return self.finish_inference_response(response, session_key).await; @@ -980,12 +969,17 @@ impl OpenSecretClient { self.perform_attestation_handshake().await?; retried_attestation = true; } - Err(Error::Api { status: 401, .. }) - if allow_refresh && !retried_refresh && !auth.using_api_key => + Err(error @ Error::Api { status: 401, .. }) + if allow_refresh && !retried_refresh => { - self.refresh_after_unauthorized(auth.token.as_deref()) - .await?; - retried_refresh = true; + if self + .recover_auth_after_unauthorized(auth_mode, &auth) + .await? + { + retried_refresh = true; + } else { + return Err(error); + } } Err(error) => return Err(error), } @@ -1086,25 +1080,33 @@ impl OpenSecretClient { } fn resolve_auth(&self, auth_mode: AuthHeaderMode) -> Result { + let credentials = self.session_manager.get_credential_snapshot()?; match auth_mode { AuthHeaderMode::None => Ok(ResolvedAuth { token: None, using_api_key: false, + generation: 0, }), AuthHeaderMode::Jwt => Ok(ResolvedAuth { - token: self.session_manager.get_access_token()?, + token: credentials + .tokens + .as_ref() + .map(|tokens| tokens.access_token.clone()), using_api_key: false, + generation: credentials.token_generation, }), AuthHeaderMode::ApiKeyOrJwt => { - if let Some(api_key) = self.session_manager.get_api_key()? { + if let Some(api_key) = credentials.api_key { Ok(ResolvedAuth { token: Some(api_key), using_api_key: true, + generation: credentials.api_key_generation, }) } else { Ok(ResolvedAuth { - token: self.session_manager.get_access_token()?, + token: credentials.tokens.map(|tokens| tokens.access_token), using_api_key: false, + generation: credentials.token_generation, }) } } @@ -1371,9 +1373,11 @@ impl OpenSecretClient { } async fn refresh_token_inner(&self) -> Result<()> { - let refresh_token = self - .session_manager - .get_refresh_token()? + let credentials = self.session_manager.get_credential_snapshot()?; + let refresh_token = credentials + .tokens + .as_ref() + .and_then(|tokens| tokens.refresh_token.clone()) .ok_or_else(|| Error::Authentication("No refresh token available".to_string()))?; let request = RefreshRequest { refresh_token }; @@ -1382,25 +1386,40 @@ impl OpenSecretClient { .encrypted_api_call("/refresh", "POST", Some(request)) .await?; - self.session_manager - .set_tokens(response.access_token, Some(response.refresh_token))?; + // A synchronous set_tokens or a logout/clear may have replaced these + // credentials while the HTTP refresh was in flight. Drop this stale + // response instead of reinstalling credentials the caller superseded. + self.session_manager.set_tokens_if_generation( + credentials.token_generation, + response.access_token, + Some(response.refresh_token), + )?; Ok(()) } - async fn refresh_after_unauthorized(&self, failed_access_token: Option<&str>) -> Result<()> { + async fn recover_auth_after_unauthorized( + &self, + auth_mode: AuthHeaderMode, + failed_auth: &ResolvedAuth, + ) -> Result { let _refresh_guard = self.refresh_lock.lock().await; - // Another request may already have replaced the credential while this - // request was in flight or waiting for the refresh lock. In that case, - // retry with the replacement instead of rotating the refresh token a - // second time. - let current_access_token = self.session_manager.get_access_token()?; - if current_access_token.as_deref() != failed_access_token { - return Ok(()); + // Another request or the application may already have replaced the + // exact credential source while this request was in flight or waiting + // for the refresh lock. Retry with that replacement, including an API + // key/JWT source switch, rather than refreshing the wrong credential. + let current_auth = self.resolve_auth(auth_mode)?; + if current_auth != *failed_auth { + return Ok(true); } - self.refresh_token_inner().await + if current_auth.using_api_key { + return Ok(false); + } + + self.refresh_token_inner().await?; + Ok(true) } pub async fn refresh_token(&self) -> Result<()> { @@ -1409,9 +1428,15 @@ impl OpenSecretClient { } async fn logout_inner(&self, push_device_id: Option) -> Result<()> { - let refresh_token = self - .session_manager - .get_refresh_token()? + // Serialize logout with refresh so an internal token rotation cannot + // race the clear. Application-supplied credentials remain lock-free + // and win through the generation check below. + let _refresh_guard = self.refresh_lock.lock().await; + let credentials = self.session_manager.get_credential_snapshot()?; + let refresh_token = credentials + .tokens + .as_ref() + .and_then(|tokens| tokens.refresh_token.clone()) .ok_or_else(|| Error::Authentication("No refresh token available".to_string()))?; let request = LogoutRequest { @@ -1423,8 +1448,10 @@ impl OpenSecretClient { .encrypted_api_call("/logout", "POST", Some(request)) .await?; - // Clear all session data - self.session_manager.clear_all()?; + // Do not clear credentials installed by the application while the + // logout request was in flight (for example, a rapid account switch). + self.session_manager + .clear_all_if_generation(credentials.generation)?; Ok(()) } @@ -2382,7 +2409,14 @@ mod tests { use crate::PushNotificationKeyPair; use futures::StreamExt; use serde_json::json; - use std::time::Duration; + use std::{ + collections::HashMap, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex as StdMutex, + }, + time::Duration, + }; use wiremock::{ matchers::{header, method, path, query_param}, Match, Mock, MockServer, Request, Respond, ResponseTemplate, @@ -2492,6 +2526,70 @@ mod tests { } } + #[derive(Clone)] + struct PerNonceAttestationResponder { + server_secrets: Arc>>, + next_key: Arc, + } + + impl Respond for PerNonceAttestationResponder { + fn respond(&self, request: &Request) -> ResponseTemplate { + let nonce = request.url.path().rsplit('/').next().unwrap_or_default(); + let key_byte = self.next_key.fetch_add(1, Ordering::SeqCst) as u8 + 1; + let server_secret_key = [key_byte; 32]; + let server_public_key = + x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from(server_secret_key)); + self.server_secrets + .lock() + .unwrap() + .insert(nonce.to_string(), server_secret_key); + + ResponseTemplate::new(200).set_body_json(json!({ + "attestation_document": build_mock_attestation_document( + nonce, + server_public_key.as_bytes(), + ) + })) + } + } + + #[derive(Clone)] + struct PerNonceKeyExchangeResponder { + server_secrets: Arc>>, + session_key: [u8; 32], + } + + impl Respond for PerNonceKeyExchangeResponder { + fn respond(&self, request: &Request) -> ResponseTemplate { + let body: KeyExchangeRequest = serde_json::from_slice(request.body.as_ref()).unwrap(); + let server_secret_key = *self + .server_secrets + .lock() + .unwrap() + .get(&body.nonce) + .expect("key exchange nonce must match an attestation"); + let client_public_bytes = BASE64.decode(body.client_public_key.as_bytes()).unwrap(); + let client_public_key = x25519_dalek::PublicKey::from( + <[u8; 32]>::try_from(client_public_bytes.as_slice()).unwrap(), + ); + let server_secret = x25519_dalek::StaticSecret::from(server_secret_key); + let shared_secret = + crypto::perform_static_key_exchange(&server_secret, &client_public_key); + let encrypted_session_key = BASE64 + .encode(crypto::encrypt_data(shared_secret.as_bytes(), &self.session_key).unwrap()); + let session_id = Uuid::from_u128(server_secret_key[0] as u128).to_string(); + + // Keep both key exchanges in flight long enough for both distinct + // attestation results to be processed by the client. + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(100)) + .set_body_json(json!({ + "encrypted_session_key": encrypted_session_key, + "session_id": session_id, + })) + } + } + fn build_mock_attestation_document(nonce: &str, server_public_key: &[u8; 32]) -> String { let payload = CborValue::Map(vec![ ( @@ -2665,6 +2763,44 @@ mod tests { assert!(client.use_mock_attestation); } + #[tokio::test] + async fn concurrent_attestation_handshakes_keep_each_nonce_public_key() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let server_secrets = Arc::new(StdMutex::new(HashMap::new())); + let next_key = Arc::new(AtomicUsize::new(0)); + let session_key = [42u8; 32]; + + Mock::given(method("GET")) + .and(PathPrefixMatcher("/attestation/")) + .respond_with(PerNonceAttestationResponder { + server_secrets: Arc::clone(&server_secrets), + next_key, + }) + .expect(2) + .mount(&mock_server) + .await; + Mock::given(method("POST")) + .and(path("/key_exchange")) + .respond_with(PerNonceKeyExchangeResponder { + server_secrets: Arc::clone(&server_secrets), + session_key, + }) + .expect(2) + .mount(&mock_server) + .await; + + let (first, second) = tokio::join!( + client.perform_attestation_handshake(), + client.perform_attestation_handshake() + ); + + first.unwrap(); + second.unwrap(); + assert_eq!(server_secrets.lock().unwrap().len(), 2); + assert!(client.get_session_id().unwrap().is_some()); + } + #[tokio::test] async fn test_register_push_device_uses_v1_push_endpoint() { let mock_server = MockServer::start().await; @@ -3586,6 +3722,68 @@ mod tests { mock_server.verify().await; } + #[tokio::test] + async fn inference_401_retries_when_auth_source_changes_from_api_key_to_jwt() { + let mock_server = MockServer::start().await; + let client = + OpenSecretClient::new_with_api_key(mock_server.uri(), "old_api_key".to_string()) + .unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [46u8; 32]; + let request_body = Bytes::from_static(br#"{"model":"test","messages":[]}"#); + let response_body = Bytes::from_static(br#"{"id":"jwt-completion"}"#); + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + client + .session_manager + .set_tokens("jwt_access".to_string(), Some("jwt_refresh".to_string())) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(header("authorization", "Bearer old_api_key")) + .respond_with( + ResponseTemplate::new(401) + .set_delay(Duration::from_millis(100)) + .set_body_string("api key rejected"), + ) + .expect(1) + .mount(&mock_server) + .await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(header("authorization", "Bearer jwt_access")) + .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/chat/completions") + .body(request_body) + .unwrap(); + let (response, clear_result) = + tokio::join!(client.send_inference_request(request), async { + tokio::time::sleep(Duration::from_millis(25)).await; + client.clear_api_key() + }); + + clear_result.unwrap(); + let response = response.unwrap(); + assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!( + collect_response_body(response.into_body()).await.unwrap(), + response_body + ); + mock_server.verify().await; + } + #[tokio::test] async fn inference_transport_decrypts_non_success_encrypted_body() { let mock_server = MockServer::start().await; @@ -3730,10 +3928,10 @@ mod tests { 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" + ": heartbeat\r\nevent: chunk\r\nid: provider-7\r\nretry: 1500\r\ndata: {encrypted_payload}\r\n\r\n: provider-heartbeat\n\ndata:\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", + ": heartbeat\r\nevent: chunk\r\nid: provider-7\r\nretry: 1500\r\ndata: {}\r\n\r\n: provider-heartbeat\n\ndata:\n\ndata: [DONE]\n\n", String::from_utf8_lossy(decrypted_payload) ); let chunks = encrypted_sse @@ -3757,6 +3955,42 @@ mod tests { ); } + #[tokio::test] + async fn inference_sse_transport_rejects_plaintext_completion_chunk() { + let session_key = [47u8; 32]; + let source: OpenSecretResponseBody = Box::pin(futures::stream::iter([Ok( + Bytes::from_static( + br#"data: {"id":"chatcmpl-injected","choices":[{"delta":{"content":"untrusted"}}]}\n\n"#, + ), + )])); + + let error = collect_response_body(decrypt_sse_stream(source, session_key)) + .await + .unwrap_err(); + + assert!(matches!( + error, + Error::InvalidResponse(message) if message.contains("not valid encrypted payload") + )); + } + + #[tokio::test] + async fn inference_sse_transport_rejects_short_base64_payload() { + let session_key = [48u8; 32]; + let source: OpenSecretResponseBody = Box::pin(futures::stream::iter([Ok( + Bytes::from_static(b"data: YWJj\n\n"), + )])); + + let error = collect_response_body(decrypt_sse_stream(source, session_key)) + .await + .unwrap_err(); + + assert!(matches!( + error, + Error::InvalidResponse(message) if message.contains("encrypted payload minimum") + )); + } + #[tokio::test] async fn inference_sse_transport_reports_corrupt_ciphertext() { let session_key = [31u8; 32]; @@ -4002,6 +4236,215 @@ mod tests { ); } + #[tokio::test] + async fn delayed_manual_refresh_does_not_overwrite_newly_set_tokens() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [43u8; 32]; + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + client + .session_manager + .set_tokens("old_access".to_string(), Some("old_refresh".to_string())) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/refresh")) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ "refresh_token": "old_refresh" }), + }) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(100)) + .set_body_json(encrypted_response( + &session_key, + &json!({ + "access_token": "stale_refreshed_access", + "refresh_token": "stale_refreshed_refresh", + }), + )), + ) + .expect(1) + .mount(&mock_server) + .await; + + let (refresh_result, ()) = tokio::join!(client.refresh_token(), async { + tokio::time::sleep(Duration::from_millis(25)).await; + client + .set_tokens("app_access".to_string(), Some("app_refresh".to_string())) + .unwrap(); + }); + + refresh_result.unwrap(); + assert_eq!( + client.get_access_token().unwrap().as_deref(), + Some("app_access") + ); + assert_eq!( + client.get_refresh_token().unwrap().as_deref(), + Some("app_refresh") + ); + assert!(client.get_session_id().unwrap().is_none()); + } + + #[tokio::test] + async fn delayed_manual_refresh_does_not_restore_cleared_credentials() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [44u8; 32]; + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + client + .session_manager + .set_tokens("old_access".to_string(), Some("old_refresh".to_string())) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/refresh")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(100)) + .set_body_json(encrypted_response( + &session_key, + &json!({ + "access_token": "stale_refreshed_access", + "refresh_token": "stale_refreshed_refresh", + }), + )), + ) + .expect(1) + .mount(&mock_server) + .await; + + let (refresh_result, clear_result) = tokio::join!(client.refresh_token(), async { + tokio::time::sleep(Duration::from_millis(25)).await; + client.session_manager.clear_all() + }); + + refresh_result.unwrap(); + clear_result.unwrap(); + assert!(client.get_access_token().unwrap().is_none()); + assert!(client.get_refresh_token().unwrap().is_none()); + assert!(client.get_session_id().unwrap().is_none()); + } + + #[tokio::test] + async fn concurrent_manual_refresh_and_logout_finish_logged_out() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [45u8; 32]; + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + client + .session_manager + .set_tokens("old_access".to_string(), Some("old_refresh".to_string())) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/refresh")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(100)) + .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("/logout")) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ + "refresh_token": "fresh_refresh" + }), + }) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(encrypted_response(&session_key, &json!({}))), + ) + .expect(1) + .mount(&mock_server) + .await; + + let (refresh_result, logout_result) = tokio::join!(client.refresh_token(), async { + tokio::time::sleep(Duration::from_millis(25)).await; + client.logout().await + }); + + refresh_result.unwrap(); + logout_result.unwrap(); + assert!(client.get_access_token().unwrap().is_none()); + assert!(client.get_refresh_token().unwrap().is_none()); + assert!(client.get_session_id().unwrap().is_none()); + } + + #[tokio::test] + async fn delayed_logout_does_not_clear_newly_set_tokens() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [49u8; 32]; + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + client + .session_manager + .set_tokens("old_access".to_string(), Some("old_refresh".to_string())) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/logout")) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ + "refresh_token": "old_refresh" + }), + }) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(100)) + .set_body_json(encrypted_response(&session_key, &json!({}))), + ) + .expect(1) + .mount(&mock_server) + .await; + + let (logout_result, ()) = tokio::join!(client.logout(), async { + tokio::time::sleep(Duration::from_millis(25)).await; + client + .set_tokens("new_access".to_string(), Some("new_refresh".to_string())) + .unwrap(); + }); + + logout_result.unwrap(); + assert_eq!( + client.get_access_token().unwrap().as_deref(), + Some("new_access") + ); + assert_eq!( + client.get_refresh_token().unwrap().as_deref(), + Some("new_refresh") + ); + assert!(client.get_session_id().unwrap().is_none()); + } + #[tokio::test] async fn test_init_main_agent_uses_authenticated_encrypted_v1_endpoint() { let mock_server = MockServer::start().await; @@ -4295,4 +4738,216 @@ mod tests { assert!(stream.next().await.is_none()); } + + #[tokio::test] + async fn concurrent_agent_stream_401s_share_one_refresh_and_decrypt() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [50u8; 32]; + let message_id = Uuid::new_v4(); + 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/agent/chat")) + .and(header("authorization", "Bearer expired_access")) + .and(header("x-session-id", session_id.to_string())) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ "input": "stream retry" }), + }) + .respond_with( + ResponseTemplate::new(401) + .set_delay(Duration::from_millis(50)) + .set_body_string("jwt expired"), + ) + .expect(2) + .mount(&mock_server) + .await; + + Mock::given(method("POST")) + .and(path("/refresh")) + .and(MissingHeaderMatcher("authorization")) + .and(header("x-session-id", session_id.to_string())) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ "refresh_token": "refresh_token" }), + }) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(100)) + .set_body_json(encrypted_response( + &session_key, + &json!({ + "access_token": "fresh_access", + "refresh_token": "fresh_refresh", + }), + )), + ) + .expect(1) + .mount(&mock_server) + .await; + + let sse_body = format!( + "{}data: [DONE]\n\n", + encrypted_sse_data( + &session_key, + &json!({ + "message_id": message_id, + "message": "stream recovered" + }) + ) + .replacen("data:", "event: agent.message\ndata:", 1), + ); + Mock::given(method("POST")) + .and(path("/v1/agent/chat")) + .and(header("authorization", "Bearer fresh_access")) + .and(header("x-session-id", session_id.to_string())) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ "input": "stream retry" }), + }) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse_body), + ) + .expect(2) + .mount(&mock_server) + .await; + + let (first, second) = tokio::join!( + client.agent_chat("stream retry"), + client.agent_chat("stream retry") + ); + let mut first = first.unwrap(); + let mut second = second.unwrap(); + + for stream in [&mut first, &mut second] { + match stream.next().await.unwrap().unwrap() { + AgentSseEvent::Message(event) => { + assert_eq!(event.message_id, message_id); + assert_eq!(event.message, "stream recovered"); + } + other => panic!("Expected recovered agent message, got {other:?}"), + } + assert!(stream.next().await.is_none()); + } + assert_eq!( + client.get_access_token().unwrap().as_deref(), + Some("fresh_access") + ); + assert_eq!( + client.get_refresh_token().unwrap().as_deref(), + Some("fresh_refresh") + ); + mock_server.verify().await; + } + + #[tokio::test] + async fn agent_stream_stale_session_reattests_and_decrypts() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let stale_session_id = Uuid::new_v4(); + let stale_session_key = [51u8; 32]; + let server_secret_key = [52u8; 32]; + let server_public_key = + x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from(server_secret_key)); + let fresh_session_id = Uuid::new_v4(); + let fresh_session_key = [53u8; 32]; + let message_id = Uuid::new_v4(); + client + .session_manager + .set_session(stale_session_id, stale_session_key) + .unwrap(); + client + .session_manager + .set_tokens( + "access_token".to_string(), + Some("refresh_token".to_string()), + ) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/v1/agent/chat")) + .and(header("authorization", "Bearer access_token")) + .and(header("x-session-id", stale_session_id.to_string())) + .and(EncryptedJsonBodyMatcher { + session_key: stale_session_key, + expected: json!({ "input": "stale stream" }), + }) + .respond_with(ResponseTemplate::new(400).set_body_string("stale session")) + .expect(1) + .mount(&mock_server) + .await; + 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: fresh_session_key, + session_id: fresh_session_id.to_string(), + }) + .expect(1) + .mount(&mock_server) + .await; + + let sse_body = format!( + "{}data: [DONE]\n\n", + encrypted_sse_data( + &fresh_session_key, + &json!({ + "message_id": message_id, + "message": "fresh session" + }) + ) + .replacen("data:", "event: agent.message\ndata:", 1), + ); + Mock::given(method("POST")) + .and(path("/v1/agent/chat")) + .and(header("authorization", "Bearer access_token")) + .and(header("x-session-id", fresh_session_id.to_string())) + .and(EncryptedJsonBodyMatcher { + session_key: fresh_session_key, + expected: json!({ "input": "stale stream" }), + }) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse_body), + ) + .expect(1) + .mount(&mock_server) + .await; + + let mut stream = client.agent_chat("stale stream").await.unwrap(); + match stream.next().await.unwrap().unwrap() { + AgentSseEvent::Message(event) => { + assert_eq!(event.message_id, message_id); + assert_eq!(event.message, "fresh session"); + } + other => panic!("Expected agent message after re-attestation, got {other:?}"), + } + assert!(stream.next().await.is_none()); + assert_eq!(client.get_session_id().unwrap(), Some(fresh_session_id)); + mock_server.verify().await; + } } diff --git a/rust/src/session.rs b/rust/src/session.rs index 61cc868..7e032a4 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -3,52 +3,76 @@ use crate::types::{SessionState, TokenPair}; use std::sync::{Arc, RwLock}; use uuid::Uuid; +#[derive(Debug, Clone)] +pub(crate) struct CredentialSnapshot { + pub(crate) tokens: Option, + pub(crate) api_key: Option, + pub(crate) generation: u64, + pub(crate) token_generation: u64, + pub(crate) api_key_generation: u64, +} + +#[derive(Debug, Default)] +struct CredentialState { + tokens: Option, + api_key: Option, + generation: u64, + token_generation: u64, + api_key_generation: u64, +} + pub struct SessionManager { session: Arc>>, - tokens: Arc>>, - api_key: Arc>>, + credentials: Arc>, } impl SessionManager { pub fn new() -> Self { Self { session: Arc::new(RwLock::new(None)), - tokens: Arc::new(RwLock::new(None)), - api_key: Arc::new(RwLock::new(None)), + credentials: Arc::new(RwLock::new(CredentialState::default())), } } pub fn new_with_api_key(api_key: String) -> Self { Self { session: Arc::new(RwLock::new(None)), - tokens: Arc::new(RwLock::new(None)), - api_key: Arc::new(RwLock::new(Some(api_key))), + credentials: Arc::new(RwLock::new(CredentialState { + api_key: Some(api_key), + generation: 1, + api_key_generation: 1, + ..CredentialState::default() + })), } } pub fn set_api_key(&self, api_key: String) -> Result<()> { - let mut api_key_guard = self.api_key.write().map_err(|e| { - Error::Authentication(format!("Failed to acquire API key write lock: {}", e)) + let mut credentials = self.credentials.write().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials write lock: {}", e)) })?; - *api_key_guard = Some(api_key); + credentials.api_key = Some(api_key); + credentials.generation = credentials.generation.wrapping_add(1); + credentials.api_key_generation = credentials.api_key_generation.wrapping_add(1); Ok(()) } pub fn get_api_key(&self) -> Result> { - let api_key_guard = self.api_key.read().map_err(|e| { - Error::Authentication(format!("Failed to acquire API key read lock: {}", e)) + let credentials = self.credentials.read().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials read lock: {}", e)) })?; - Ok(api_key_guard.clone()) + Ok(credentials.api_key.clone()) } pub fn clear_api_key(&self) -> Result<()> { - let mut api_key_guard = self.api_key.write().map_err(|e| { - Error::Authentication(format!("Failed to acquire API key write lock: {}", e)) + let mut credentials = self.credentials.write().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials write lock: {}", e)) })?; - *api_key_guard = None; + credentials.api_key = None; + credentials.generation = credentials.generation.wrapping_add(1); + credentials.api_key_generation = credentials.api_key_generation.wrapping_add(1); Ok(()) } @@ -86,70 +110,144 @@ impl SessionManager { } pub fn set_tokens(&self, access_token: String, refresh_token: Option) -> Result<()> { - let mut tokens_guard = self.tokens.write().map_err(|e| { - Error::Authentication(format!("Failed to acquire tokens write lock: {}", e)) + let mut credentials = self.credentials.write().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials write lock: {}", e)) })?; - *tokens_guard = Some(TokenPair { + credentials.tokens = Some(TokenPair { access_token, refresh_token, }); + credentials.generation = credentials.generation.wrapping_add(1); + credentials.token_generation = credentials.token_generation.wrapping_add(1); Ok(()) } + pub(crate) fn set_tokens_if_generation( + &self, + expected_token_generation: u64, + access_token: String, + refresh_token: Option, + ) -> Result { + let mut credentials = self.credentials.write().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials write lock: {}", e)) + })?; + + if credentials.token_generation != expected_token_generation { + return Ok(false); + } + + credentials.tokens = Some(TokenPair { + access_token, + refresh_token, + }); + credentials.generation = credentials.generation.wrapping_add(1); + credentials.token_generation = credentials.token_generation.wrapping_add(1); + Ok(true) + } + + pub(crate) fn get_credential_snapshot(&self) -> Result { + let credentials = self.credentials.read().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials read lock: {}", e)) + })?; + + Ok(CredentialSnapshot { + tokens: credentials.tokens.clone(), + api_key: credentials.api_key.clone(), + generation: credentials.generation, + token_generation: credentials.token_generation, + api_key_generation: credentials.api_key_generation, + }) + } + pub fn get_tokens(&self) -> Result> { - let tokens_guard = self.tokens.read().map_err(|e| { - Error::Authentication(format!("Failed to acquire tokens read lock: {}", e)) + let credentials = self.credentials.read().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials read lock: {}", e)) })?; - Ok(tokens_guard.clone()) + Ok(credentials.tokens.clone()) } pub fn get_access_token(&self) -> Result> { - let tokens_guard = self.tokens.read().map_err(|e| { - Error::Authentication(format!("Failed to acquire tokens read lock: {}", e)) + let credentials = self.credentials.read().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials read lock: {}", e)) })?; - Ok(tokens_guard.as_ref().map(|t| t.access_token.clone())) + Ok(credentials + .tokens + .as_ref() + .map(|tokens| tokens.access_token.clone())) } pub fn get_refresh_token(&self) -> Result> { - let tokens_guard = self.tokens.read().map_err(|e| { - Error::Authentication(format!("Failed to acquire tokens read lock: {}", e)) + let credentials = self.credentials.read().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials read lock: {}", e)) })?; - Ok(tokens_guard.as_ref().and_then(|t| t.refresh_token.clone())) + Ok(credentials + .tokens + .as_ref() + .and_then(|tokens| tokens.refresh_token.clone())) } pub fn update_access_token(&self, access_token: String) -> Result<()> { - let mut tokens_guard = self.tokens.write().map_err(|e| { - Error::Authentication(format!("Failed to acquire tokens write lock: {}", e)) + let mut credentials = self.credentials.write().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials write lock: {}", e)) })?; - if let Some(tokens) = tokens_guard.as_mut() { + if let Some(tokens) = credentials.tokens.as_mut() { tokens.access_token = access_token; - Ok(()) } else { - Err(Error::Authentication("No tokens to update".to_string())) + return Err(Error::Authentication("No tokens to update".to_string())); } + + credentials.generation = credentials.generation.wrapping_add(1); + credentials.token_generation = credentials.token_generation.wrapping_add(1); + Ok(()) } pub fn clear_tokens(&self) -> Result<()> { - let mut tokens_guard = self.tokens.write().map_err(|e| { - Error::Authentication(format!("Failed to acquire tokens write lock: {}", e)) + let mut credentials = self.credentials.write().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials write lock: {}", e)) })?; - *tokens_guard = None; + credentials.tokens = None; + credentials.generation = credentials.generation.wrapping_add(1); + credentials.token_generation = credentials.token_generation.wrapping_add(1); Ok(()) } pub fn clear_all(&self) -> Result<()> { self.clear_session()?; - self.clear_tokens()?; - self.clear_api_key()?; + let mut credentials = self.credentials.write().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials write lock: {}", e)) + })?; + credentials.tokens = None; + credentials.api_key = None; + credentials.generation = credentials.generation.wrapping_add(1); + credentials.token_generation = credentials.token_generation.wrapping_add(1); + credentials.api_key_generation = credentials.api_key_generation.wrapping_add(1); Ok(()) } + + pub(crate) fn clear_all_if_generation(&self, expected_generation: u64) -> Result { + let mut credentials = self.credentials.write().map_err(|e| { + Error::Authentication(format!("Failed to acquire credentials write lock: {}", e)) + })?; + + if credentials.generation != expected_generation { + return Ok(false); + } + + self.clear_session()?; + credentials.tokens = None; + credentials.api_key = None; + credentials.generation = credentials.generation.wrapping_add(1); + credentials.token_generation = credentials.token_generation.wrapping_add(1); + credentials.api_key_generation = credentials.api_key_generation.wrapping_add(1); + Ok(true) + } } impl Default for SessionManager { From e13b794c69ba7d18a9ae3141eb1f0e7ca5978761 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Fri, 17 Jul 2026 07:03:19 +0000 Subject: [PATCH 3/7] fix(rust): enforce signed enclave identity --- rust/README.md | 29 ++- rust/src/attestation.rs | 6 + rust/src/client.rs | 170 ++++++++++++++- rust/src/lib.rs | 2 + rust/src/pcr.rs | 451 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 646 insertions(+), 12 deletions(-) create mode 100644 rust/src/pcr.rs diff --git a/rust/README.md b/rust/README.md index beed83c..01d203f 100644 --- a/rust/README.md +++ b/rust/README.md @@ -26,7 +26,7 @@ http = "1" ## Quick Start ```rust -use opensecret::{OpenSecretClient, Result}; +use opensecret::{OpenSecretClient, Pcr0TrustPolicy, Result}; use uuid::Uuid; #[tokio::main] @@ -51,6 +51,28 @@ async fn main() -> Result<()> { } ``` +Production clients verify both the AWS Nitro attestation and the enclave's +PCR0 deployment identity. `OpenSecretClient::new` uses pinned official PCR0 +values and OpenSecret's signed production and development histories. Custom +deployments can add a static allowlist without replacing official trust: + +```rust +let policy = Pcr0TrustPolicy::official().with_additional_pcr0s([ + "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", +])?; +let client = OpenSecretClient::new_with_pcr0_trust_policy( + "https://api.opensecret.cloud", + policy, +)?; +``` + +Use `Pcr0TrustPolicy::from_static_allowlist(...)` to disable remote history and +trust only an explicit custom set. Remote entries are size/time bounded and +must verify against the SDK's hardcoded OpenSecret P-384 signing key. Exact +localhost, loopback, and unspecified-address development endpoints continue to +use mock attestation; Android also supports the exact emulator alias +`10.0.2.2`. Other endpoints must use HTTPS. + ## Inference APIs `send_inference_request` is the lossless inference API. The caller owns the @@ -139,7 +161,10 @@ let response = client.login_with_id( Tokens are automatically stored after login/registration. You can: ```rust -// Get current tokens +// Get one coherent access/refresh pair snapshot +let tokens = client.get_tokens()?; + +// Individual reads remain available when a coherent pair is not required let access_token = client.get_access_token()?; let refresh_token = client.get_refresh_token()?; diff --git a/rust/src/attestation.rs b/rust/src/attestation.rs index 278ebe9..d5f545e 100644 --- a/rust/src/attestation.rs +++ b/rust/src/attestation.rs @@ -21,6 +21,12 @@ pub struct AttestationDocument { pub nonce: Option>, } +/// Low-level AWS Nitro document verifier. +/// +/// This verifies the certificate chain, document signature, and nonce. Nitro +/// authenticity alone does not identify an OpenSecret deployment. Production +/// callers should use `OpenSecretClient`, which additionally enforces its +/// configured `Pcr0TrustPolicy` before key exchange. pub struct AttestationVerifier { expected_pcrs: Option>>, allow_debug: bool, diff --git a/rust/src/client.rs b/rust/src/client.rs index 24640c7..bbbb587 100644 --- a/rust/src/client.rs +++ b/rust/src/client.rs @@ -3,6 +3,7 @@ use crate::{ cbor::{self, Value as CborValue}, crypto::{self}, error::{Error, Result}, + pcr::Pcr0TrustPolicy, session::SessionManager, types::*, }; @@ -16,7 +17,7 @@ use reqwest::{ Client, }; use serde::{de::DeserializeOwned, Deserialize, Serialize}; -use std::pin::Pin; +use std::{net::IpAddr, pin::Pin}; use tokio::sync::Mutex; use uuid::Uuid; @@ -47,6 +48,7 @@ pub struct OpenSecretClient { session_manager: SessionManager, refresh_lock: Mutex<()>, use_mock_attestation: bool, + pcr0_trust_policy: Pcr0TrustPolicy, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -376,13 +378,63 @@ async fn collect_response_body(mut body: OpenSecretResponseBody) -> Result Result { + let parsed = reqwest::Url::parse(base_url) + .map_err(|error| Error::Configuration(format!("Invalid base URL: {error}")))?; + if !matches!(parsed.scheme(), "http" | "https") { + return Err(Error::Configuration( + "Base URL must use HTTP or HTTPS".to_string(), + )); + } + if !parsed.username().is_empty() || parsed.password().is_some() { + return Err(Error::Configuration( + "Base URL must not contain credentials".to_string(), + )); + } + if parsed.query().is_some() || parsed.fragment().is_some() { + return Err(Error::Configuration( + "Base URL must not contain a query or fragment".to_string(), + )); + } + let host = parsed + .host_str() + .ok_or_else(|| Error::Configuration("Base URL must include a host".to_string()))?; + let host = host.trim_end_matches('.'); + + let is_mock_host = if host.eq_ignore_ascii_case("localhost") { + true + } else { + let address_host = host + .strip_prefix('[') + .and_then(|host| host.strip_suffix(']')) + .unwrap_or(host); + address_host.parse::().is_ok_and(|address| { + address.is_loopback() + || address.is_unspecified() + || (cfg!(target_os = "android") && address == IpAddr::from([10, 0, 2, 2])) + }) + }; + + if parsed.scheme() != "https" && !is_mock_host { + return Err(Error::Configuration( + "Non-local base URLs must use HTTPS".to_string(), + )); + } + Ok(is_mock_host) +} + impl OpenSecretClient { pub fn new(base_url: impl Into) -> Result { + Self::new_with_pcr0_trust_policy(base_url, Pcr0TrustPolicy::official()) + } + + /// Construct a client with an explicit production PCR0 trust policy. + pub fn new_with_pcr0_trust_policy( + base_url: impl Into, + pcr0_trust_policy: Pcr0TrustPolicy, + ) -> Result { let base_url = base_url.into(); - let use_mock = base_url.contains("localhost") - || base_url.contains("127.0.0.1") - || base_url.contains("0.0.0.0") - || base_url.contains("10.0.2.2"); + let use_mock = uses_mock_attestation(&base_url)?; Ok(Self { client: Client::new(), @@ -390,15 +442,22 @@ impl OpenSecretClient { session_manager: SessionManager::new(), refresh_lock: Mutex::new(()), use_mock_attestation: use_mock, + pcr0_trust_policy, }) } pub fn new_with_api_key(base_url: impl Into, api_key: String) -> Result { + Self::new_with_api_key_and_pcr0_trust_policy(base_url, api_key, Pcr0TrustPolicy::official()) + } + + /// Construct an API-key client with an explicit production PCR0 policy. + pub fn new_with_api_key_and_pcr0_trust_policy( + base_url: impl Into, + api_key: String, + pcr0_trust_policy: Pcr0TrustPolicy, + ) -> Result { let base_url = base_url.into(); - let use_mock = base_url.contains("localhost") - || base_url.contains("127.0.0.1") - || base_url.contains("0.0.0.0") - || base_url.contains("10.0.2.2"); + let use_mock = uses_mock_attestation(&base_url)?; Ok(Self { client: Client::new(), @@ -406,6 +465,7 @@ impl OpenSecretClient { session_manager: SessionManager::new_with_api_key(api_key), refresh_lock: Mutex::new(()), use_mock_attestation: use_mock, + pcr0_trust_policy, }) } @@ -427,7 +487,17 @@ impl OpenSecretClient { // Step 2: Parse and verify attestation document let doc = if !self.use_mock_attestation { let verifier = AttestationVerifier::new(); - verifier.verify_attestation_document(&attestation_doc.attestation_document, &nonce)? + let doc = verifier + .verify_attestation_document(&attestation_doc.attestation_document, &nonce)?; + let pcr0 = doc.pcrs.get(&0).ok_or_else(|| { + Error::AttestationVerificationFailed( + "Missing PCR0 in attestation document".to_string(), + ) + })?; + self.pcr0_trust_policy + .verify_pcr0(&self.client, pcr0) + .await?; + doc } else { // For mock mode, extract without full verification self.parse_mock_attestation(&attestation_doc.attestation_document)? @@ -1468,6 +1538,15 @@ impl OpenSecretClient { self.session_manager.get_access_token() } + /// Return one coherent snapshot of the current JWT token pair. + /// + /// Prefer this over separate access- and refresh-token reads when the pair + /// will be persisted or copied into another client: an automatic refresh + /// can replace both values between two independent lock acquisitions. + pub fn get_tokens(&self) -> Result> { + self.session_manager.get_tokens() + } + pub fn get_refresh_token(&self) -> Result> { self.session_manager.get_refresh_token() } @@ -2763,6 +2842,77 @@ mod tests { assert!(client.use_mock_attestation); } + #[test] + fn mock_attestation_uses_the_parsed_host_not_url_substrings() { + for url in [ + "https://localhost.example.com", + "https://example.com/localhost", + "https://example.com/127.0.0.1", + ] { + let client = OpenSecretClient::new(url).unwrap(); + assert!(!client.use_mock_attestation, "unexpected mock URL: {url}"); + } + + assert!( + OpenSecretClient::new("http://localhost:3000") + .unwrap() + .use_mock_attestation + ); + assert!( + OpenSecretClient::new("http://127.0.0.1:3000") + .unwrap() + .use_mock_attestation + ); + assert!( + OpenSecretClient::new("http://[::1]:3000") + .unwrap() + .use_mock_attestation + ); + } + + #[test] + fn base_url_validation_rejects_malformed_or_ambiguous_urls() { + for url in [ + "not a URL", + "file:///tmp/opensecret", + "https://localhost@example.com", + "https://example.com?redirect=localhost", + "https://example.com/#localhost", + ] { + assert!( + OpenSecretClient::new(url).is_err(), + "unexpectedly accepted base URL: {url}" + ); + } + } + + #[test] + fn android_emulator_alias_is_not_a_desktop_mock_bypass() { + let client = OpenSecretClient::new("http://10.0.2.2:3000"); + if cfg!(target_os = "android") { + assert!(client.unwrap().use_mock_attestation); + } else { + assert!(client.is_err()); + assert!( + !OpenSecretClient::new("https://10.0.2.2:3000") + .unwrap() + .use_mock_attestation + ); + } + } + + #[test] + fn get_tokens_returns_one_coherent_pair_snapshot() { + let client = OpenSecretClient::new("http://localhost:3000").unwrap(); + client + .set_tokens("access".to_string(), Some("refresh".to_string())) + .unwrap(); + + let tokens = client.get_tokens().unwrap().unwrap(); + assert_eq!(tokens.access_token, "access"); + assert_eq!(tokens.refresh_token.as_deref(), Some("refresh")); + } + #[tokio::test] async fn concurrent_attestation_handshakes_keep_each_nonce_public_key() { let mock_server = MockServer::start().await; diff --git a/rust/src/lib.rs b/rust/src/lib.rs index dd1859c..0dd922f 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -3,11 +3,13 @@ mod cbor; pub mod client; pub mod crypto; pub mod error; +pub mod pcr; pub mod push; pub mod session; pub mod types; pub use client::{InferenceRequest, InferenceResponse, OpenSecretClient, OpenSecretResponseBody}; pub use error::{Error, Result}; +pub use pcr::Pcr0TrustPolicy; pub use push::*; pub use types::*; diff --git a/rust/src/pcr.rs b/rust/src/pcr.rs new file mode 100644 index 0000000..ceec167 --- /dev/null +++ b/rust/src/pcr.rs @@ -0,0 +1,451 @@ +use crate::{error::Error, Result}; +use base64::{engine::general_purpose::STANDARD as BASE64, Engine}; +use reqwest::{Client, Url}; +use ring::signature; +use serde::Deserialize; +use std::collections::HashSet; +use std::time::Duration; + +const PCR0_HEX_LEN: usize = 96; +const PCR0_BYTES_LEN: usize = 48; +const MAX_REMOTE_HISTORY_BYTES: usize = 1024 * 1024; +const MAX_REMOTE_HISTORY_ENTRIES: usize = 2048; +const MAX_REMOTE_HISTORY_URLS: usize = 4; +const MAX_REMOTE_HISTORY_URL_BYTES: usize = 2048; +const REMOTE_HISTORY_TIMEOUT: Duration = Duration::from_secs(5); + +/// OpenSecret's P-384 PCR-history verification key in SPKI DER form. +/// +/// This key is the trust root for remote history entries. Replacing a history +/// URL cannot expand trust without a signature made by the corresponding +/// private key. +const PCR_HISTORY_VERIFICATION_KEY_B64: &str = + "MHYwEAYHKoZIzj0CAQYFK4EEACIDYgAEHiUY9kFWK1GqBGzczohhwEwElXzgWLDZa9R6wBx3JOBocgSt9+UIzZlJbPDjYeGBfDUXh7Z62BG2vVsh2NgclLB5S7A2ucBBtb1wd8vSQHP8jpdPhZX1slauPgbnROIP"; + +pub const OFFICIAL_PRODUCTION_PCR_HISTORY_URL: &str = + "https://raw.githubusercontent.com/OpenSecretCloud/opensecret/master/pcrProdHistory.json"; +pub const OFFICIAL_DEVELOPMENT_PCR_HISTORY_URL: &str = + "https://raw.githubusercontent.com/OpenSecretCloud/opensecret/master/pcrDevHistory.json"; + +const OFFICIAL_PRODUCTION_PCR0S: &[&str] = &[ + "eeddbb58f57c38894d6d5af5e575fbe791c5bf3bbcfb5df8da8cfcf0c2e1da1913108e6a762112444740b88c163d7f4b", + "74ed417f88cb0ca76c4a3d10f278bd010f1d3f95eafb254d4732511bb50e404507a4049b779c5230137e4091a5582271", + "9043fcab93b972d3c14ad2dc8fa78ca7ad374fc937c02435681772a003f7a72876bc4d578089b5c4cf3fe9b480f1aabb", + "52c3595b151d93d8b159c257301bfd5aa6f49210de0c55a6cd6df5ebeee44e4206cab950500f5d188f7fa14e6d900b75", + "91cb67311e910cce68cd5b7d0de77aa40610d87c6681439b44c46c3ff786ae643956ab2c812478a1da8745b259f07a45", + "859065ac81b81d3735130ba08b8af72a7256b603fefb74faabae25ed28cca6edcaa7c10ea32b5948d675c18a9b0f2b1d", + "acd82a7d3943e23e95a9dc3ce0b0107ea358d6287f9e3afa245622f7c7e3e0a66142a928b6efcc02f594a95366d3a99d", +]; + +const OFFICIAL_DEVELOPMENT_PCR0S: &[&str] = &[ + "62c0407056217a4c10764ed9045694c29fa93255d3cc04c2f989cdd9a1f8050c8b169714c71f1118ebce2fcc9951d1a9", + "cb95519905443f9f66f05f63c548b61ad1561a27fd5717b69285861aaea3c3063fe12a2571773b67fea3c6c11b4d8ec6", + "deb5895831b5e4286f5a2dcf5e9c27383821446f8df2b465f141d10743599be20ba3bb381ce063bf7139cc89f7f61d4c", + "70ba26c6af1ec3b57ce80e1adcc0ee96d70224d4c7a078f427895cdf68e1c30f09b5ac4c456588d872f3f21ff77c036b", + "669404ea71435b8f498b48db7816a5c2ab1d258b1a77685b11d84d15a73189504d79c4dee13a658de9f4a0cbfc39cfe8", + "a791bf92c25ffdfd372660e460a0e238c6778c090672df6509ae4bc065cf8668b6baac6b6a11d554af53ee0ff0172ad5", + "c4285443b87b9b12a6cea3bef1064ec060f652b235a297095975af8f134e5ed65f92d70d4616fdec80af9dff48bb9f35", +]; + +/// PCR0 deployment-identity policy enforced after Nitro document validation. +/// +/// The default policy trusts OpenSecret's pinned PCR0 values and falls back to +/// the signed official production and development histories. Use +/// [`Self::from_static_allowlist`] for a custom deployment that must not use +/// remote history. +#[derive(Debug, Clone)] +pub struct Pcr0TrustPolicy { + trusted_pcr0s: HashSet, + remote_history_urls: Vec, +} + +impl Pcr0TrustPolicy { + /// Return the official OpenSecret policy used by default. + pub fn official() -> Self { + let trusted_pcr0s = OFFICIAL_PRODUCTION_PCR0S + .iter() + .chain(OFFICIAL_DEVELOPMENT_PCR0S) + .map(|pcr0| (*pcr0).to_string()) + .collect(); + let remote_history_urls = [ + OFFICIAL_PRODUCTION_PCR_HISTORY_URL, + OFFICIAL_DEVELOPMENT_PCR_HISTORY_URL, + ] + .into_iter() + .map(|url| Url::parse(url).expect("official PCR history URL must be valid")) + .collect(); + + Self { + trusted_pcr0s, + remote_history_urls, + } + } + + /// Build a remote-disabled policy containing only caller-supplied PCR0s. + pub fn from_static_allowlist(pcr0s: I) -> Result + where + I: IntoIterator, + S: AsRef, + { + let mut policy = Self { + trusted_pcr0s: HashSet::new(), + remote_history_urls: Vec::new(), + }; + policy.add_pcr0s(pcr0s)?; + if policy.trusted_pcr0s.is_empty() { + return Err(Error::Configuration( + "PCR0 static allowlist must not be empty".to_string(), + )); + } + Ok(policy) + } + + /// Add caller-supplied PCR0 values to this policy. + pub fn with_additional_pcr0s(mut self, pcr0s: I) -> Result + where + I: IntoIterator, + S: AsRef, + { + self.add_pcr0s(pcr0s)?; + Ok(self) + } + + /// Disable signed remote history and retain only this policy's static set. + pub fn without_remote_history(mut self) -> Self { + self.remote_history_urls.clear(); + self + } + + /// Replace the default remote history locations. + /// + /// Every accepted entry must still verify against OpenSecret's hardcoded + /// signing key. HTTPS is required except for an exact loopback host, which + /// is allowed to support deterministic local testing and signed mirrors. + pub fn with_remote_history_urls(mut self, urls: I) -> Result + where + I: IntoIterator, + S: AsRef, + { + let parsed = urls + .into_iter() + .map(|url| parse_remote_history_url(url.as_ref())) + .collect::>>()?; + if parsed.is_empty() || parsed.len() > MAX_REMOTE_HISTORY_URLS { + return Err(Error::Configuration(format!( + "PCR history requires between 1 and {MAX_REMOTE_HISTORY_URLS} URLs" + ))); + } + self.remote_history_urls = parsed; + Ok(self) + } + + fn add_pcr0s(&mut self, pcr0s: I) -> Result<()> + where + I: IntoIterator, + S: AsRef, + { + for pcr0 in pcr0s { + let pcr0 = pcr0.as_ref(); + validate_pcr0_hex(pcr0)?; + self.trusted_pcr0s.insert(pcr0.to_string()); + } + Ok(()) + } + + pub(crate) async fn verify_pcr0(&self, client: &Client, pcr0: &[u8]) -> Result<()> { + if pcr0.len() != PCR0_BYTES_LEN { + return Err(Error::AttestationVerificationFailed(format!( + "PCR0 must be {PCR0_BYTES_LEN} bytes" + ))); + } + let pcr0_hex = hex::encode(pcr0); + if self.trusted_pcr0s.contains(&pcr0_hex) { + return Ok(()); + } + + for url in &self.remote_history_urls { + let history = match fetch_remote_history(client, url).await { + Ok(history) => history, + Err(_) => continue, + }; + if history + .iter() + .any(|entry| entry.pcr0 == pcr0_hex && entry.has_valid_signature()) + { + return Ok(()); + } + } + + Err(Error::AttestationVerificationFailed( + "PCR0 is not approved by the configured trust policy".to_string(), + )) + } +} + +impl Default for Pcr0TrustPolicy { + fn default() -> Self { + Self::official() + } +} + +#[derive(Debug, Deserialize)] +struct PcrHistoryEntry { + #[serde(rename = "PCR0")] + pcr0: String, + #[serde(rename = "PCR1")] + pcr1: String, + #[serde(rename = "PCR2")] + pcr2: String, + timestamp: u64, + signature: String, +} + +impl PcrHistoryEntry { + fn validate(&self) -> Result<()> { + validate_pcr0_hex(&self.pcr0)?; + validate_pcr0_hex(&self.pcr1)?; + validate_pcr0_hex(&self.pcr2)?; + if self.timestamp == 0 { + return Err(Error::AttestationVerificationFailed( + "PCR history timestamp must be nonzero".to_string(), + )); + } + let signature = BASE64.decode(&self.signature)?; + if signature.len() != 96 { + return Err(Error::AttestationVerificationFailed( + "PCR history signature must be 96 bytes".to_string(), + )); + } + Ok(()) + } + + fn has_valid_signature(&self) -> bool { + let Ok(signature_bytes) = BASE64.decode(&self.signature) else { + return false; + }; + let Ok(spki) = BASE64.decode(PCR_HISTORY_VERIFICATION_KEY_B64) else { + return false; + }; + let Some(public_key) = spki.get(spki.len().saturating_sub(97)..) else { + return false; + }; + if public_key.first() != Some(&0x04) { + return false; + } + + signature::UnparsedPublicKey::new(&signature::ECDSA_P384_SHA384_FIXED, public_key) + .verify(self.pcr0.as_bytes(), &signature_bytes) + .is_ok() + } +} + +async fn fetch_remote_history(client: &Client, url: &Url) -> Result> { + tokio::time::timeout( + REMOTE_HISTORY_TIMEOUT, + fetch_remote_history_inner(client, url), + ) + .await + .map_err(|_| { + Error::AttestationVerificationFailed("PCR history request timed out".to_string()) + })? +} + +async fn fetch_remote_history_inner(client: &Client, url: &Url) -> Result> { + let mut response = client.get(url.clone()).send().await?; + if !response.status().is_success() { + return Err(Error::AttestationVerificationFailed( + "PCR history request failed".to_string(), + )); + } + if response + .content_length() + .is_some_and(|length| length > MAX_REMOTE_HISTORY_BYTES as u64) + { + return Err(Error::AttestationVerificationFailed( + "PCR history response is too large".to_string(), + )); + } + + let mut body = Vec::new(); + while let Some(chunk) = response.chunk().await? { + if body.len().saturating_add(chunk.len()) > MAX_REMOTE_HISTORY_BYTES { + return Err(Error::AttestationVerificationFailed( + "PCR history response is too large".to_string(), + )); + } + body.extend_from_slice(&chunk); + } + + let entries: Vec = serde_json::from_slice(&body)?; + if entries.is_empty() || entries.len() > MAX_REMOTE_HISTORY_ENTRIES { + return Err(Error::AttestationVerificationFailed(format!( + "PCR history requires between 1 and {MAX_REMOTE_HISTORY_ENTRIES} entries" + ))); + } + for entry in &entries { + entry.validate()?; + } + Ok(entries) +} + +fn validate_pcr0_hex(pcr0: &str) -> Result<()> { + if pcr0.len() != PCR0_HEX_LEN + || !pcr0 + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + { + return Err(Error::Configuration( + "PCR0 values must be 96 lowercase hexadecimal characters".to_string(), + )); + } + Ok(()) +} + +fn parse_remote_history_url(value: &str) -> Result { + if value.len() > MAX_REMOTE_HISTORY_URL_BYTES { + return Err(Error::Configuration( + "PCR history URL is too long".to_string(), + )); + } + let url = Url::parse(value) + .map_err(|error| Error::Configuration(format!("Invalid PCR history URL: {error}")))?; + if !url.username().is_empty() || url.password().is_some() || url.fragment().is_some() { + return Err(Error::Configuration( + "PCR history URL must not contain credentials or a fragment".to_string(), + )); + } + let is_loopback = url.host_str().is_some_and(|host| { + let host = host.trim_end_matches('.'); + let address_host = host + .strip_prefix('[') + .and_then(|host| host.strip_suffix(']')) + .unwrap_or(host); + host.eq_ignore_ascii_case("localhost") + || address_host + .parse::() + .is_ok_and(|address| address.is_loopback()) + }); + if url.scheme() != "https" && !(url.scheme() == "http" && is_loopback) { + return Err(Error::Configuration( + "PCR history URL must use HTTPS (HTTP is allowed only for loopback)".to_string(), + )); + } + Ok(url) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use wiremock::{matchers::path, Mock, MockServer, ResponseTemplate}; + + const SIGNED_PCR0: &str = + "3637534c33a8bafc5034d5763e441a481f161bbbe888e375ce14b016c7497dc4e550afe866bd8e65969b409d54766481"; + const SIGNED_PCR0_SIGNATURE: &str = + "GZTXC0Xt0+yAaAatmMUd37pUJpF0nRAOj3Df9qxDOvDvRkiTF8UbGlzlL4kIOi/nd7dXAaEqYnY7OlpyngHBED2CSTpRRwV0xGo109epfqUKWWudrFaXpMsJ+GRKJLFO"; + + fn pcr_bytes(value: &str) -> Vec { + hex::decode(value).unwrap() + } + + fn history(signature: &str) -> serde_json::Value { + json!([{ + "PCR0": SIGNED_PCR0, + "PCR1": "e45de6f4e9809176f6adc68df999f87f32a602361247d5819d1edf11ac5a403cfbb609943705844251af85713a17c83a", + "PCR2": "fe0a6f7c29c7c4999571869f880b6d5086b377deaaf359e19ae824edacd6a9d90247b793a5f2d73c0e74e2f9630aeb4a", + "timestamp": 1743710235_u64, + "signature": signature, + "futureMetadata": { "release": "ignored" }, + }]) + } + + #[tokio::test] + async fn static_allowlist_approves_only_exact_pcr0() { + let policy = Pcr0TrustPolicy::from_static_allowlist([SIGNED_PCR0]).unwrap(); + policy + .verify_pcr0(&Client::new(), &pcr_bytes(SIGNED_PCR0)) + .await + .unwrap(); + + let error = policy + .verify_pcr0(&Client::new(), &[0x42; PCR0_BYTES_LEN]) + .await + .unwrap_err(); + assert!(matches!(error, Error::AttestationVerificationFailed(_))); + } + + #[tokio::test] + async fn signed_remote_history_approves_matching_pcr0() { + let server = MockServer::start().await; + Mock::given(path("/history.json")) + .respond_with(ResponseTemplate::new(200).set_body_json(history(SIGNED_PCR0_SIGNATURE))) + .expect(1) + .mount(&server) + .await; + let policy = Pcr0TrustPolicy::from_static_allowlist(["000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"]) + .unwrap() + .with_remote_history_urls([format!("{}/history.json", server.uri())]) + .unwrap(); + + policy + .verify_pcr0(&Client::new(), &pcr_bytes(SIGNED_PCR0)) + .await + .unwrap(); + } + + #[tokio::test] + async fn invalid_remote_signature_fails_closed() { + let server = MockServer::start().await; + let invalid_signature = BASE64.encode([0u8; 96]); + Mock::given(path("/history.json")) + .respond_with(ResponseTemplate::new(200).set_body_json(history(&invalid_signature))) + .expect(1) + .mount(&server) + .await; + let policy = Pcr0TrustPolicy::from_static_allowlist(["000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"]) + .unwrap() + .with_remote_history_urls([format!("{}/history.json", server.uri())]) + .unwrap(); + + let error = policy + .verify_pcr0(&Client::new(), &pcr_bytes(SIGNED_PCR0)) + .await + .unwrap_err(); + assert!(matches!(error, Error::AttestationVerificationFailed(_))); + } + + #[tokio::test] + async fn malformed_remote_history_fails_closed() { + let server = MockServer::start().await; + Mock::given(path("/history.json")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!([{ + "PCR0": "too-short", + "PCR1": "too-short", + "PCR2": "too-short", + "timestamp": 1, + "signature": SIGNED_PCR0_SIGNATURE, + }]))) + .expect(1) + .mount(&server) + .await; + let policy = Pcr0TrustPolicy::from_static_allowlist(["000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"]) + .unwrap() + .with_remote_history_urls([format!("{}/history.json", server.uri())]) + .unwrap(); + + assert!(policy + .verify_pcr0(&Client::new(), &pcr_bytes(SIGNED_PCR0)) + .await + .is_err()); + } + + #[test] + fn rejects_unsafe_remote_history_urls() { + assert!(Pcr0TrustPolicy::official() + .with_remote_history_urls(["http://example.com/history.json"]) + .is_err()); + assert!(Pcr0TrustPolicy::official() + .with_remote_history_urls(["https://user@example.com/history.json"]) + .is_err()); + } +} From 37273121e2d1cb8f21c67e211ad95b13088fb861 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Fri, 17 Jul 2026 07:08:20 +0000 Subject: [PATCH 4/7] fix(rust): reject PCR history redirects --- rust/src/client.rs | 4 +--- rust/src/pcr.rs | 51 ++++++++++++++++++++++++++++++---------------- 2 files changed, 35 insertions(+), 20 deletions(-) diff --git a/rust/src/client.rs b/rust/src/client.rs index bbbb587..604826d 100644 --- a/rust/src/client.rs +++ b/rust/src/client.rs @@ -494,9 +494,7 @@ impl OpenSecretClient { "Missing PCR0 in attestation document".to_string(), ) })?; - self.pcr0_trust_policy - .verify_pcr0(&self.client, pcr0) - .await?; + self.pcr0_trust_policy.verify_pcr0(pcr0).await?; doc } else { // For mock mode, extract without full verification diff --git a/rust/src/pcr.rs b/rust/src/pcr.rs index ceec167..0e79e7e 100644 --- a/rust/src/pcr.rs +++ b/rust/src/pcr.rs @@ -1,6 +1,6 @@ use crate::{error::Error, Result}; use base64::{engine::general_purpose::STANDARD as BASE64, Engine}; -use reqwest::{Client, Url}; +use reqwest::{redirect::Policy, Client, Url}; use ring::signature; use serde::Deserialize; use std::collections::HashSet; @@ -152,7 +152,7 @@ impl Pcr0TrustPolicy { Ok(()) } - pub(crate) async fn verify_pcr0(&self, client: &Client, pcr0: &[u8]) -> Result<()> { + pub(crate) async fn verify_pcr0(&self, pcr0: &[u8]) -> Result<()> { if pcr0.len() != PCR0_BYTES_LEN { return Err(Error::AttestationVerificationFailed(format!( "PCR0 must be {PCR0_BYTES_LEN} bytes" @@ -163,8 +163,12 @@ impl Pcr0TrustPolicy { return Ok(()); } + // URL validation applies to the complete network request. Do not let a + // valid HTTPS history location redirect to HTTP, loopback, or another + // destination that has not passed `parse_remote_history_url`. + let history_client = Client::builder().redirect(Policy::none()).build()?; for url in &self.remote_history_urls { - let history = match fetch_remote_history(client, url).await { + let history = match fetch_remote_history(&history_client, url).await { Ok(history) => history, Err(_) => continue, }; @@ -362,13 +366,10 @@ mod tests { #[tokio::test] async fn static_allowlist_approves_only_exact_pcr0() { let policy = Pcr0TrustPolicy::from_static_allowlist([SIGNED_PCR0]).unwrap(); - policy - .verify_pcr0(&Client::new(), &pcr_bytes(SIGNED_PCR0)) - .await - .unwrap(); + policy.verify_pcr0(&pcr_bytes(SIGNED_PCR0)).await.unwrap(); let error = policy - .verify_pcr0(&Client::new(), &[0x42; PCR0_BYTES_LEN]) + .verify_pcr0(&[0x42; PCR0_BYTES_LEN]) .await .unwrap_err(); assert!(matches!(error, Error::AttestationVerificationFailed(_))); @@ -387,10 +388,7 @@ mod tests { .with_remote_history_urls([format!("{}/history.json", server.uri())]) .unwrap(); - policy - .verify_pcr0(&Client::new(), &pcr_bytes(SIGNED_PCR0)) - .await - .unwrap(); + policy.verify_pcr0(&pcr_bytes(SIGNED_PCR0)).await.unwrap(); } #[tokio::test] @@ -408,7 +406,7 @@ mod tests { .unwrap(); let error = policy - .verify_pcr0(&Client::new(), &pcr_bytes(SIGNED_PCR0)) + .verify_pcr0(&pcr_bytes(SIGNED_PCR0)) .await .unwrap_err(); assert!(matches!(error, Error::AttestationVerificationFailed(_))); @@ -433,10 +431,29 @@ mod tests { .with_remote_history_urls([format!("{}/history.json", server.uri())]) .unwrap(); - assert!(policy - .verify_pcr0(&Client::new(), &pcr_bytes(SIGNED_PCR0)) - .await - .is_err()); + assert!(policy.verify_pcr0(&pcr_bytes(SIGNED_PCR0)).await.is_err()); + } + + #[tokio::test] + async fn remote_history_redirects_are_not_followed() { + let server = MockServer::start().await; + let redirected_url = format!("{}/redirected.json", server.uri()); + Mock::given(path("/history.json")) + .respond_with(ResponseTemplate::new(302).insert_header("location", redirected_url)) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/redirected.json")) + .respond_with(ResponseTemplate::new(200).set_body_json(history(SIGNED_PCR0_SIGNATURE))) + .expect(0) + .mount(&server) + .await; + let policy = Pcr0TrustPolicy::from_static_allowlist(["000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"]) + .unwrap() + .with_remote_history_urls([format!("{}/history.json", server.uri())]) + .unwrap(); + + assert!(policy.verify_pcr0(&pcr_bytes(SIGNED_PCR0)).await.is_err()); } #[test] From 07f8e53df465763a9a11d839a8aeaec6f9a0302e Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Fri, 17 Jul 2026 07:38:07 +0000 Subject: [PATCH 5/7] feat: add web search and extraction APIs --- rust/src/client.rs | 227 +++++++++++++++++++++++++++ rust/src/types.rs | 191 ++++++++++++++++++++++ src/lib/api.ts | 103 ++++++++++++ src/lib/index.ts | 14 ++ src/lib/main.tsx | 10 ++ src/lib/test/integration/web.test.ts | 156 ++++++++++++++++++ 6 files changed, 701 insertions(+) create mode 100644 src/lib/test/integration/web.test.ts diff --git a/rust/src/client.rs b/rust/src/client.rs index 604826d..395e61c 100644 --- a/rust/src/client.rs +++ b/rust/src/client.rs @@ -2167,6 +2167,20 @@ impl OpenSecretClient { Ok(Box::pin(event_stream)) } + // Web API Methods + + /// Searches the public web through OpenSecret's configured search provider. + pub async fn web_search(&self, request: WebSearchRequest) -> Result { + self.authenticated_api_call("/v1/web/search", "POST", Some(request)) + .await + } + + /// Extracts sanitized Markdown from public URLs through OpenSecret's configured provider. + pub async fn web_extract(&self, request: WebExtractRequest) -> Result { + self.authenticated_api_call("/v1/web/extract", "POST", Some(request)) + .await + } + async fn agent_chat_stream( &self, endpoint: String, @@ -4593,6 +4607,219 @@ mod tests { assert!(client.get_session_id().unwrap().is_none()); } + #[tokio::test] + async fn web_search_uses_authenticated_encrypted_endpoint() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [51u8; 32]; + + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + client + .session_manager + .set_tokens( + "web_access_token".to_string(), + Some("web_refresh_token".to_string()), + ) + .unwrap(); + + let request = WebSearchRequest { + query: "rust confidential computing".to_string(), + workflow: Some(WebSearchWorkflow::News), + page: Some(2), + limit: Some(25), + safe_search: Some(false), + timeout: Some(2.5), + lens_id: None, + lens: Some(WebSearchLens { + sites_included: Some(vec!["example.com".to_string()]), + keywords_included: Some(vec!["enclave".to_string()]), + time_relative: Some(WebSearchTimeRelative::Week), + search_region: Some("US".to_string()), + ..Default::default() + }), + filters: Some(WebSearchFilters { + region: Some("US".to_string()), + after: None, + before: None, + }), + }; + let response = WebSearchResponse { + trace_id: Some("trace-search-1".to_string()), + results: vec![WebSearchResult { + category: "news".to_string(), + url: "https://example.com/enclave".to_string(), + title: "Enclave update".to_string(), + snippet: Some("A short description.".to_string()), + published_at: Some("2026-07-16T12:00:00Z".to_string()), + }], + }; + + Mock::given(method("POST")) + .and(path("/v1/web/search")) + .and(header("authorization", "Bearer web_access_token")) + .and(header("x-session-id", session_id.to_string())) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ + "query": "rust confidential computing", + "workflow": "news", + "page": 2, + "limit": 25, + "safe_search": false, + "timeout": 2.5, + "lens": { + "sites_included": ["example.com"], + "keywords_included": ["enclave"], + "time_relative": "week", + "search_region": "US" + }, + "filters": { + "region": "US" + } + }), + }) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(encrypted_response(&session_key, &response)), + ) + .expect(1) + .mount(&mock_server) + .await; + + let actual = client.web_search(request).await.unwrap(); + + assert_eq!(actual, response); + mock_server.verify().await; + } + + #[tokio::test] + async fn web_extract_preserves_order_and_partial_page_errors() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [52u8; 32]; + + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + client + .session_manager + .set_tokens( + "web_access_token".to_string(), + Some("web_refresh_token".to_string()), + ) + .unwrap(); + + let first_url = "https://example.com/first".to_string(); + let second_url = "https://example.com/second".to_string(); + let request = WebExtractRequest { + urls: vec![first_url.clone(), second_url.clone()], + timeout: Some(4.5), + }; + let response = WebExtractResponse { + trace_id: Some("trace-extract-1".to_string()), + pages: vec![ + WebExtractPage { + url: first_url.clone(), + markdown: Some("# First\n\nExtracted text.".to_string()), + error: None, + }, + WebExtractPage { + url: second_url.clone(), + markdown: None, + error: Some(WebExtractPageError { + code: "no_content".to_string(), + message: "No readable content was found.".to_string(), + }), + }, + ], + }; + + Mock::given(method("POST")) + .and(path("/v1/web/extract")) + .and(header("authorization", "Bearer web_access_token")) + .and(header("x-session-id", session_id.to_string())) + .and(EncryptedJsonBodyMatcher { + session_key, + expected: json!({ + "urls": [first_url, second_url], + "timeout": 4.5 + }), + }) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(encrypted_response(&session_key, &response)), + ) + .expect(1) + .mount(&mock_server) + .await; + + let actual = client.web_extract(request).await.unwrap(); + + assert_eq!(actual, response); + assert_eq!(actual.pages[0].url, "https://example.com/first"); + assert_eq!( + actual.pages[1] + .error + .as_ref() + .map(|error| error.code.as_str()), + Some("no_content") + ); + mock_server.verify().await; + } + + #[tokio::test] + async fn web_validation_error_does_not_retry_attestation() { + let mock_server = MockServer::start().await; + let client = OpenSecretClient::new(mock_server.uri()).unwrap(); + let session_id = Uuid::new_v4(); + let session_key = [53u8; 32]; + + client + .session_manager + .set_session(session_id, session_key) + .unwrap(); + client + .session_manager + .set_tokens( + "web_access_token".to_string(), + Some("web_refresh_token".to_string()), + ) + .unwrap(); + + Mock::given(method("POST")) + .and(path("/v1/web/search")) + .and(header("authorization", "Bearer web_access_token")) + .and(header("x-session-id", session_id.to_string())) + .respond_with(ResponseTemplate::new(422).set_body_json(json!({ + "status": 422, + "code": "invalid_request", + "message": "The web request is invalid." + }))) + .expect(1) + .mount(&mock_server) + .await; + + let error = client + .web_search(WebSearchRequest::new("maple privacy")) + .await + .unwrap_err(); + + match error { + Error::Api { status, message } => { + assert_eq!(status, 422); + assert!(message.contains("invalid_request")); + } + other => panic!("expected API validation error, got {other:?}"), + } + mock_server.verify().await; + } + #[tokio::test] async fn test_init_main_agent_uses_authenticated_encrypted_v1_endpoint() { let mock_server = MockServer::start().await; diff --git a/rust/src/types.rs b/rust/src/types.rs index 8159059..5a8c498 100644 --- a/rust/src/types.rs +++ b/rust/src/types.rs @@ -890,6 +890,159 @@ pub struct EmbeddingUsage { pub total_tokens: i32, } +// Web API Types + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WebSearchWorkflow { + #[default] + Search, + Images, + Videos, + News, + Podcasts, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WebSearchTimeRelative { + Day, + Week, + Month, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct WebSearchLens { + #[serde(skip_serializing_if = "Option::is_none")] + pub sites_included: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub sites_excluded: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub keywords_included: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub keywords_excluded: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub file_type: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub time_after: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub time_before: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub time_relative: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub search_region: Option, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct WebSearchFilters { + #[serde(skip_serializing_if = "Option::is_none")] + pub region: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub after: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub before: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct WebSearchRequest { + pub query: String, + /// Search result class. The server defaults to [`WebSearchWorkflow::Search`]. + #[serde(skip_serializing_if = "Option::is_none")] + pub workflow: Option, + /// One-based result page, from 1 through 10. + #[serde(skip_serializing_if = "Option::is_none")] + pub page: Option, + /// Maximum results to return, from 1 through 50. The server defaults to 10. + #[serde(skip_serializing_if = "Option::is_none")] + pub limit: Option, + /// Whether to omit potentially unsafe content. The server defaults to true. + #[serde(skip_serializing_if = "Option::is_none")] + pub safe_search: Option, + /// Search collection timeout in seconds, from 0.5 through 4. + #[serde(skip_serializing_if = "Option::is_none")] + pub timeout: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub lens_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub lens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub filters: Option, +} + +impl WebSearchRequest { + pub fn new(query: impl Into) -> Self { + Self { + query: query.into(), + workflow: None, + page: None, + limit: None, + safe_search: None, + timeout: None, + lens_id: None, + lens: None, + filters: None, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct WebSearchResult { + pub category: String, + pub url: String, + pub title: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub snippet: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub published_at: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct WebSearchResponse { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub trace_id: Option, + pub results: Vec, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct WebExtractRequest { + /// Public HTTPS URLs to extract, from 1 through 10 entries. + pub urls: Vec, + /// Bulk extraction timeout in seconds, from 0.5 through 10. + #[serde(skip_serializing_if = "Option::is_none")] + pub timeout: Option, +} + +impl WebExtractRequest { + pub fn new(urls: impl IntoIterator>) -> Self { + Self { + urls: urls.into_iter().map(Into::into).collect(), + timeout: None, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct WebExtractPageError { + pub code: String, + pub message: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct WebExtractPage { + pub url: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub markdown: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct WebExtractResponse { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub trace_id: Option, + pub pages: Vec, +} + // Agent API Types #[derive(Debug, Clone, Serialize, Deserialize)] @@ -1168,6 +1321,44 @@ mod tests { assert_eq!(response.refresh_token, None); } + #[test] + fn web_requests_omit_unspecified_options() { + let search = WebSearchRequest::new("maple privacy"); + assert_eq!( + serde_json::to_value(search).unwrap(), + json!({ "query": "maple privacy" }) + ); + + let extract = WebExtractRequest::new(["https://example.com/article"]); + assert_eq!( + serde_json::to_value(extract).unwrap(), + json!({ "urls": ["https://example.com/article"] }) + ); + } + + #[test] + fn web_responses_tolerate_absent_optional_fields() { + let search: WebSearchResponse = serde_json::from_value(json!({ + "results": [{ + "category": "search", + "url": "https://example.com", + "title": "Example" + }] + })) + .unwrap(); + assert_eq!(search.trace_id, None); + assert_eq!(search.results[0].snippet, None); + assert_eq!(search.results[0].published_at, None); + + let extract: WebExtractResponse = serde_json::from_value(json!({ + "pages": [{ "url": "https://example.com" }] + })) + .unwrap(); + assert_eq!(extract.trace_id, None); + assert_eq!(extract.pages[0].markdown, None); + assert_eq!(extract.pages[0].error, None); + } + #[test] fn embedding_response_deserializes_float_vectors() { let response: EmbeddingResponse = serde_json::from_value(json!({ diff --git a/src/lib/api.ts b/src/lib/api.ts index a70e7c1..0438287 100644 --- a/src/lib/api.ts +++ b/src/lib/api.ts @@ -2801,6 +2801,109 @@ export async function setDefaultInstruction(instructionId: string): Promise { + return authenticatedApiCall( + `${apiUrl}/v1/web/search`, + "POST", + request, + "Failed to search the web" + ); +} + +/** Extracts sanitized Markdown from public URLs through OpenSecret's configured provider. */ +export async function webExtract(request: WebExtractRequest): Promise { + return authenticatedApiCall( + `${apiUrl}/v1/web/extract`, + "POST", + request, + "Failed to extract web pages" + ); +} + // ============================================================================ // Agent API Types // ============================================================================ diff --git a/src/lib/index.ts b/src/lib/index.ts index 7c0a793..c393014 100644 --- a/src/lib/index.ts +++ b/src/lib/index.ts @@ -48,6 +48,17 @@ export type { ConversationProjectUpdateRequest, ConversationProjectListParams, ConversationProjectDeleteResponse, + WebSearchWorkflow, + WebSearchTimeRelative, + WebSearchLens, + WebSearchFilters, + WebSearchRequest, + WebSearchResult, + WebSearchResponse, + WebExtractRequest, + WebExtractPageError, + WebExtractPage, + WebExtractResponse, AgentCreatedBy, MainAgentResponse, CreateSubagentRequest, @@ -85,6 +96,9 @@ export { deleteConversationProject } from "./api"; +// Export provider-neutral web API functions +export { webSearch, webExtract } from "./api"; + // Export Agent API functions export { getMainAgent, diff --git a/src/lib/main.tsx b/src/lib/main.tsx index 0752e16..ec53316 100644 --- a/src/lib/main.tsx +++ b/src/lib/main.tsx @@ -660,6 +660,12 @@ export type OpenSecretContextType = { */ transcribeAudio: typeof api.transcribeAudio; + /** Searches the public web through OpenSecret's configured search provider. */ + webSearch: typeof api.webSearch; + + /** Extracts sanitized Markdown from public URLs through OpenSecret's configured provider. */ + webExtract: typeof api.webExtract; + /** * Lists user's responses with pagination * @param params - Optional parameters for pagination and filtering @@ -955,6 +961,8 @@ export const OpenSecretContext = createContext({ listApiKeys: api.listApiKeys, deleteApiKey: api.deleteApiKey, transcribeAudio: api.transcribeAudio, + webSearch: api.webSearch, + webExtract: api.webExtract, fetchResponsesList: api.fetchResponsesList, fetchResponse: api.fetchResponse, cancelResponse: api.cancelResponse, @@ -1373,6 +1381,8 @@ export function OpenSecretProvider({ listApiKeys: api.listApiKeys, deleteApiKey: api.deleteApiKey, transcribeAudio: api.transcribeAudio, + webSearch: api.webSearch, + webExtract: api.webExtract, fetchResponsesList: api.fetchResponsesList, fetchResponse: api.fetchResponse, cancelResponse: api.cancelResponse, diff --git a/src/lib/test/integration/web.test.ts b/src/lib/test/integration/web.test.ts new file mode 100644 index 0000000..3b43071 --- /dev/null +++ b/src/lib/test/integration/web.test.ts @@ -0,0 +1,156 @@ +import { afterEach, beforeEach, expect, mock, test } from "bun:test"; +import { encode } from "@stablelib/base64"; +import { decryptMessage, encryptMessage } from "../../encryption"; +import { + setApiUrl, + webExtract, + webSearch, + type WebExtractRequest, + type WebExtractResponse, + type WebSearchRequest, + type WebSearchResponse +} from "../../api"; + +const apiUrl = "https://api.example.com"; +const accessToken = "web-access-token"; +const sessionId = "web-session-id"; +const sessionKey = new Uint8Array(32).fill(19); +const originalFetch = globalThis.fetch; + +beforeEach(() => { + window.localStorage.clear(); + window.sessionStorage.clear(); + window.localStorage.setItem("access_token", accessToken); + window.sessionStorage.setItem("sessionKey", encode(sessionKey)); + window.sessionStorage.setItem("sessionId", sessionId); + setApiUrl(apiUrl); +}); + +afterEach(() => { + globalThis.fetch = originalFetch; +}); + +test("webSearch sends an authenticated encrypted request and decrypts results", async () => { + const request: WebSearchRequest = { + query: "rust confidential computing", + workflow: "news", + page: 2, + limit: 25, + safe_search: false, + timeout: 2.5, + lens: { + sites_included: ["example.com"], + keywords_included: ["enclave"], + time_relative: "week", + search_region: "US" + }, + filters: { + region: "US" + } + }; + const response: WebSearchResponse = { + trace_id: "trace-search-1", + results: [ + { + category: "news", + url: "https://example.com/enclave", + title: "Enclave update", + snippet: "A short description.", + published_at: "2026-07-16T12:00:00Z" + } + ] + }; + + globalThis.fetch = mock(async (input: string | URL | Request, init?: RequestInit) => { + expect(input.toString()).toBe(`${apiUrl}/v1/web/search`); + expect(init?.method).toBe("POST"); + expect(init?.headers).toMatchObject({ + Authorization: `Bearer ${accessToken}`, + "x-session-id": sessionId + }); + + const body = JSON.parse(String(init?.body)) as { encrypted: string }; + expect(JSON.parse(decryptMessage(sessionKey, body.encrypted))).toEqual(request); + + return new Response( + JSON.stringify({ encrypted: encryptMessage(sessionKey, JSON.stringify(response)) }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + }) as typeof fetch; + + await expect(webSearch(request)).resolves.toEqual(response); +}); + +test("webExtract preserves ordered pages and typed partial failures", async () => { + const request: WebExtractRequest = { + urls: ["https://example.com/first", "https://example.com/second"], + timeout: 4.5 + }; + const response: WebExtractResponse = { + trace_id: "trace-extract-1", + pages: [ + { + url: request.urls[0], + markdown: "# First\n\nExtracted text." + }, + { + url: request.urls[1], + error: { + code: "no_content", + message: "No readable content was found." + } + } + ] + }; + + globalThis.fetch = mock(async (input: string | URL | Request, init?: RequestInit) => { + expect(input.toString()).toBe(`${apiUrl}/v1/web/extract`); + expect(init?.method).toBe("POST"); + expect(init?.headers).toMatchObject({ + Authorization: `Bearer ${accessToken}`, + "x-session-id": sessionId + }); + + const body = JSON.parse(String(init?.body)) as { encrypted: string }; + expect(JSON.parse(decryptMessage(sessionKey, body.encrypted))).toEqual(request); + + return new Response( + JSON.stringify({ encrypted: encryptMessage(sessionKey, JSON.stringify(response)) }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + }) as typeof fetch; + + const result = await webExtract(request); + + expect(result).toEqual(response); + expect(result.pages.map((page) => page.url)).toEqual(request.urls); + expect(result.pages[1].error?.code).toBe("no_content"); +}); + +test("web validation errors surface without an attestation retry", async () => { + let requestCount = 0; + + globalThis.fetch = mock(async (input: string | URL | Request) => { + requestCount += 1; + expect(input.toString()).toBe(`${apiUrl}/v1/web/search`); + + return new Response( + JSON.stringify({ + status: 422, + code: "invalid_request", + message: "The web request is invalid." + }), + { status: 422, headers: { "Content-Type": "application/json" } } + ); + }) as typeof fetch; + + try { + await webSearch({ query: "maple privacy", limit: 51 }); + throw new Error("expected webSearch to reject invalid input"); + } catch (error) { + expect(error).toBeInstanceOf(Error); + expect((error as Error).message).toBe("The web request is invalid."); + } + + expect(requestCount).toBe(1); +}); From 2ba4593304b6ab70fa607a9ab3f48ea014e5004c Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Fri, 17 Jul 2026 07:51:26 +0000 Subject: [PATCH 6/7] test: cover live web API integration --- rust/tests/web_integration.rs | 73 +++++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 rust/tests/web_integration.rs diff --git a/rust/tests/web_integration.rs b/rust/tests/web_integration.rs new file mode 100644 index 0000000..08debc4 --- /dev/null +++ b/rust/tests/web_integration.rs @@ -0,0 +1,73 @@ +use opensecret::{OpenSecretClient, Result, WebExtractRequest, WebSearchRequest}; +use uuid::Uuid; + +async fn authenticated_client() -> Result { + let base_url = std::env::var("VITE_OPEN_SECRET_API_URL") + .unwrap_or_else(|_| "http://localhost:3000".to_string()); + let email = std::env::var("VITE_TEST_EMAIL").expect("VITE_TEST_EMAIL must be set"); + let password = std::env::var("VITE_TEST_PASSWORD").expect("VITE_TEST_PASSWORD must be set"); + let client_id = std::env::var("VITE_TEST_CLIENT_ID") + .expect("VITE_TEST_CLIENT_ID must be set") + .parse::() + .expect("VITE_TEST_CLIENT_ID must be a UUID"); + + let client = OpenSecretClient::new(base_url)?; + client.perform_attestation_handshake().await?; + + if client + .login(email.clone(), password.clone(), client_id) + .await + .is_err() + { + client + .register(email, password, client_id, Some("Web API Test".to_string())) + .await?; + } + + Ok(client) +} + +#[tokio::test] +#[ignore = "Requires a live OpenSecret backend and Kagi API access"] +async fn live_web_search_and_extract() -> Result<()> { + let client = authenticated_client().await?; + + let mut search_request = WebSearchRequest::new( + std::env::var("VITE_TEST_WEB_QUERY") + .unwrap_or_else(|_| "Maple private AI assistant".to_string()), + ); + search_request.limit = Some(5); + + let search = client.web_search(search_request).await?; + assert!( + !search.results.is_empty(), + "Kagi search returned no results" + ); + assert!(search + .results + .iter() + .all(|result| result.url.starts_with("https://"))); + + let extract_url = std::env::var("VITE_TEST_WEB_EXTRACT_URL") + .unwrap_or_else(|_| "https://kagi.com/api/pricing".to_string()); + let extract = client + .web_extract(WebExtractRequest::new([extract_url.clone()])) + .await?; + + assert_eq!(extract.pages.len(), 1); + assert_eq!(extract.pages[0].url, extract_url); + assert!( + extract.pages[0].error.is_none(), + "Kagi extraction failed: {:?}", + extract.pages[0].error + ); + let markdown = extract.pages[0] + .markdown + .as_deref() + .expect("successful extraction should contain markdown"); + assert!(!markdown.trim().is_empty()); + assert!(!markdown.contains("![")); + assert!(!markdown.to_ascii_lowercase().contains(" Date: Fri, 17 Jul 2026 15:54:16 +0000 Subject: [PATCH 7/7] test: restore SDK state after web API mocks --- src/lib/test/integration/web.test.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/lib/test/integration/web.test.ts b/src/lib/test/integration/web.test.ts index 3b43071..5afa12d 100644 --- a/src/lib/test/integration/web.test.ts +++ b/src/lib/test/integration/web.test.ts @@ -2,6 +2,7 @@ import { afterEach, beforeEach, expect, mock, test } from "bun:test"; import { encode } from "@stablelib/base64"; import { decryptMessage, encryptMessage } from "../../encryption"; import { + getApiUrl, setApiUrl, webExtract, webSearch, @@ -16,6 +17,7 @@ const accessToken = "web-access-token"; const sessionId = "web-session-id"; const sessionKey = new Uint8Array(32).fill(19); const originalFetch = globalThis.fetch; +const originalApiUrl = getApiUrl(); beforeEach(() => { window.localStorage.clear(); @@ -28,6 +30,9 @@ beforeEach(() => { afterEach(() => { globalThis.fetch = originalFetch; + setApiUrl(originalApiUrl); + window.localStorage.clear(); + window.sessionStorage.clear(); }); test("webSearch sends an authenticated encrypted request and decrypts results", async () => {