mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
Add trusted access context for MCP metadata (#40992)
## What changed - Add `TrustedAccessContext` to fetch account-bound cyber verified-access grants and attach them as host-owned `openai/entitlementContext` metadata. - Replace caller-supplied entitlement metadata and return an `unknown` status when authentication, identity consistency, response validation, size limits, or the lookup deadline cannot be verified. - Add `HttpResponseBodyStream::from_chunks` for lightweight in-memory streaming clients. ## Testing - Cover grant mapping, unsupported authentication, account and workspace changes, malformed or oversized responses, duplicate programs, and lookup timeouts. GitOrigin-RevId: 343b6a6608c98c45a34939ef507c2183516a11c1
This commit is contained in:
soyeon-oai
@@ -25,6 +25,7 @@ pub use runtime::SandboxState;
|
||||
pub use runtime::apply_http_headers_helper;
|
||||
pub use tool_catalog_cache::McpToolCatalogCache;
|
||||
pub use tools::ToolInfo;
|
||||
pub use trusted_access::TrustedAccessContext;
|
||||
|
||||
/// Backward-compatible name for the shared Codex Apps tools runtime.
|
||||
pub type CodexAppsToolsCache = codex_connectors::ConnectorRuntimeManager<ToolInfo>;
|
||||
@@ -112,3 +113,4 @@ pub(crate) mod runtime;
|
||||
pub(crate) mod server;
|
||||
mod tool_catalog_cache;
|
||||
pub(crate) mod tools;
|
||||
mod trusted_access;
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use codex_exec_server::HttpClient;
|
||||
use codex_exec_server::HttpHeader;
|
||||
use codex_exec_server::HttpRedirectPolicy;
|
||||
use codex_exec_server::HttpRequestParams;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_protocol::auth::AuthMode;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use serde_json::Map;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
|
||||
const ENTITLEMENT_CONTEXT_KEY: &str = "openai/entitlementContext";
|
||||
const MAX_VERIFIED_ACCESS_RESPONSE_BYTES: usize = 1024 * 1024;
|
||||
const TRUSTED_ACCESS_TIMEOUT: Duration = Duration::from_millis(2_500);
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct VerifiedAccessResponse {
|
||||
programs: Vec<Value>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct VerifiedAccessProgram {
|
||||
state: VerifiedAccessState,
|
||||
grants: Vec<VerifiedAccessGrant>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum VerifiedAccessState {
|
||||
Active,
|
||||
Inactive,
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct VerifiedAccessGrant {
|
||||
level: TrustedAccessLevel,
|
||||
source: VerifiedAccessSource,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum TrustedAccessLevel {
|
||||
Tac1,
|
||||
Tac2,
|
||||
Tac3,
|
||||
Government,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
enum VerifiedAccessSource {
|
||||
Individual,
|
||||
Organization,
|
||||
}
|
||||
|
||||
/// Fetches account-bound verified access for trusted, host-owned MCP metadata.
|
||||
/// Callers must authorize the receiving plugin and tool before attaching it.
|
||||
pub struct TrustedAccessContext {
|
||||
auth: CodexAuth,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
chatgpt_base_url: String,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
}
|
||||
|
||||
impl TrustedAccessContext {
|
||||
pub fn new(
|
||||
auth: CodexAuth,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
chatgpt_base_url: String,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Self {
|
||||
Self {
|
||||
auth,
|
||||
auth_manager,
|
||||
chatgpt_base_url,
|
||||
http_client,
|
||||
}
|
||||
}
|
||||
|
||||
/// Replaces caller-supplied entitlement metadata with a fresh verified result.
|
||||
pub async fn add_context(&self, meta: Option<Value>) -> Option<Value> {
|
||||
let mut meta = match meta {
|
||||
Some(Value::Object(meta)) => meta,
|
||||
None => Map::new(),
|
||||
other => return other,
|
||||
};
|
||||
meta.remove(ENTITLEMENT_CONTEXT_KEY);
|
||||
|
||||
let status = tokio::time::timeout(TRUSTED_ACCESS_TIMEOUT, self.fetch_status())
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.unwrap_or_else(|| {
|
||||
json!({
|
||||
"schemaVersion": 1,
|
||||
"status": "unknown",
|
||||
"grants": [],
|
||||
"stale": false
|
||||
})
|
||||
});
|
||||
meta.insert(
|
||||
ENTITLEMENT_CONTEXT_KEY.to_string(),
|
||||
json!({
|
||||
"schemaVersion": 1,
|
||||
"entitlements": { "cyber_trusted_access": status }
|
||||
}),
|
||||
);
|
||||
Some(Value::Object(meta))
|
||||
}
|
||||
|
||||
async fn fetch_status(&self) -> Option<Value> {
|
||||
let auth = self.auth_manager.auth().await?;
|
||||
let expected_account_id = self
|
||||
.auth
|
||||
.get_account_id()
|
||||
.filter(|account_id| !account_id.trim().is_empty())?;
|
||||
let account_id = auth
|
||||
.get_account_id()
|
||||
.filter(|account_id| !account_id.trim().is_empty())?;
|
||||
if !matches!(
|
||||
self.auth.api_auth_mode(),
|
||||
AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens
|
||||
) || !matches!(
|
||||
auth.api_auth_mode(),
|
||||
AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens
|
||||
) || account_id != expected_account_id
|
||||
|| auth.get_chatgpt_user_id() != self.auth.get_chatgpt_user_id()
|
||||
|| auth.is_workspace_account() != self.auth.is_workspace_account()
|
||||
|| auth.is_fedramp_account() != self.auth.is_fedramp_account()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let headers = codex_model_provider::auth_provider_from_auth(&auth)
|
||||
.to_auth_headers()
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
Some(HttpHeader {
|
||||
name: name.as_str().to_string(),
|
||||
value: value.to_str().ok()?.to_string(),
|
||||
value_env_var: None,
|
||||
})
|
||||
})
|
||||
.collect::<Option<Vec<_>>>()?;
|
||||
let (response, mut response_body) = self
|
||||
.http_client
|
||||
.http_request_stream(HttpRequestParams {
|
||||
method: "GET".to_string(),
|
||||
url: format!(
|
||||
"{}/accounts/verified_access",
|
||||
self.chatgpt_base_url.trim_end_matches('/')
|
||||
),
|
||||
headers,
|
||||
body: None,
|
||||
timeout_ms: Some(TRUSTED_ACCESS_TIMEOUT.as_millis() as u64),
|
||||
redirect_policy: HttpRedirectPolicy::Stop,
|
||||
request_id: "trusted-access-status".to_string(),
|
||||
stream_response: true,
|
||||
})
|
||||
.await
|
||||
.ok()?;
|
||||
if response.status != 200 {
|
||||
return None;
|
||||
}
|
||||
let mut response_bytes = Vec::new();
|
||||
while let Some(chunk) = response_body.recv().await.ok()? {
|
||||
if chunk.len() > MAX_VERIFIED_ACCESS_RESPONSE_BYTES - response_bytes.len() {
|
||||
return None;
|
||||
}
|
||||
response_bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
let response: VerifiedAccessResponse = serde_json::from_slice(&response_bytes).ok()?;
|
||||
let mut programs = response
|
||||
.programs
|
||||
.into_iter()
|
||||
.filter(|program| program.get("program").and_then(Value::as_str) == Some("cyber"));
|
||||
let program = programs.next()?;
|
||||
if programs.next().is_some() {
|
||||
return None;
|
||||
}
|
||||
let program: VerifiedAccessProgram = serde_json::from_value(program).ok()?;
|
||||
|
||||
if self.auth_manager.auth_cached().is_none_or(|current| {
|
||||
!matches!(
|
||||
current.api_auth_mode(),
|
||||
AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens
|
||||
) || current.get_account_id().as_deref() != Some(account_id.as_str())
|
||||
|| current.get_chatgpt_user_id() != auth.get_chatgpt_user_id()
|
||||
|| current.is_workspace_account() != auth.is_workspace_account()
|
||||
|| current.is_fedramp_account() != auth.is_fedramp_account()
|
||||
}) {
|
||||
return None;
|
||||
}
|
||||
|
||||
let status = match program.state {
|
||||
VerifiedAccessState::Active if !program.grants.is_empty() => "granted",
|
||||
VerifiedAccessState::Inactive if program.grants.is_empty() => "not_granted",
|
||||
VerifiedAccessState::Unavailable if program.grants.is_empty() => "unknown",
|
||||
VerifiedAccessState::Active
|
||||
| VerifiedAccessState::Inactive
|
||||
| VerifiedAccessState::Unavailable => return None,
|
||||
};
|
||||
let grants = program
|
||||
.grants
|
||||
.into_iter()
|
||||
.map(|grant| {
|
||||
let source = match grant.source {
|
||||
VerifiedAccessSource::Individual => "user",
|
||||
VerifiedAccessSource::Organization => "current_account",
|
||||
};
|
||||
json!({ "level": grant.level, "source": source })
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
Some(json!({
|
||||
"schemaVersion": 1,
|
||||
"status": status,
|
||||
"grants": grants,
|
||||
"stale": false
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "trusted_access_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,658 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
use codex_exec_server::ByteChunk;
|
||||
use codex_exec_server::ExecServerError;
|
||||
use codex_exec_server::HttpClient;
|
||||
use codex_exec_server::HttpHeader;
|
||||
use codex_exec_server::HttpRedirectPolicy;
|
||||
use codex_exec_server::HttpRequestParams;
|
||||
use codex_exec_server::HttpRequestResponse;
|
||||
use codex_exec_server::HttpResponseBodyStream;
|
||||
use codex_login::AuthCredentialsStoreMode;
|
||||
use codex_login::AuthHeaders;
|
||||
use codex_login::AuthKeyringBackendKind;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::ExternalAuth;
|
||||
use codex_login::ExternalAuthFuture;
|
||||
use codex_login::ExternalAuthRefreshContext;
|
||||
use futures::FutureExt;
|
||||
use futures::future::BoxFuture;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use super::MAX_VERIFIED_ACCESS_RESPONSE_BYTES;
|
||||
use super::TrustedAccessContext;
|
||||
|
||||
struct RecordingHttpClient {
|
||||
requests: Mutex<Vec<HttpRequestParams>>,
|
||||
status: u16,
|
||||
response: Vec<u8>,
|
||||
response_chunks: Option<Vec<Vec<u8>>>,
|
||||
response_gate: Option<(Notify, Notify)>,
|
||||
}
|
||||
|
||||
impl RecordingHttpClient {
|
||||
fn new(status: u16, response: Value) -> Self {
|
||||
Self {
|
||||
requests: Mutex::new(Vec::new()),
|
||||
status,
|
||||
response: serde_json::to_vec(&response).expect("serialize response"),
|
||||
response_chunks: None,
|
||||
response_gate: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpClient for RecordingHttpClient {
|
||||
fn http_request(
|
||||
&self,
|
||||
params: HttpRequestParams,
|
||||
) -> BoxFuture<'_, Result<HttpRequestResponse, ExecServerError>> {
|
||||
self.requests.lock().expect("record request").push(params);
|
||||
let response = HttpRequestResponse {
|
||||
status: self.status,
|
||||
headers: Vec::new(),
|
||||
body: ByteChunk(self.response.clone()),
|
||||
};
|
||||
async move {
|
||||
if let Some((requested, release)) = &self.response_gate {
|
||||
requested.notify_one();
|
||||
release.notified().await;
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
|
||||
fn http_request_stream(
|
||||
&self,
|
||||
params: HttpRequestParams,
|
||||
) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> {
|
||||
self.requests.lock().expect("record request").push(params);
|
||||
let response = HttpRequestResponse {
|
||||
status: self.status,
|
||||
headers: Vec::new(),
|
||||
body: ByteChunk(Vec::new()),
|
||||
};
|
||||
let chunks = self
|
||||
.response_chunks
|
||||
.clone()
|
||||
.unwrap_or_else(|| vec![self.response.clone()]);
|
||||
async move {
|
||||
if let Some((requested, release)) = &self.response_gate {
|
||||
requested.notify_one();
|
||||
release.notified().await;
|
||||
}
|
||||
Ok((response, HttpResponseBodyStream::from_chunks(chunks)))
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
}
|
||||
|
||||
struct StaticExternalAuth(CodexAuth);
|
||||
|
||||
impl ExternalAuth for StaticExternalAuth {
|
||||
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
Box::pin(async { Ok(self.0.clone()) })
|
||||
}
|
||||
|
||||
fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
self.resolve()
|
||||
}
|
||||
}
|
||||
|
||||
struct SequencedExternalAuth(Mutex<VecDeque<CodexAuth>>);
|
||||
|
||||
impl ExternalAuth for SequencedExternalAuth {
|
||||
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
Box::pin(async {
|
||||
let mut sequence = self.0.lock().expect("auth sequence");
|
||||
Ok(if sequence.len() > 1 {
|
||||
sequence.pop_front().expect("next auth")
|
||||
} else {
|
||||
sequence.front().expect("last auth").clone()
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
self.resolve()
|
||||
}
|
||||
}
|
||||
|
||||
fn header_auth() -> CodexAuth {
|
||||
CodexAuth::Headers(AuthHeaders::new(
|
||||
[
|
||||
(
|
||||
"authorization".parse().unwrap(),
|
||||
"Bearer synthetic-pat".parse().unwrap(),
|
||||
),
|
||||
(
|
||||
"chatgpt-account-id".parse().unwrap(),
|
||||
"account-a".parse().unwrap(),
|
||||
),
|
||||
]
|
||||
.into_iter()
|
||||
.collect(),
|
||||
))
|
||||
}
|
||||
|
||||
fn chatgpt_auth(account_id: &str) -> CodexAuth {
|
||||
CodexAuth::from_external_chatgpt_tokens(
|
||||
"header.e30.same",
|
||||
account_id,
|
||||
/*chatgpt_plan_type*/ None,
|
||||
)
|
||||
.expect("test auth")
|
||||
}
|
||||
|
||||
fn fedramp_chatgpt_auth(account_id: &str) -> CodexAuth {
|
||||
CodexAuth::from_external_chatgpt_tokens(
|
||||
"header.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF9hY2NvdW50X2lzX2ZlZHJhbXAiOnRydWV9fQ.same",
|
||||
account_id,
|
||||
/*chatgpt_plan_type*/ None,
|
||||
)
|
||||
.expect("FedRAMP test auth")
|
||||
}
|
||||
|
||||
fn context(auth: CodexAuth, client: Arc<RecordingHttpClient>) -> TrustedAccessContext {
|
||||
TrustedAccessContext::new(
|
||||
auth.clone(),
|
||||
AuthManager::from_auth_for_testing(auth),
|
||||
"https://chatgpt.com/backend-api".to_string(),
|
||||
client,
|
||||
)
|
||||
}
|
||||
|
||||
fn cyber_response(state: &str, grants: Value) -> Value {
|
||||
json!({ "programs": [{ "program": "cyber", "state": state, "grants": grants }] })
|
||||
}
|
||||
|
||||
fn expected_metadata(status: &str, grants: Value) -> Value {
|
||||
json!({ "openai/entitlementContext": {
|
||||
"schemaVersion": 1,
|
||||
"entitlements": { "cyber_trusted_access": {
|
||||
"schemaVersion": 1,
|
||||
"status": status,
|
||||
"grants": grants,
|
||||
"stale": false
|
||||
} }
|
||||
} })
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delivers_account_bound_verified_access_to_the_calling_plugin() {
|
||||
let client = Arc::new(RecordingHttpClient::new(
|
||||
/*status*/ 200,
|
||||
cyber_response(
|
||||
"active",
|
||||
json!([
|
||||
{ "level": "tac2", "source": "individual" },
|
||||
{ "level": "government", "source": "organization" }
|
||||
]),
|
||||
),
|
||||
));
|
||||
let mut context = context(
|
||||
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
|
||||
client.clone(),
|
||||
);
|
||||
context.chatgpt_base_url.push('/');
|
||||
let metadata = context
|
||||
.add_context(Some(json!({ "threadId": "thread-1" })))
|
||||
.await;
|
||||
let mut expected = expected_metadata(
|
||||
"granted",
|
||||
json!([
|
||||
{ "level": "tac2", "source": "user" },
|
||||
{ "level": "government", "source": "current_account" }
|
||||
]),
|
||||
);
|
||||
expected["threadId"] = json!("thread-1");
|
||||
assert_eq!(metadata, Some(expected));
|
||||
|
||||
let mut requests = client.requests.lock().expect("inspect recorded requests");
|
||||
for request in requests.iter_mut() {
|
||||
request
|
||||
.headers
|
||||
.sort_by(|left, right| left.name.cmp(&right.name));
|
||||
}
|
||||
assert_eq!(
|
||||
*requests,
|
||||
vec![HttpRequestParams {
|
||||
method: "GET".to_string(),
|
||||
url: "https://chatgpt.com/backend-api/accounts/verified_access".to_string(),
|
||||
headers: vec![
|
||||
HttpHeader {
|
||||
name: "authorization".to_string(),
|
||||
value: "Bearer Access Token".to_string(),
|
||||
value_env_var: None,
|
||||
},
|
||||
HttpHeader {
|
||||
name: "chatgpt-account-id".to_string(),
|
||||
value: "account_id".to_string(),
|
||||
value_env_var: None,
|
||||
},
|
||||
],
|
||||
body: None,
|
||||
timeout_ms: Some(2_500),
|
||||
redirect_policy: HttpRedirectPolicy::Stop,
|
||||
request_id: "trusted-access-status".to_string(),
|
||||
stream_response: true,
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_auth_without_a_nonempty_account_id() -> anyhow::Result<()> {
|
||||
for account_id in [None, Some(""), Some(" ")] {
|
||||
let home = tempfile::tempdir()?;
|
||||
std::fs::write(
|
||||
home.path().join("auth.json"),
|
||||
serde_json::to_vec(&json!({
|
||||
"auth_mode": "chatgpt",
|
||||
"tokens": {
|
||||
"id_token": "header.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF91c2VyX2lkIjoidXNlci1hIn19.signature",
|
||||
"access_token": "synthetic-access-token",
|
||||
"refresh_token": "synthetic-refresh-token",
|
||||
"account_id": account_id
|
||||
},
|
||||
"last_refresh": "2099-01-01T00:00:00Z"
|
||||
}))?,
|
||||
)?;
|
||||
let auth = CodexAuth::from_auth_storage(
|
||||
home.path(),
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
AuthKeyringBackendKind::default(),
|
||||
&codex_login::test_support::transport_default_auth_route_config(),
|
||||
)
|
||||
.await?
|
||||
.expect("managed ChatGPT auth");
|
||||
let client = Arc::new(RecordingHttpClient::new(
|
||||
/*status*/ 200,
|
||||
cyber_response(
|
||||
"active",
|
||||
json!([{ "level": "tac1", "source": "individual" }]),
|
||||
),
|
||||
));
|
||||
assert_eq!(
|
||||
context(auth, client.clone())
|
||||
.add_context(/*meta*/ None)
|
||||
.await,
|
||||
Some(expected_metadata("unknown", json!([]))),
|
||||
"account_id={account_id:?}"
|
||||
);
|
||||
assert!(client.requests.lock().expect("inspect requests").is_empty());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_api_key_authentication_without_sending_credentials() {
|
||||
let client = Arc::new(RecordingHttpClient::new(
|
||||
/*status*/ 200,
|
||||
json!({ "programs": [] }),
|
||||
));
|
||||
let context = context(CodexAuth::from_api_key("test-api-key"), client.clone());
|
||||
|
||||
let metadata = context
|
||||
.add_context(Some(json!({
|
||||
"threadId": "thread-1",
|
||||
"openai/entitlementContext": { "forged": true }
|
||||
})))
|
||||
.await;
|
||||
|
||||
let mut expected = expected_metadata("unknown", json!([]));
|
||||
expected["threadId"] = json!("thread-1");
|
||||
assert_eq!(metadata, Some(expected));
|
||||
assert!(client.requests.lock().expect("inspect requests").is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_initial_unsupported_auth_after_switch_to_chatgpt() {
|
||||
let client = Arc::new(RecordingHttpClient::new(
|
||||
/*status*/ 200,
|
||||
cyber_response(
|
||||
"active",
|
||||
json!([{ "level": "tac1", "source": "individual" }]),
|
||||
),
|
||||
));
|
||||
let mut context = context(header_auth(), client.clone());
|
||||
context.auth_manager = AuthManager::from_auth_for_testing(chatgpt_auth("account-a"));
|
||||
|
||||
assert_eq!(
|
||||
context.add_context(/*meta*/ None).await,
|
||||
Some(expected_metadata("unknown", json!([])))
|
||||
);
|
||||
assert!(client.requests.lock().expect("inspect requests").is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn maps_verified_access_states_and_rejects_invalid_provider_responses() {
|
||||
let cases = [
|
||||
(
|
||||
200,
|
||||
json!({ "programs": [
|
||||
{ "program": "future", "state": "pending", "grants": [{ "level": "premium", "source": "subscription" }] },
|
||||
{ "program": "cyber", "state": "active", "grants": [{ "level": "tac1", "source": "individual" }] }
|
||||
] }),
|
||||
expected_metadata("granted", json!([{ "level": "tac1", "source": "user" }])),
|
||||
),
|
||||
(
|
||||
200,
|
||||
cyber_response("inactive", json!([])),
|
||||
expected_metadata("not_granted", json!([])),
|
||||
),
|
||||
(
|
||||
200,
|
||||
cyber_response("unavailable", json!([])),
|
||||
expected_metadata("unknown", json!([])),
|
||||
),
|
||||
(
|
||||
200,
|
||||
json!({ "programs": [{ "program": "other", "state": "active", "grants": [] }] }),
|
||||
expected_metadata("unknown", json!([])),
|
||||
),
|
||||
(
|
||||
200,
|
||||
cyber_response(
|
||||
"active",
|
||||
json!([{ "level": "admin", "source": "individual" }]),
|
||||
),
|
||||
expected_metadata("unknown", json!([])),
|
||||
),
|
||||
(
|
||||
200,
|
||||
cyber_response("active", json!([])),
|
||||
expected_metadata("unknown", json!([])),
|
||||
),
|
||||
(
|
||||
403,
|
||||
json!({ "error": "forbidden" }),
|
||||
expected_metadata("unknown", json!([])),
|
||||
),
|
||||
];
|
||||
|
||||
for (response_status, response, expected) in cases {
|
||||
let response_description = response.to_string();
|
||||
let context = context(
|
||||
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
|
||||
Arc::new(RecordingHttpClient::new(response_status, response)),
|
||||
);
|
||||
let metadata = context.add_context(/*meta*/ None).await.expect("metadata");
|
||||
assert_eq!(
|
||||
metadata, expected,
|
||||
"HTTP {response_status}: {response_description}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replaces_untrusted_entitlements_on_malformed_response() {
|
||||
let mut client = RecordingHttpClient::new(/*status*/ 200, Value::Null);
|
||||
client.response = b"{\"programs\":[".to_vec();
|
||||
let context = context(chatgpt_auth("account-a"), Arc::new(client));
|
||||
let mut expected = expected_metadata("unknown", json!([]));
|
||||
expected["threadId"] = json!("thread-1");
|
||||
|
||||
assert_eq!(
|
||||
context
|
||||
.add_context(Some(json!({
|
||||
"threadId": "thread-1",
|
||||
"openai/entitlementContext": { "forged": true }
|
||||
})))
|
||||
.await,
|
||||
Some(expected)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_verified_access_response_larger_than_one_mebibyte() {
|
||||
let valid_response = serde_json::to_vec(&cyber_response(
|
||||
"active",
|
||||
json!([{ "level": "tac1", "source": "individual" }]),
|
||||
))
|
||||
.expect("serialize response");
|
||||
let mut client = RecordingHttpClient::new(/*status*/ 200, Value::Null);
|
||||
client.response_chunks = Some(vec![
|
||||
vec![b' '; MAX_VERIFIED_ACCESS_RESPONSE_BYTES],
|
||||
valid_response,
|
||||
]);
|
||||
client.response = client
|
||||
.response_chunks
|
||||
.as_ref()
|
||||
.expect("response chunks")
|
||||
.concat();
|
||||
|
||||
assert_eq!(
|
||||
context(chatgpt_auth("account-a"), Arc::new(client))
|
||||
.add_context(/*meta*/ None)
|
||||
.await,
|
||||
Some(expected_metadata("unknown", json!([])))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn returns_unknown_at_the_lookup_deadline() {
|
||||
let mut client = RecordingHttpClient::new(/*status*/ 200, Value::Null);
|
||||
client.response_gate = Some((Notify::new(), Notify::new()));
|
||||
let client = Arc::new(client);
|
||||
let context = context(chatgpt_auth("account-a"), client.clone());
|
||||
let lookup = context.add_context(/*meta*/ None);
|
||||
tokio::pin!(lookup);
|
||||
|
||||
assert!(futures::poll!(lookup.as_mut()).is_pending());
|
||||
assert_eq!(client.requests.lock().expect("recorded requests").len(), 1);
|
||||
tokio::time::advance(Duration::from_millis(2_499)).await;
|
||||
assert!(futures::poll!(lookup.as_mut()).is_pending());
|
||||
tokio::time::advance(Duration::from_millis(1)).await;
|
||||
assert_eq!(
|
||||
lookup
|
||||
.now_or_never()
|
||||
.expect("lookup must finish at the deadline"),
|
||||
Some(expected_metadata("unknown", json!([])))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_duplicate_cyber_programs() {
|
||||
let active = json!({
|
||||
"program": "cyber", "state": "active",
|
||||
"grants": [{ "level": "tac1", "source": "individual" }]
|
||||
});
|
||||
for other in [
|
||||
active.clone(),
|
||||
json!({ "program": "cyber", "state": "inactive", "grants": [] }),
|
||||
json!({ "program": "cyber" }),
|
||||
] {
|
||||
for programs in [json!([active, other]), json!([other, active])] {
|
||||
let context = context(
|
||||
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
|
||||
Arc::new(RecordingHttpClient::new(
|
||||
/*status*/ 200,
|
||||
json!({ "programs": programs }),
|
||||
)),
|
||||
);
|
||||
assert_eq!(
|
||||
context.add_context(/*meta*/ None).await,
|
||||
Some(expected_metadata("unknown", json!([]))),
|
||||
"duplicate programs: {programs}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_identity_changes_before_sending_credentials() {
|
||||
for (description, initial_auth, selected_auth) in [
|
||||
(
|
||||
"account switch",
|
||||
chatgpt_auth("account-a"),
|
||||
chatgpt_auth("account-b"),
|
||||
),
|
||||
(
|
||||
"standard to FedRAMP",
|
||||
chatgpt_auth("account-a"),
|
||||
fedramp_chatgpt_auth("account-a"),
|
||||
),
|
||||
(
|
||||
"FedRAMP to standard",
|
||||
fedramp_chatgpt_auth("account-a"),
|
||||
chatgpt_auth("account-a"),
|
||||
),
|
||||
] {
|
||||
let client = Arc::new(RecordingHttpClient::new(
|
||||
/*status*/ 200,
|
||||
cyber_response(
|
||||
"active",
|
||||
json!([{ "level": "tac1", "source": "individual" }]),
|
||||
),
|
||||
));
|
||||
let mut context = context(initial_auth, client.clone());
|
||||
context.auth_manager = AuthManager::from_auth_for_testing(selected_auth);
|
||||
assert_eq!(
|
||||
context.add_context(/*meta*/ None).await,
|
||||
Some(expected_metadata("unknown", json!([]))),
|
||||
"identity change before the request: {description}"
|
||||
);
|
||||
assert!(
|
||||
client.requests.lock().expect("inspect requests").is_empty(),
|
||||
"identity change before the request: {description}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn uses_the_checked_auth_snapshot_for_request_headers() -> anyhow::Result<()> {
|
||||
let client = Arc::new(RecordingHttpClient::new(
|
||||
/*status*/ 200,
|
||||
cyber_response("inactive", json!([])),
|
||||
));
|
||||
let context = context(chatgpt_auth("account-a"), client.clone());
|
||||
let refreshed = CodexAuth::from_external_chatgpt_tokens(
|
||||
"header.e30.refreshed",
|
||||
"account-a",
|
||||
/*chatgpt_plan_type*/ None,
|
||||
)?;
|
||||
context
|
||||
.auth_manager
|
||||
.set_external_auth(Arc::new(SequencedExternalAuth(Mutex::new(VecDeque::from(
|
||||
[context.auth.clone(), refreshed, header_auth()],
|
||||
)))))
|
||||
.await?;
|
||||
|
||||
assert_eq!(
|
||||
context.add_context(/*meta*/ None).await,
|
||||
Some(expected_metadata("not_granted", json!([])))
|
||||
);
|
||||
let requests = client.requests.lock().expect("recorded requests");
|
||||
assert_eq!(requests.len(), 1);
|
||||
let authorization = requests[0]
|
||||
.headers
|
||||
.iter()
|
||||
.find(|header| header.name == "authorization");
|
||||
assert_eq!(
|
||||
authorization,
|
||||
Some(&HttpHeader {
|
||||
name: "authorization".to_string(),
|
||||
value: "Bearer header.e30.refreshed".to_string(),
|
||||
value_env_var: None,
|
||||
})
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_identity_changes_while_request_is_in_flight() -> anyhow::Result<()> {
|
||||
let personal = chatgpt_auth("account-a");
|
||||
let workspace =
|
||||
CodexAuth::from_external_chatgpt_tokens("header.e30.same", "account-a", Some("team"))?;
|
||||
for (description, initial_auth, selected_auth) in [
|
||||
(
|
||||
"account switch",
|
||||
personal.clone(),
|
||||
Some(chatgpt_auth("account-b")),
|
||||
),
|
||||
("logout", personal.clone(), None),
|
||||
(
|
||||
"personal to workspace",
|
||||
personal.clone(),
|
||||
Some(workspace.clone()),
|
||||
),
|
||||
("workspace to personal", workspace, Some(personal)),
|
||||
(
|
||||
"auth mode change",
|
||||
chatgpt_auth("account-a"),
|
||||
Some(header_auth()),
|
||||
),
|
||||
(
|
||||
"standard to FedRAMP",
|
||||
chatgpt_auth("account-a"),
|
||||
Some(fedramp_chatgpt_auth("account-a")),
|
||||
),
|
||||
(
|
||||
"FedRAMP to standard",
|
||||
fedramp_chatgpt_auth("account-a"),
|
||||
Some(chatgpt_auth("account-a")),
|
||||
),
|
||||
] {
|
||||
let mut client = RecordingHttpClient::new(
|
||||
/*status*/ 200,
|
||||
cyber_response(
|
||||
"active",
|
||||
json!([{ "level": "tac1", "source": "individual" }]),
|
||||
),
|
||||
);
|
||||
client.response_gate = Some((Notify::new(), Notify::new()));
|
||||
let client = Arc::new(client);
|
||||
let request_is_fedramp = initial_auth.is_fedramp_account();
|
||||
let context = context(initial_auth, client.clone());
|
||||
let auth_manager = &context.auth_manager;
|
||||
auth_manager
|
||||
.set_external_auth(Arc::new(StaticExternalAuth(context.auth.clone())))
|
||||
.await?;
|
||||
|
||||
let (metadata, auth_change) = tokio::join!(context.add_context(/*meta*/ None), async {
|
||||
let (requested, release) = client.response_gate.as_ref().expect("gated response");
|
||||
tokio::time::timeout(Duration::from_secs(5), requested.notified()).await?;
|
||||
{
|
||||
let requests = client.requests.lock().expect("inspect request");
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert!(requests[0].headers.iter().any(|header| {
|
||||
header.name.eq_ignore_ascii_case("chatgpt-account-id")
|
||||
&& header.value == "account-a"
|
||||
}));
|
||||
assert_eq!(
|
||||
requests[0]
|
||||
.headers
|
||||
.iter()
|
||||
.any(|header| header.name.eq_ignore_ascii_case("x-openai-fedramp")),
|
||||
request_is_fedramp
|
||||
);
|
||||
}
|
||||
if let Some(auth) = selected_auth {
|
||||
auth_manager
|
||||
.set_external_auth(Arc::new(StaticExternalAuth(auth)))
|
||||
.await?;
|
||||
} else {
|
||||
auth_manager.clear_external_auth();
|
||||
}
|
||||
release.notify_one();
|
||||
Ok::<(), anyhow::Error>(())
|
||||
});
|
||||
auth_change?;
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
Some(expected_metadata("unknown", json!([]))),
|
||||
"auth change after the request: {description}"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -78,6 +78,24 @@ pub struct HttpResponseBodyStream {
|
||||
}
|
||||
|
||||
impl HttpResponseBodyStream {
|
||||
/// Creates an in-memory response stream from pre-buffered chunks.
|
||||
///
|
||||
/// This is useful for [`crate::HttpClient`] implementations that already
|
||||
/// own the response bytes, including lightweight test clients.
|
||||
#[doc(hidden)]
|
||||
pub fn from_chunks(chunks: Vec<Vec<u8>>) -> Self {
|
||||
let body = futures::stream::iter(
|
||||
chunks
|
||||
.into_iter()
|
||||
.map(|chunk| Ok::<Bytes, HttpError>(chunk.into())),
|
||||
);
|
||||
Self {
|
||||
inner: HttpResponseBodyStreamInner::Local {
|
||||
body: Box::pin(body),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn local(response: HttpResponse) -> Self {
|
||||
Self {
|
||||
inner: HttpResponseBodyStreamInner::Local {
|
||||
|
||||
Reference in New Issue
Block a user