From d0b19e41e52035af64452c0b97d7cdcfd54f3623 Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Tue, 2 Jun 2026 19:55:08 +0000 Subject: [PATCH] [codex-client] add native integrity state store [ci changed_files] --- codex-rs/Cargo.lock | 1 + codex-rs/codex-client/Cargo.toml | 1 + codex-rs/codex-client/src/chatgpt_hosts.rs | 33 +++ codex-rs/codex-client/src/lib.rs | 8 + codex-rs/codex-client/src/native_integrity.rs | 240 ++++++++++++++++++ .../src/native_integrity_tests.rs | 120 +++++++++ 6 files changed, 403 insertions(+) create mode 100644 codex-rs/codex-client/src/native_integrity.rs create mode 100644 codex-rs/codex-client/src/native_integrity_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 27b66293e1..24890c9bc5 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2319,6 +2319,7 @@ dependencies = [ "async-trait", "bytes", "codex-utils-cargo-bin", + "codex-utils-path", "codex-utils-rustls-provider", "eventsource-stream", "futures", diff --git a/codex-rs/codex-client/Cargo.toml b/codex-rs/codex-client/Cargo.toml index 184505eb55..6c00908a33 100644 --- a/codex-rs/codex-client/Cargo.toml +++ b/codex-rs/codex-client/Cargo.toml @@ -7,6 +7,7 @@ version.workspace = true [dependencies] async-trait = { workspace = true } bytes = { workspace = true } +codex-utils-path = { workspace = true } eventsource-stream = { workspace = true } futures = { workspace = true } http = { workspace = true } diff --git a/codex-rs/codex-client/src/chatgpt_hosts.rs b/codex-rs/codex-client/src/chatgpt_hosts.rs index dd0b99589c..7ab6819e79 100644 --- a/codex-rs/codex-client/src/chatgpt_hosts.rs +++ b/codex-rs/codex-client/src/chatgpt_hosts.rs @@ -10,6 +10,18 @@ pub fn is_allowed_chatgpt_host(host: &str) -> bool { .any(|suffix| host.ends_with(suffix)) } +/// Returns whether `url` is an HTTPS or secure WebSocket ChatGPT URL that Codex +/// may treat as first-party traffic. +pub fn is_allowed_chatgpt_request_url(url: &str) -> bool { + let Ok(url) = reqwest::Url::parse(url) else { + return false; + }; + if !matches!(url.scheme(), "https" | "wss") { + return false; + } + url.host_str().is_some_and(is_allowed_chatgpt_host) +} + #[cfg(test)] mod tests { use super::*; @@ -36,4 +48,25 @@ mod tests { assert!(!is_allowed_chatgpt_host(host)); } } + + #[test] + fn recognizes_secure_chatgpt_request_urls() { + for url in [ + "https://chatgpt.com/backend-api/codex/responses", + "https://preview.chatgpt.com/backend-api/codex/models", + "wss://chatgpt-staging.com/backend-api/codex/responses", + ] { + assert!(is_allowed_chatgpt_request_url(url)); + } + + for url in [ + "http://chatgpt.com/backend-api/codex/responses", + "ws://chatgpt.com/backend-api/codex/responses", + "https://api.openai.com/v1/responses", + "https://chatgpt.com.evil.example/backend-api", + "not a url", + ] { + assert!(!is_allowed_chatgpt_request_url(url)); + } + } } diff --git a/codex-rs/codex-client/src/lib.rs b/codex-rs/codex-client/src/lib.rs index 0f503fb3e2..785c9bc0e2 100644 --- a/codex-rs/codex-client/src/lib.rs +++ b/codex-rs/codex-client/src/lib.rs @@ -3,6 +3,7 @@ mod chatgpt_hosts; mod custom_ca; mod default_client; mod error; +mod native_integrity; mod request; mod retry; mod sse; @@ -11,6 +12,7 @@ mod transport; pub use crate::chatgpt_cloudflare_cookies::with_chatgpt_cloudflare_cookie_store; pub use crate::chatgpt_hosts::is_allowed_chatgpt_host; +pub use crate::chatgpt_hosts::is_allowed_chatgpt_request_url; pub use crate::custom_ca::BuildCustomCaTransportError; /// Test-only subprocess hook for custom CA coverage. /// @@ -25,6 +27,12 @@ pub use crate::default_client::CodexHttpClient; pub use crate::default_client::CodexRequestBuilder; pub use crate::error::StreamError; pub use crate::error::TransportError; +pub use crate::native_integrity::INTEGRITY_STATE_HEADER_NAME; +pub use crate::native_integrity::INTEGRITY_STATE_UPDATE_HEADER_NAME; +pub use crate::native_integrity::NativeIntegrityStateContext; +pub use crate::native_integrity::NativeIntegrityStateFile; +pub use crate::native_integrity::NativeIntegrityStateStore; +pub use crate::native_integrity::NativeIntegritySurface; pub use crate::request::PreparedRequestBody; pub use crate::request::Request; pub use crate::request::RequestBody; diff --git a/codex-rs/codex-client/src/native_integrity.rs b/codex-rs/codex-client/src/native_integrity.rs new file mode 100644 index 0000000000..63af8466a2 --- /dev/null +++ b/codex-rs/codex-client/src/native_integrity.rs @@ -0,0 +1,240 @@ +use std::fmt; +use std::fs; +use std::io; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::PoisonError; +use std::sync::RwLock; + +use codex_utils_path::write_atomically; +use serde::Deserialize; +use serde::Serialize; + +pub const INTEGRITY_STATE_HEADER_NAME: &str = "X-OAI-IS"; +pub const INTEGRITY_STATE_UPDATE_HEADER_NAME: &str = "X-OAI-IS-Update"; +const MAX_INTEGRITY_STATE_ENVELOPE_BYTES: usize = 2048; + +const STATE_DIR_NAME: &str = "state"; + +/// Native Codex surface that owns one local integrity-state file. +/// +/// Unknown app-server clients intentionally share the CLI state file. This +/// preserves the default classification while first-party clients opt into a +/// more specific surface by setting `clientInfo.name` during initialization. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum NativeIntegritySurface { + CodexCli, + CodexTui, + CodexExec, + CodexVscode, + CodexDesktop, + CodexDesktopSsh, + CodexRemoteControl, +} + +impl NativeIntegritySurface { + pub fn from_app_server_client_name(client_name: &str) -> Self { + match client_name { + "codex-tui" => Self::CodexTui, + "codex_exec" => Self::CodexExec, + "codex_vscode" => Self::CodexVscode, + "codex_desktop" => Self::CodexDesktop, + "codex_desktop_ssh" => Self::CodexDesktopSsh, + "codex_remote_control" => Self::CodexRemoteControl, + _ => Self::CodexCli, + } + } + + pub const fn as_str(self) -> &'static str { + match self { + Self::CodexCli => "codex_cli", + Self::CodexTui => "codex_tui", + Self::CodexExec => "codex_exec", + Self::CodexVscode => "codex_vscode", + Self::CodexDesktop => "codex_desktop", + Self::CodexDesktopSsh => "codex_desktop_ssh", + Self::CodexRemoteControl => "codex_remote_control", + } + } +} + +impl fmt::Display for NativeIntegritySurface { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct NativeIntegrityStateFile { + pub state: String, +} + +#[derive(Clone, Debug)] +pub struct NativeIntegrityStateStore { + codex_home: PathBuf, +} + +impl NativeIntegrityStateStore { + pub fn new(codex_home: PathBuf) -> Self { + Self { codex_home } + } + + fn state_path(&self, surface: NativeIntegritySurface) -> PathBuf { + self.codex_home + .join(STATE_DIR_NAME) + .join(format!("{surface}.json")) + } + + pub fn load( + &self, + surface: NativeIntegritySurface, + ) -> io::Result> { + let path = self.state_path(surface); + let contents = match fs::read(&path) { + Ok(contents) => contents, + Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(error), + }; + let state_file: NativeIntegrityStateFile = + serde_json::from_slice(&contents).map_err(io::Error::other)?; + validate_integrity_state_envelope(&state_file.state)?; + Ok(Some(state_file)) + } + + pub fn replace(&self, surface: NativeIntegritySurface, state: String) -> io::Result<()> { + validate_integrity_state_envelope(&state)?; + self.write(surface, &NativeIntegrityStateFile { state }) + } + + pub fn clear(&self, surface: NativeIntegritySurface) -> io::Result<()> { + match fs::remove_file(self.state_path(surface)) { + Ok(()) => Ok(()), + Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(error), + } + } + + /// Stores `next_state` only if the local file still contains `expected_state`. + /// + /// This is intentionally lock-free. Concurrent writers may both observe + /// the same prior value, in which case atomic replacement keeps the file + /// well-formed and the last writer wins. + pub fn compare_and_store( + &self, + surface: NativeIntegritySurface, + expected_state: &str, + next_state: String, + ) -> io::Result { + validate_integrity_state_envelope(&next_state)?; + let Some(current) = self.load(surface)? else { + return Ok(false); + }; + if current.state != expected_state { + return Ok(false); + } + + self.write(surface, &NativeIntegrityStateFile { state: next_state })?; + Ok(true) + } + + fn write( + &self, + surface: NativeIntegritySurface, + state_file: &NativeIntegrityStateFile, + ) -> io::Result<()> { + let path = self.state_path(surface); + let contents = serde_json::to_string_pretty(state_file).map_err(io::Error::other)?; + write_atomically(&path, &contents) + } +} + +/// Shared surface selection for one native Codex client session. +/// +/// App-server clients may update the selected surface after the model client is +/// constructed. Requests snapshot the current value when they create their +/// auth decorator, so clones of the model client stay in sync without moving +/// in-flight rotations to another surface. +#[derive(Clone, Debug)] +pub struct NativeIntegrityStateContext { + store: NativeIntegrityStateStore, + surface: Arc>, +} + +impl NativeIntegrityStateContext { + pub fn new(codex_home: PathBuf, surface: NativeIntegritySurface) -> Self { + Self { + store: NativeIntegrityStateStore::new(codex_home), + surface: Arc::new(RwLock::new(surface)), + } + } + + pub fn surface(&self) -> NativeIntegritySurface { + *self.surface.read().unwrap_or_else(PoisonError::into_inner) + } + + pub fn set_surface(&self, surface: NativeIntegritySurface) -> bool { + let mut selected_surface = self.surface.write().unwrap_or_else(PoisonError::into_inner); + if *selected_surface == surface { + return false; + } + *selected_surface = surface; + true + } + + pub fn load_for_surface( + &self, + surface: NativeIntegritySurface, + ) -> io::Result> { + self.store.load(surface) + } + + pub fn compare_and_store_for_surface( + &self, + surface: NativeIntegritySurface, + expected_state: &str, + next_state: String, + ) -> io::Result { + self.store + .compare_and_store(surface, expected_state, next_state) + } +} + +fn is_valid_integrity_state_envelope(value: &str) -> bool { + if value.is_empty() || value.len() > MAX_INTEGRITY_STATE_ENVELOPE_BYTES || value.trim() != value + { + return false; + } + + let mut parts = value.split('.'); + if parts.next() != Some("ois1") { + return false; + } + + let valid_part = |part: &str| { + !part.is_empty() + && part + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_') + }; + + matches!( + (parts.next(), parts.next(), parts.next(), parts.next()), + (Some(header), Some(nonce), Some(ciphertext), None) + if valid_part(header) && valid_part(nonce) && valid_part(ciphertext) + ) +} + +fn validate_integrity_state_envelope(value: &str) -> io::Result<()> { + if is_valid_integrity_state_envelope(value) { + Ok(()) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidData, + "integrity state envelope is malformed", + )) + } +} + +#[cfg(test)] +#[path = "native_integrity_tests.rs"] +mod tests; diff --git a/codex-rs/codex-client/src/native_integrity_tests.rs b/codex-rs/codex-client/src/native_integrity_tests.rs new file mode 100644 index 0000000000..0f42a363e0 --- /dev/null +++ b/codex-rs/codex-client/src/native_integrity_tests.rs @@ -0,0 +1,120 @@ +use super::*; +use pretty_assertions::assert_eq; +use tempfile::TempDir; + +const STATE_N: &str = "ois1.a.b.c"; +const STATE_N_PLUS_ONE: &str = "ois1.d.e.f"; + +#[test] +fn maps_app_server_client_names_to_bounded_surfaces() { + assert_eq!( + [ + "codex-tui", + "codex_exec", + "codex_vscode", + "codex_desktop", + "codex_desktop_ssh", + "codex_remote_control", + "third_party_client", + ] + .map(NativeIntegritySurface::from_app_server_client_name), + [ + NativeIntegritySurface::CodexTui, + NativeIntegritySurface::CodexExec, + NativeIntegritySurface::CodexVscode, + NativeIntegritySurface::CodexDesktop, + NativeIntegritySurface::CodexDesktopSsh, + NativeIntegritySurface::CodexRemoteControl, + NativeIntegritySurface::CodexCli, + ] + ); +} + +#[test] +fn stores_state_in_one_file_per_surface() { + let codex_home = TempDir::new().expect("tempdir"); + let store = NativeIntegrityStateStore::new(codex_home.path().to_path_buf()); + + store + .replace(NativeIntegritySurface::CodexCli, STATE_N.to_string()) + .expect("CLI state should store"); + store + .replace( + NativeIntegritySurface::CodexDesktop, + STATE_N_PLUS_ONE.to_string(), + ) + .expect("desktop state should store"); + + assert_eq!( + store + .load(NativeIntegritySurface::CodexCli) + .expect("CLI state should load") + .expect("CLI state should exist"), + NativeIntegrityStateFile { + state: STATE_N.to_string(), + } + ); + assert_eq!( + store + .load(NativeIntegritySurface::CodexDesktop) + .expect("desktop state should load") + .expect("desktop state should exist"), + NativeIntegrityStateFile { + state: STATE_N_PLUS_ONE.to_string(), + } + ); + assert_eq!( + store.state_path(NativeIntegritySurface::CodexCli), + codex_home.path().join("state/codex_cli.json") + ); + store + .clear(NativeIntegritySurface::CodexDesktop) + .expect("desktop state should clear"); + assert_eq!( + store + .load(NativeIntegritySurface::CodexDesktop) + .expect("desktop state should load"), + None + ); +} + +#[test] +fn compare_and_store_rejects_a_stale_prior_value() { + let codex_home = TempDir::new().expect("tempdir"); + let store = NativeIntegrityStateStore::new(codex_home.path().to_path_buf()); + store + .replace(NativeIntegritySurface::CodexCli, STATE_N.to_string()) + .expect("state should store"); + + assert!( + !store + .compare_and_store( + NativeIntegritySurface::CodexCli, + "ois1.stale.state.value", + STATE_N_PLUS_ONE.to_string(), + ) + .expect("compare should succeed") + ); + assert!( + store + .compare_and_store( + NativeIntegritySurface::CodexCli, + STATE_N, + STATE_N_PLUS_ONE.to_string(), + ) + .expect("compare should succeed") + ); +} + +#[test] +fn validates_integrity_state_envelopes() { + assert!(is_valid_integrity_state_envelope(STATE_N)); + assert!(!is_valid_integrity_state_envelope("")); + assert!(!is_valid_integrity_state_envelope("state")); + assert!(!is_valid_integrity_state_envelope(&format!(" {STATE_N}"))); + assert!(!is_valid_integrity_state_envelope("ois1.a.b")); + assert!(!is_valid_integrity_state_envelope(&format!( + "ois1.a.b.{}", + "c".repeat(MAX_INTEGRITY_STATE_ENVELOPE_BYTES + 1) + ))); +}