From d2b2ee88fc9a0509be79b21de5695f4829cf43f4 Mon Sep 17 00:00:00 2001 From: mikhail-oai Date: Tue, 10 Mar 2026 12:24:46 -0400 Subject: [PATCH] generalize json serialization and unify MCP/Auth codepaths. --- codex-rs/Cargo.lock | 3 + codex-rs/keyring-store/Cargo.toml | 5 + codex-rs/keyring-store/src/lib.rs | 7 + codex-rs/keyring-store/src/split_json.rs | 718 +++++++++++++++++++++++ codex-rs/login/src/auth/storage.rs | 300 ++-------- codex-rs/login/src/auth/storage_tests.rs | 159 ++--- codex-rs/rmcp-client/src/oauth.rs | 202 ++++++- 7 files changed, 1003 insertions(+), 391 deletions(-) create mode 100644 codex-rs/keyring-store/src/split_json.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c8eae3be0c..2cc28531fb 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2158,6 +2158,9 @@ name = "codex-keyring-store" version = "0.0.0" dependencies = [ "keyring", + "pretty_assertions", + "serde", + "serde_json", "tracing", ] diff --git a/codex-rs/keyring-store/Cargo.toml b/codex-rs/keyring-store/Cargo.toml index dd3686ee59..6a546b55fc 100644 --- a/codex-rs/keyring-store/Cargo.toml +++ b/codex-rs/keyring-store/Cargo.toml @@ -9,8 +9,13 @@ workspace = true [dependencies] keyring = { workspace = true, features = ["crypto-rust"] } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } tracing = { workspace = true } +[dev-dependencies] +pretty_assertions = { workspace = true } + [target.'cfg(target_os = "linux")'.dependencies] keyring = { workspace = true, features = ["linux-native-async-persistent"] } diff --git a/codex-rs/keyring-store/src/lib.rs b/codex-rs/keyring-store/src/lib.rs index a6b28babdd..3cd55a50ca 100644 --- a/codex-rs/keyring-store/src/lib.rs +++ b/codex-rs/keyring-store/src/lib.rs @@ -5,6 +5,13 @@ use std::fmt; use std::fmt::Debug; use tracing::trace; +mod split_json; + +pub use split_json::SplitJsonKeyringError; +pub use split_json::delete_split_json_from_keyring; +pub use split_json::load_split_json_from_keyring; +pub use split_json::save_split_json_to_keyring; + #[derive(Debug)] pub enum CredentialStoreError { Other(KeyringError), diff --git a/codex-rs/keyring-store/src/split_json.rs b/codex-rs/keyring-store/src/split_json.rs new file mode 100644 index 0000000000..282aa57297 --- /dev/null +++ b/codex-rs/keyring-store/src/split_json.rs @@ -0,0 +1,718 @@ +use crate::CredentialStoreError; +use crate::KeyringStore; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Map; +use serde_json::Value; +use std::fmt; +use std::fmt::Write as _; +use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; +use tracing::warn; + +const LAYOUT_VERSION: &str = "v1"; +const ACTIVE_REVISION_ENTRY: &str = "active"; +const MANIFEST_ENTRY: &str = "manifest"; +const VALUE_ENTRY_PREFIX: &str = "value"; +const ROOT_PATH_SENTINEL: &str = "root"; + +#[derive(Debug, Clone)] +pub struct SplitJsonKeyringError { + message: String, +} + +impl SplitJsonKeyringError { + fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } +} + +impl fmt::Display for SplitJsonKeyringError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.message) + } +} + +impl std::error::Error for SplitJsonKeyringError {} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +enum JsonNodeKind { + Null, + Bool, + Number, + String, + Object, + Array, +} + +impl JsonNodeKind { + fn from_value(value: &Value) -> Self { + match value { + Value::Null => Self::Null, + Value::Bool(_) => Self::Bool, + Value::Number(_) => Self::Number, + Value::String(_) => Self::String, + Value::Object(_) => Self::Object, + Value::Array(_) => Self::Array, + } + } + + fn is_container(self) -> bool { + matches!(self, Self::Object | Self::Array) + } + + fn empty_value(self) -> Option { + match self { + Self::Object => Some(Value::Object(Map::new())), + Self::Array => Some(Value::Array(Vec::new())), + Self::Null | Self::Bool | Self::Number | Self::String => None, + } + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +struct SplitJsonNode { + path: String, + kind: JsonNodeKind, +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +struct SplitJsonManifest { + nodes: Vec, +} + +type SplitJsonLeafValues = Vec<(String, Vec)>; + +pub fn load_split_json_from_keyring( + keyring_store: &K, + service: &str, + base_key: &str, +) -> Result, SplitJsonKeyringError> { + let Some(revision) = load_active_revision(keyring_store, service, base_key)? else { + return Ok(None); + }; + let manifest = load_manifest(keyring_store, service, base_key, &revision)?; + inflate_split_json(keyring_store, service, base_key, &revision, &manifest).map(Some) +} + +pub fn save_split_json_to_keyring( + keyring_store: &K, + service: &str, + base_key: &str, + value: &Value, +) -> Result<(), SplitJsonKeyringError> { + let previous_revision = match load_active_revision(keyring_store, service, base_key) { + Ok(revision) => revision, + Err(err) => { + warn!("failed to read previous split JSON revision from keyring: {err}"); + None + } + }; + let revision = next_revision(); + let (manifest, leaf_values) = flatten_split_json(value)?; + + for (path, bytes) in leaf_values { + let key = value_key(base_key, &revision, &path); + save_secret_to_keyring( + keyring_store, + service, + &key, + &bytes, + &format!("JSON value at {path}"), + )?; + } + + let manifest_key = revision_key(base_key, &revision, MANIFEST_ENTRY); + let manifest_bytes = serde_json::to_vec(&manifest).map_err(|err| { + SplitJsonKeyringError::new(format!("failed to serialize JSON manifest: {err}")) + })?; + save_secret_to_keyring( + keyring_store, + service, + &manifest_key, + &manifest_bytes, + "JSON manifest", + )?; + + let active_key = layout_key(base_key, ACTIVE_REVISION_ENTRY); + save_secret_to_keyring( + keyring_store, + service, + &active_key, + revision.as_bytes(), + "active split JSON revision", + )?; + + if let Some(previous_revision) = previous_revision + && previous_revision != revision + && let Err(err) = + delete_revision_entries(keyring_store, service, base_key, &previous_revision) + { + warn!("failed to remove stale split JSON revision from keyring: {err}"); + } + + Ok(()) +} + +pub fn delete_split_json_from_keyring( + keyring_store: &K, + service: &str, + base_key: &str, +) -> Result { + let Some(revision) = load_active_revision(keyring_store, service, base_key)? else { + return Ok(false); + }; + + let mut removed = delete_revision_entries(keyring_store, service, base_key, &revision)?; + let active_key = layout_key(base_key, ACTIVE_REVISION_ENTRY); + removed |= delete_keyring_entry( + keyring_store, + service, + &active_key, + "active split JSON revision", + )?; + Ok(removed) +} + +fn flatten_split_json( + value: &Value, +) -> Result<(SplitJsonManifest, SplitJsonLeafValues), SplitJsonKeyringError> { + let mut nodes = Vec::new(); + let mut leaf_values = Vec::new(); + collect_nodes("", value, &mut nodes, &mut leaf_values)?; + nodes.sort_by(|left, right| { + path_depth(&left.path) + .cmp(&path_depth(&right.path)) + .then_with(|| left.path.cmp(&right.path)) + }); + leaf_values.sort_by(|left, right| left.0.cmp(&right.0)); + Ok((SplitJsonManifest { nodes }, leaf_values)) +} + +fn collect_nodes( + path: &str, + value: &Value, + nodes: &mut Vec, + leaf_values: &mut SplitJsonLeafValues, +) -> Result<(), SplitJsonKeyringError> { + let kind = JsonNodeKind::from_value(value); + nodes.push(SplitJsonNode { + path: path.to_string(), + kind, + }); + + match value { + Value::Object(map) => { + let mut keys = map.keys().cloned().collect::>(); + keys.sort(); + for key in keys { + let child_path = append_json_pointer_token(path, &key); + let child_value = map.get(&key).ok_or_else(|| { + SplitJsonKeyringError::new(format!( + "missing object value for path {child_path}" + )) + })?; + collect_nodes(&child_path, child_value, nodes, leaf_values)?; + } + } + Value::Array(items) => { + for (index, item) in items.iter().enumerate() { + let child_path = append_json_pointer_token(path, &index.to_string()); + collect_nodes(&child_path, item, nodes, leaf_values)?; + } + } + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => { + let bytes = serde_json::to_vec(value).map_err(|err| { + SplitJsonKeyringError::new(format!( + "failed to serialize JSON value at {path}: {err}" + )) + })?; + leaf_values.push((path.to_string(), bytes)); + } + } + + Ok(()) +} + +fn inflate_split_json( + keyring_store: &K, + service: &str, + base_key: &str, + revision: &str, + manifest: &SplitJsonManifest, +) -> Result { + let root_node = manifest + .nodes + .iter() + .find(|node| node.path.is_empty()) + .ok_or_else(|| SplitJsonKeyringError::new("missing root JSON node in keyring manifest"))?; + + let mut result = if let Some(value) = root_node.kind.empty_value() { + value + } else { + load_value(keyring_store, service, base_key, revision, "")? + }; + + let mut nodes = manifest.nodes.clone(); + nodes.sort_by(|left, right| { + path_depth(&left.path) + .cmp(&path_depth(&right.path)) + .then_with(|| left.path.cmp(&right.path)) + }); + + for node in nodes.into_iter().filter(|node| !node.path.is_empty()) { + let value = if let Some(value) = node.kind.empty_value() { + value + } else { + load_value(keyring_store, service, base_key, revision, &node.path)? + }; + insert_value_at_pointer(&mut result, &node.path, value)?; + } + + Ok(result) +} + +fn load_value( + keyring_store: &K, + service: &str, + base_key: &str, + revision: &str, + path: &str, +) -> Result { + let key = value_key(base_key, revision, path); + let bytes = load_secret_from_keyring( + keyring_store, + service, + &key, + &format!("JSON value at {path}"), + )? + .ok_or_else(|| { + SplitJsonKeyringError::new(format!("missing JSON value at {path} in keyring")) + })?; + serde_json::from_slice(&bytes).map_err(|err| { + SplitJsonKeyringError::new(format!("failed to deserialize JSON value at {path}: {err}")) + }) +} + +fn insert_value_at_pointer( + root: &mut Value, + pointer: &str, + value: Value, +) -> Result<(), SplitJsonKeyringError> { + if pointer.is_empty() { + *root = value; + return Ok(()); + } + + let tokens = decode_json_pointer(pointer)?; + let Some((last, parents)) = tokens.split_last() else { + return Err(SplitJsonKeyringError::new( + "missing JSON pointer path tokens", + )); + }; + + let mut current = root; + for token in parents { + current = match current { + Value::Object(map) => map.get_mut(token).ok_or_else(|| { + SplitJsonKeyringError::new(format!( + "missing parent object entry for JSON pointer {pointer}" + )) + })?, + Value::Array(items) => { + let index = parse_array_index(token, pointer)?; + items.get_mut(index).ok_or_else(|| { + SplitJsonKeyringError::new(format!( + "missing parent array entry for JSON pointer {pointer}" + )) + })? + } + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => { + return Err(SplitJsonKeyringError::new(format!( + "encountered scalar while walking JSON pointer {pointer}" + ))); + } + }; + } + + match current { + Value::Object(map) => { + map.insert(last.to_string(), value); + Ok(()) + } + Value::Array(items) => { + let index = parse_array_index(last, pointer)?; + if index >= items.len() { + items.resize(index + 1, Value::Null); + } + items[index] = value; + Ok(()) + } + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => { + Err(SplitJsonKeyringError::new(format!( + "encountered scalar while assigning JSON pointer {pointer}" + ))) + } + } +} + +fn delete_revision_entries( + keyring_store: &K, + service: &str, + base_key: &str, + revision: &str, +) -> Result { + let manifest = load_manifest(keyring_store, service, base_key, revision)?; + let mut removed = false; + + for node in manifest.nodes { + if node.kind.is_container() { + continue; + } + let key = value_key(base_key, revision, &node.path); + removed |= delete_keyring_entry( + keyring_store, + service, + &key, + &format!("JSON value at {}", node.path), + )?; + } + + let manifest_key = revision_key(base_key, revision, MANIFEST_ENTRY); + removed |= delete_keyring_entry(keyring_store, service, &manifest_key, "JSON manifest")?; + Ok(removed) +} + +fn load_manifest( + keyring_store: &K, + service: &str, + base_key: &str, + revision: &str, +) -> Result { + let manifest_key = revision_key(base_key, revision, MANIFEST_ENTRY); + let bytes = load_secret_from_keyring(keyring_store, service, &manifest_key, "JSON manifest")? + .ok_or_else(|| SplitJsonKeyringError::new("missing JSON manifest in keyring"))?; + let manifest: SplitJsonManifest = serde_json::from_slice(&bytes).map_err(|err| { + SplitJsonKeyringError::new(format!("failed to deserialize JSON manifest: {err}")) + })?; + if manifest.nodes.is_empty() { + return Err(SplitJsonKeyringError::new("JSON manifest is empty")); + } + Ok(manifest) +} + +fn load_active_revision( + keyring_store: &K, + service: &str, + base_key: &str, +) -> Result, SplitJsonKeyringError> { + let active_key = layout_key(base_key, ACTIVE_REVISION_ENTRY); + load_utf8_secret_from_keyring( + keyring_store, + service, + &active_key, + "active split JSON revision", + ) +} + +fn load_utf8_secret_from_keyring( + keyring_store: &K, + service: &str, + key: &str, + field: &str, +) -> Result, SplitJsonKeyringError> { + let Some(secret) = load_secret_from_keyring(keyring_store, service, key, field)? else { + return Ok(None); + }; + String::from_utf8(secret).map(Some).map_err(|err| { + SplitJsonKeyringError::new(format!( + "failed to decode {field} from keyring as UTF-8: {err}" + )) + }) +} + +fn load_secret_from_keyring( + keyring_store: &K, + service: &str, + key: &str, + field: &str, +) -> Result>, SplitJsonKeyringError> { + keyring_store + .load_secret(service, key) + .map_err(|err| credential_store_error("load", field, err)) +} + +fn save_secret_to_keyring( + keyring_store: &K, + service: &str, + key: &str, + value: &[u8], + field: &str, +) -> Result<(), SplitJsonKeyringError> { + keyring_store + .save_secret(service, key, value) + .map_err(|err| credential_store_error("write", field, err)) +} + +fn delete_keyring_entry( + keyring_store: &K, + service: &str, + key: &str, + field: &str, +) -> Result { + keyring_store + .delete(service, key) + .map_err(|err| credential_store_error("delete", field, err)) +} + +fn credential_store_error( + action: &str, + field: &str, + error: CredentialStoreError, +) -> SplitJsonKeyringError { + SplitJsonKeyringError::new(format!( + "failed to {action} {field} in keyring: {}", + error.message() + )) +} + +fn next_revision() -> String { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_nanos() + .to_string() +} + +fn layout_key(base_key: &str, suffix: &str) -> String { + format!("{base_key}|{LAYOUT_VERSION}|{suffix}") +} + +fn revision_key(base_key: &str, revision: &str, suffix: &str) -> String { + format!("{base_key}|{LAYOUT_VERSION}|{revision}|{suffix}") +} + +fn value_key(base_key: &str, revision: &str, path: &str) -> String { + let encoded_path = encode_path(path); + let suffix = format!("{VALUE_ENTRY_PREFIX}|{encoded_path}"); + revision_key(base_key, revision, &suffix) +} + +fn encode_path(path: &str) -> String { + if path.is_empty() { + return ROOT_PATH_SENTINEL.to_string(); + } + + let mut encoded = String::with_capacity(path.len() * 2); + for byte in path.as_bytes() { + let _ = write!(&mut encoded, "{byte:02x}"); + } + encoded +} + +fn append_json_pointer_token(path: &str, token: &str) -> String { + let escaped = token.replace('~', "~0").replace('/', "~1"); + if path.is_empty() { + format!("/{escaped}") + } else { + format!("{path}/{escaped}") + } +} + +fn decode_json_pointer(pointer: &str) -> Result, SplitJsonKeyringError> { + if pointer.is_empty() { + return Ok(Vec::new()); + } + if !pointer.starts_with('/') { + return Err(SplitJsonKeyringError::new(format!( + "invalid JSON pointer {pointer}: expected leading slash" + ))); + } + + pointer[1..] + .split('/') + .map(unescape_json_pointer_token) + .collect() +} + +fn unescape_json_pointer_token(token: &str) -> Result { + let mut result = String::with_capacity(token.len()); + let mut chars = token.chars(); + + while let Some(ch) = chars.next() { + if ch != '~' { + result.push(ch); + continue; + } + + match chars.next() { + Some('0') => result.push('~'), + Some('1') => result.push('/'), + Some(other) => { + return Err(SplitJsonKeyringError::new(format!( + "invalid JSON pointer escape sequence ~{other}" + ))); + } + None => { + return Err(SplitJsonKeyringError::new( + "invalid JSON pointer escape sequence at end of token", + )); + } + } + } + + Ok(result) +} + +fn parse_array_index(token: &str, pointer: &str) -> Result { + token.parse::().map_err(|err| { + SplitJsonKeyringError::new(format!( + "invalid array index '{token}' in JSON pointer {pointer}: {err}" + )) + }) +} + +fn path_depth(path: &str) -> usize { + path.chars().filter(|ch| *ch == '/').count() +} + +#[cfg(test)] +mod tests { + use super::ACTIVE_REVISION_ENTRY; + use super::LAYOUT_VERSION; + use super::MANIFEST_ENTRY; + use super::ROOT_PATH_SENTINEL; + use super::VALUE_ENTRY_PREFIX; + use super::delete_split_json_from_keyring; + use super::layout_key; + use super::load_split_json_from_keyring; + use super::revision_key; + use super::save_split_json_to_keyring; + use super::value_key; + use crate::tests::MockKeyringStore; + use pretty_assertions::assert_eq; + use serde_json::json; + + const SERVICE: &str = "Test Service"; + const BASE_KEY: &str = "base"; + + #[test] + fn split_json_round_trips_nested_values() { + let store = MockKeyringStore::default(); + let expected = json!({ + "name": "codex", + "enabled": true, + "count": 3, + "nested": { + "items": [null, {"hello": "world"}], + "slash/key": "~value~", + }, + }); + + save_split_json_to_keyring(&store, SERVICE, BASE_KEY, &expected) + .expect("split JSON should save"); + + let loaded = load_split_json_from_keyring(&store, SERVICE, BASE_KEY) + .expect("split JSON should load") + .expect("split JSON should exist"); + assert_eq!(loaded, expected); + } + + #[test] + fn split_json_supports_scalar_root_values() { + let store = MockKeyringStore::default(); + let expected = json!("value"); + + save_split_json_to_keyring(&store, SERVICE, BASE_KEY, &expected) + .expect("split JSON should save"); + + let revision = store + .saved_secret_utf8(&layout_key(BASE_KEY, ACTIVE_REVISION_ENTRY)) + .expect("active revision should exist"); + let root_value_key = revision_key( + BASE_KEY, + &revision, + &format!("{VALUE_ENTRY_PREFIX}|{ROOT_PATH_SENTINEL}"), + ); + assert_eq!( + store.saved_secret_utf8(&root_value_key), + Some("\"value\"".to_string()) + ); + + let loaded = load_split_json_from_keyring(&store, SERVICE, BASE_KEY) + .expect("split JSON should load") + .expect("split JSON should exist"); + assert_eq!(loaded, expected); + } + + #[test] + fn split_json_delete_removes_saved_entries() { + let store = MockKeyringStore::default(); + let expected = json!({ + "token": "secret", + "nested": { + "id": 123, + }, + }); + + save_split_json_to_keyring(&store, SERVICE, BASE_KEY, &expected) + .expect("split JSON should save"); + + let revision = store + .saved_secret_utf8(&layout_key(BASE_KEY, ACTIVE_REVISION_ENTRY)) + .expect("active revision should exist"); + let manifest_key = revision_key(BASE_KEY, &revision, MANIFEST_ENTRY); + let token_key = value_key(BASE_KEY, &revision, "/token"); + let nested_id_key = value_key(BASE_KEY, &revision, "/nested/id"); + + let removed = delete_split_json_from_keyring(&store, SERVICE, BASE_KEY) + .expect("split JSON delete should succeed"); + + assert!(removed); + assert!(!store.contains(&layout_key(BASE_KEY, ACTIVE_REVISION_ENTRY))); + assert!(!store.contains(&manifest_key)); + assert!(!store.contains(&token_key)); + assert!(!store.contains(&nested_id_key)); + } + + #[test] + fn split_json_save_replaces_previous_revision() { + let store = MockKeyringStore::default(); + let first = json!({"value": "first"}); + let second = json!({"value": "second", "extra": 1}); + + save_split_json_to_keyring(&store, SERVICE, BASE_KEY, &first) + .expect("first split JSON save should succeed"); + let first_revision = store + .saved_secret_utf8(&layout_key(BASE_KEY, ACTIVE_REVISION_ENTRY)) + .expect("first revision should exist"); + let first_manifest_key = revision_key(BASE_KEY, &first_revision, MANIFEST_ENTRY); + let first_value_key = value_key(BASE_KEY, &first_revision, "/value"); + + save_split_json_to_keyring(&store, SERVICE, BASE_KEY, &second) + .expect("second split JSON save should succeed"); + let second_revision = store + .saved_secret_utf8(&layout_key(BASE_KEY, ACTIVE_REVISION_ENTRY)) + .expect("second revision should exist"); + let second_manifest_key = revision_key(BASE_KEY, &second_revision, MANIFEST_ENTRY); + + assert_ne!(first_revision, second_revision); + assert!(!store.contains(&first_manifest_key)); + assert!(!store.contains(&first_value_key)); + assert!(store.contains(&second_manifest_key)); + + let loaded = load_split_json_from_keyring(&store, SERVICE, BASE_KEY) + .expect("split JSON should load") + .expect("split JSON should exist"); + assert_eq!(loaded, second); + } + + #[test] + fn split_json_uses_distinct_layout_version() { + assert_eq!(LAYOUT_VERSION, "v1"); + } +} diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index c796d89ced..65fd199517 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -23,6 +23,9 @@ use crate::token_data::TokenData; use codex_app_server_protocol::AuthMode; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; +use codex_keyring_store::delete_split_json_from_keyring; +use codex_keyring_store::load_split_json_from_keyring; +use codex_keyring_store::save_split_json_to_keyring; use once_cell::sync::Lazy; /// Determine where Codex should store CLI auth credentials. @@ -133,21 +136,6 @@ impl AuthStorageBackend for FileAuthStorage { } const KEYRING_SERVICE: &str = "Codex Auth"; -const KEYRING_LAYOUT_VERSION: &str = "v2"; -const KEYRING_ACTIVE_REVISION_ENTRY: &str = "active"; -const KEYRING_MANIFEST_ENTRY: &str = "manifest"; -const KEYRING_OPENAI_API_KEY_ENTRY: &str = "OPENAI_API_KEY"; -const KEYRING_ID_TOKEN_ENTRY: &str = "tokens.id_token"; -const KEYRING_ACCESS_TOKEN_ENTRY: &str = "tokens.access_token"; -const KEYRING_REFRESH_TOKEN_ENTRY: &str = "tokens.refresh_token"; -const KEYRING_ACCOUNT_ID_ENTRY: &str = "tokens.account_id"; -const KEYRING_RECORD_ENTRIES: [&str; 5] = [ - KEYRING_MANIFEST_ENTRY, - KEYRING_OPENAI_API_KEY_ENTRY, - KEYRING_ID_TOKEN_ENTRY, - KEYRING_ACCESS_TOKEN_ENTRY, - KEYRING_REFRESH_TOKEN_ENTRY, -]; // turns codex_home path into a stable, short key string fn compute_store_key(codex_home: &Path) -> std::io::Result { @@ -163,47 +151,6 @@ fn compute_store_key(codex_home: &Path) -> std::io::Result { Ok(format!("cli|{truncated}")) } -fn keyring_layout_key(base_key: &str, suffix: &str) -> String { - format!("{base_key}|{KEYRING_LAYOUT_VERSION}|{suffix}") -} - -fn keyring_revision_key(base_key: &str, revision: &str, suffix: &str) -> String { - format!("{base_key}|{KEYRING_LAYOUT_VERSION}|{revision}|{suffix}") -} - -fn next_keyring_revision() -> String { - Utc::now() - .timestamp_nanos_opt() - .unwrap_or_else(|| Utc::now().timestamp_micros() * 1_000) - .to_string() -} - -#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)] -struct KeyringAuthManifest { - auth_mode: Option, - has_openai_api_key: bool, - has_tokens: bool, - has_account_id: bool, - last_refresh: Option>, -} - -impl From<&AuthDotJson> for KeyringAuthManifest { - fn from(auth: &AuthDotJson) -> Self { - let has_account_id = auth - .tokens - .as_ref() - .and_then(|tokens| tokens.account_id.as_ref()) - .is_some(); - Self { - auth_mode: auth.auth_mode, - has_openai_api_key: auth.openai_api_key.is_some(), - has_tokens: auth.tokens.is_some(), - has_account_id, - last_refresh: auth.last_refresh, - } - } -} - #[derive(Clone, Debug)] struct KeyringAuthStorage { codex_home: PathBuf, @@ -233,133 +180,31 @@ impl KeyringAuthStorage { } } - fn load_secret_from_keyring(&self, key: &str, field: &str) -> std::io::Result>> { - match self.keyring_store.load_secret(KEYRING_SERVICE, key) { - Ok(secret) => Ok(secret), - Err(error) => Err(std::io::Error::other(format!( - "failed to load {field} from keyring: {}", - error.message() - ))), - } - } - - fn load_utf8_secret_from_keyring( - &self, - key: &str, - field: &str, - ) -> std::io::Result> { - let Some(secret) = self.load_secret_from_keyring(key, field)? else { + fn load_split_auth_from_keyring(&self, base_key: &str) -> std::io::Result> { + let Some(value) = + load_split_json_from_keyring(self.keyring_store.as_ref(), KEYRING_SERVICE, base_key) + .map_err(|err| { + std::io::Error::other(format!( + "failed to load split CLI auth from keyring: {err}" + )) + })? + else { return Ok(None); }; - String::from_utf8(secret).map(Some).map_err(|err| { + serde_json::from_value(value).map(Some).map_err(|err| { std::io::Error::other(format!( - "failed to decode {field} from keyring as UTF-8: {err}" + "failed to deserialize CLI auth from keyring: {err}" )) }) } - fn save_secret_to_keyring(&self, key: &str, value: &[u8], field: &str) -> std::io::Result<()> { - match self.keyring_store.save_secret(KEYRING_SERVICE, key, value) { - Ok(()) => Ok(()), - Err(error) => { - let message = format!("failed to write {field} to keyring: {}", error.message()); - warn!("{message}"); - Err(std::io::Error::other(message)) - } - } - } - - fn load_active_revision(&self, base_key: &str) -> std::io::Result> { - let active_key = keyring_layout_key(base_key, KEYRING_ACTIVE_REVISION_ENTRY); - self.load_utf8_secret_from_keyring(&active_key, "active auth revision") - } - - fn load_required_utf8_secret(&self, key: &str, field: &str) -> std::io::Result { - self.load_utf8_secret_from_keyring(key, field)? - .ok_or_else(|| std::io::Error::other(format!("missing {field} in keyring"))) - } - - fn load_manifest( - &self, - base_key: &str, - revision: &str, - ) -> std::io::Result { - let manifest_key = keyring_revision_key(base_key, revision, KEYRING_MANIFEST_ENTRY); - let manifest = self - .load_secret_from_keyring(&manifest_key, "auth manifest")? - .ok_or_else(|| std::io::Error::other("missing auth manifest in keyring"))?; - serde_json::from_slice(&manifest).map_err(|err| { - std::io::Error::other(format!( - "failed to deserialize auth manifest from keyring: {err}" - )) - }) - } - - fn load_v2_from_keyring(&self, base_key: &str, revision: &str) -> std::io::Result { - let manifest = self.load_manifest(base_key, revision)?; - let openai_api_key = if manifest.has_openai_api_key { - let key = keyring_revision_key(base_key, revision, KEYRING_OPENAI_API_KEY_ENTRY); - Some(self.load_required_utf8_secret(&key, "OPENAI_API_KEY")?) - } else { - None - }; - let tokens = if manifest.has_tokens { - let id_token_key = keyring_revision_key(base_key, revision, KEYRING_ID_TOKEN_ENTRY); - let id_token = self.load_required_utf8_secret(&id_token_key, "ID token")?; - let access_token_key = - keyring_revision_key(base_key, revision, KEYRING_ACCESS_TOKEN_ENTRY); - let access_token = self.load_required_utf8_secret(&access_token_key, "access token")?; - let refresh_token_key = - keyring_revision_key(base_key, revision, KEYRING_REFRESH_TOKEN_ENTRY); - let refresh_token = - self.load_required_utf8_secret(&refresh_token_key, "refresh token")?; - let account_id = if manifest.has_account_id { - let account_id_key = - keyring_revision_key(base_key, revision, KEYRING_ACCOUNT_ID_ENTRY); - Some(self.load_required_utf8_secret(&account_id_key, "account ID")?) - } else { - None - }; - Some(TokenData { - id_token: crate::token_data::parse_chatgpt_jwt_claims(&id_token) - .map_err(std::io::Error::other)?, - access_token, - refresh_token, - account_id, - }) - } else { - None - }; - Ok(AuthDotJson { - auth_mode: manifest.auth_mode, - openai_api_key, - tokens, - last_refresh: manifest.last_refresh, - }) - } - fn load_from_keyring(&self, base_key: &str) -> std::io::Result> { - if let Some(revision) = self.load_active_revision(base_key)? { - return self.load_v2_from_keyring(base_key, &revision).map(Some); + if let Some(auth) = self.load_split_auth_from_keyring(base_key)? { + return Ok(Some(auth)); } self.load_legacy_from_keyring(base_key) } - fn write_optional_secret( - &self, - base_key: &str, - revision: &str, - entry: &str, - value: Option<&str>, - field: &str, - ) -> std::io::Result<()> { - if let Some(value) = value { - let key = keyring_revision_key(base_key, revision, entry); - self.save_secret_to_keyring(&key, value.as_bytes(), field)?; - } - Ok(()) - } - fn delete_keyring_entry(&self, key: &str) -> std::io::Result { self.keyring_store .delete(KEYRING_SERVICE, key) @@ -368,95 +213,8 @@ impl KeyringAuthStorage { }) } - fn delete_v2_revision(&self, base_key: &str, revision: &str) -> std::io::Result { - let mut removed = false; - for entry in KEYRING_RECORD_ENTRIES { - let key = keyring_revision_key(base_key, revision, entry); - removed |= self.delete_keyring_entry(&key)?; - } - let account_id_key = keyring_revision_key(base_key, revision, KEYRING_ACCOUNT_ID_ENTRY); - removed |= self.delete_keyring_entry(&account_id_key)?; - Ok(removed) - } - - fn delete_from_keyring_only(&self) -> std::io::Result { - let base_key = compute_store_key(&self.codex_home)?; - let mut removed = false; - if let Some(revision) = self.load_active_revision(&base_key)? { - removed |= self.delete_v2_revision(&base_key, &revision)?; - let active_key = keyring_layout_key(&base_key, KEYRING_ACTIVE_REVISION_ENTRY); - removed |= self.delete_keyring_entry(&active_key)?; - } - removed |= self.delete_keyring_entry(&base_key)?; - Ok(removed) - } - - fn save_v2_to_keyring(&self, base_key: &str, auth: &AuthDotJson) -> std::io::Result<()> { - let previous_revision = match self.load_active_revision(base_key) { - Ok(revision) => revision, - Err(err) => { - warn!("failed to read previous auth revision from keyring: {err}"); - None - } - }; - let revision = next_keyring_revision(); - let manifest = KeyringAuthManifest::from(auth); - - self.write_optional_secret( - base_key, - &revision, - KEYRING_OPENAI_API_KEY_ENTRY, - auth.openai_api_key.as_deref(), - "OPENAI_API_KEY", - )?; - if let Some(tokens) = auth.tokens.as_ref() { - self.write_optional_secret( - base_key, - &revision, - KEYRING_ID_TOKEN_ENTRY, - Some(&tokens.id_token.raw_jwt), - "ID token", - )?; - self.write_optional_secret( - base_key, - &revision, - KEYRING_ACCESS_TOKEN_ENTRY, - Some(&tokens.access_token), - "access token", - )?; - self.write_optional_secret( - base_key, - &revision, - KEYRING_REFRESH_TOKEN_ENTRY, - Some(&tokens.refresh_token), - "refresh token", - )?; - self.write_optional_secret( - base_key, - &revision, - KEYRING_ACCOUNT_ID_ENTRY, - tokens.account_id.as_deref(), - "account ID", - )?; - } - - let manifest_key = keyring_revision_key(base_key, &revision, KEYRING_MANIFEST_ENTRY); - let manifest_bytes = serde_json::to_vec(&manifest).map_err(std::io::Error::other)?; - self.save_secret_to_keyring(&manifest_key, &manifest_bytes, "auth manifest")?; - - let active_key = keyring_layout_key(base_key, KEYRING_ACTIVE_REVISION_ENTRY); - self.save_secret_to_keyring(&active_key, revision.as_bytes(), "active auth revision")?; - - if let Some(previous_revision) = previous_revision - && previous_revision != revision - && let Err(err) = self.delete_v2_revision(base_key, &previous_revision) - { - warn!("failed to remove stale auth revision from keyring: {err}"); - } - if let Err(err) = self.delete_keyring_entry(base_key) { - warn!("failed to remove legacy auth entry from keyring: {err}"); - } - Ok(()) + fn delete_legacy_from_keyring_only(&self, base_key: &str) -> std::io::Result { + self.delete_keyring_entry(base_key) } } @@ -468,7 +226,17 @@ impl AuthStorageBackend for KeyringAuthStorage { fn save(&self, auth: &AuthDotJson) -> std::io::Result<()> { let base_key = compute_store_key(&self.codex_home)?; - self.save_v2_to_keyring(&base_key, auth)?; + let value = serde_json::to_value(auth).map_err(std::io::Error::other)?; + save_split_json_to_keyring( + self.keyring_store.as_ref(), + KEYRING_SERVICE, + &base_key, + &value, + ) + .map_err(|err| std::io::Error::other(format!("failed to write auth to keyring: {err}")))?; + if let Err(err) = self.delete_legacy_from_keyring_only(&base_key) { + warn!("failed to remove legacy auth entries from keyring: {err}"); + } if let Err(err) = delete_file_if_exists(&self.codex_home) { warn!("failed to remove CLI auth fallback file: {err}"); } @@ -476,9 +244,15 @@ impl AuthStorageBackend for KeyringAuthStorage { } fn delete(&self) -> std::io::Result { - let keyring_removed = self.delete_from_keyring_only()?; + let base_key = compute_store_key(&self.codex_home)?; + let split_removed = + delete_split_json_from_keyring(self.keyring_store.as_ref(), KEYRING_SERVICE, &base_key) + .map_err(|err| { + std::io::Error::other(format!("failed to delete auth from keyring: {err}")) + })?; + let legacy_removed = self.delete_legacy_from_keyring_only(&base_key)?; let file_removed = delete_file_if_exists(&self.codex_home)?; - Ok(keyring_removed || file_removed) + Ok(split_removed || legacy_removed || file_removed) } } diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index 4030f6613d..6e24eef146 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -2,7 +2,6 @@ use super::*; use crate::token_data::IdTokenInfo; use anyhow::Context; use base64::Engine; -use pretty_assertions::assert_eq; use serde_json::json; use tempfile::tempdir; @@ -49,6 +48,45 @@ impl KeyringStore for SaveSecretErrorKeyringStore { } } +#[derive(Clone, Debug)] +struct LoadSecretErrorKeyringStore { + inner: MockKeyringStore, +} + +impl KeyringStore for LoadSecretErrorKeyringStore { + fn load(&self, service: &str, account: &str) -> Result, CredentialStoreError> { + self.inner.load(service, account) + } + + fn load_secret( + &self, + _service: &str, + _account: &str, + ) -> Result>, CredentialStoreError> { + Err(CredentialStoreError::new(KeyringError::Invalid( + "error".into(), + "load".into(), + ))) + } + + fn save(&self, service: &str, account: &str, value: &str) -> Result<(), CredentialStoreError> { + self.inner.save(service, account, value) + } + + fn save_secret( + &self, + service: &str, + account: &str, + value: &[u8], + ) -> Result<(), CredentialStoreError> { + self.inner.save_secret(service, account, value) + } + + fn delete(&self, service: &str, account: &str) -> Result { + self.inner.delete(service, account) + } +} + #[tokio::test] async fn file_storage_load_returns_auth_dot_json() -> anyhow::Result<()> { let codex_home = tempdir()?; @@ -141,15 +179,12 @@ fn seed_keyring_and_fallback_auth_file_for_delete( storage: &KeyringAuthStorage, codex_home: &Path, auth: &AuthDotJson, -) -> anyhow::Result<(String, String, PathBuf)> { +) -> anyhow::Result<(String, PathBuf)> { storage.save(auth)?; let base_key = compute_store_key(codex_home)?; - let revision = storage - .load_active_revision(&base_key)? - .context("active auth revision should exist")?; let auth_file = get_auth_file(codex_home); std::fs::write(&auth_file, "stale")?; - Ok((base_key, revision, auth_file)) + Ok((base_key, auth_file)) } fn seed_keyring_with_auth( @@ -172,53 +207,15 @@ fn assert_keyring_saved_auth_and_removed_fallback( codex_home: &Path, expected: &AuthDotJson, ) { - let active_key = keyring_layout_key(base_key, KEYRING_ACTIVE_REVISION_ENTRY); - let revision = mock_keyring - .saved_secret_utf8(&active_key) - .expect("active auth revision should exist"); assert!( mock_keyring.saved_value(base_key).is_none(), "legacy keyring entry should not be used for split auth storage" ); - let manifest_key = keyring_revision_key(base_key, &revision, KEYRING_MANIFEST_ENTRY); - let manifest_bytes = mock_keyring - .saved_secret(&manifest_key) - .expect("auth manifest should exist"); - let manifest: KeyringAuthManifest = - serde_json::from_slice(&manifest_bytes).expect("manifest should deserialize"); - assert_eq!(manifest, KeyringAuthManifest::from(expected)); - - let openai_api_key_key = - keyring_revision_key(base_key, &revision, KEYRING_OPENAI_API_KEY_ENTRY); - assert_eq!( - mock_keyring.saved_secret_utf8(&openai_api_key_key), - expected.openai_api_key - ); - - if let Some(tokens) = expected.tokens.as_ref() { - let id_token_key = keyring_revision_key(base_key, &revision, KEYRING_ID_TOKEN_ENTRY); - assert_eq!( - mock_keyring.saved_secret_utf8(&id_token_key), - Some(tokens.id_token.raw_jwt.clone()) - ); - let access_token_key = - keyring_revision_key(base_key, &revision, KEYRING_ACCESS_TOKEN_ENTRY); - assert_eq!( - mock_keyring.saved_secret_utf8(&access_token_key), - Some(tokens.access_token.clone()) - ); - let refresh_token_key = - keyring_revision_key(base_key, &revision, KEYRING_REFRESH_TOKEN_ENTRY); - assert_eq!( - mock_keyring.saved_secret_utf8(&refresh_token_key), - Some(tokens.refresh_token.clone()) - ); - let account_id_key = keyring_revision_key(base_key, &revision, KEYRING_ACCOUNT_ID_ENTRY); - assert_eq!( - mock_keyring.saved_secret_utf8(&account_id_key), - tokens.account_id.clone() - ); - } + let loaded = load_split_json_from_keyring(mock_keyring, KEYRING_SERVICE, base_key) + .expect("split auth should load from keyring") + .expect("split auth should exist"); + let expected_json = serde_json::to_value(expected).expect("auth should serialize"); + assert_eq!(loaded, expected_json); let auth_file = get_auth_file(codex_home); assert!( !auth_file.exists(), @@ -292,7 +289,7 @@ fn keyring_auth_storage_load_supports_legacy_single_entry() -> anyhow::Result<() } #[test] -fn keyring_auth_storage_load_returns_deserialized_v2_auth() -> anyhow::Result<()> { +fn keyring_auth_storage_load_returns_deserialized_split_auth() -> anyhow::Result<()> { let codex_home = tempdir()?; let mock_keyring = MockKeyringStore::default(); let storage = KeyringAuthStorage::new(codex_home.path().to_path_buf(), Arc::new(mock_keyring)); @@ -353,28 +350,15 @@ fn keyring_auth_storage_delete_removes_keyring_and_file() -> anyhow::Result<()> Arc::new(mock_keyring.clone()), ); let auth = auth_with_prefix("delete"); - let (base_key, revision, auth_file) = + let (base_key, auth_file) = seed_keyring_and_fallback_auth_file_for_delete(&storage, codex_home.path(), &auth)?; let removed = storage.delete()?; assert!(removed, "delete should report removal"); - let active_key = keyring_layout_key(&base_key, KEYRING_ACTIVE_REVISION_ENTRY); assert!( - !mock_keyring.contains(&active_key), - "active revision should be removed" - ); - for entry in KEYRING_RECORD_ENTRIES { - let key = keyring_revision_key(&base_key, &revision, entry); - assert!( - !mock_keyring.contains(&key), - "keyring entry should be removed" - ); - } - let account_id_key = keyring_revision_key(&base_key, &revision, KEYRING_ACCOUNT_ID_ENTRY); - assert!( - !mock_keyring.contains(&account_id_key), - "account id entry should be removed" + load_split_json_from_keyring(&mock_keyring, KEYRING_SERVICE, &base_key)?.is_none(), + "split auth should be removed" ); assert!( !auth_file.exists(), @@ -424,16 +408,10 @@ fn auto_auth_storage_load_uses_file_when_keyring_empty() -> anyhow::Result<()> { fn auto_auth_storage_load_falls_back_when_keyring_errors() -> anyhow::Result<()> { let codex_home = tempdir()?; let mock_keyring = MockKeyringStore::default(); - let storage = AutoAuthStorage::new( - codex_home.path().to_path_buf(), - Arc::new(mock_keyring.clone()), - ); - let key = compute_store_key(codex_home.path())?; - let active_key = keyring_layout_key(&key, KEYRING_ACTIVE_REVISION_ENTRY); - mock_keyring.set_error( - &active_key, - KeyringError::Invalid("error".into(), "load".into()), - ); + let failing_keyring = LoadSecretErrorKeyringStore { + inner: mock_keyring, + }; + let storage = AutoAuthStorage::new(codex_home.path().to_path_buf(), Arc::new(failing_keyring)); let expected = auth_with_prefix("fallback"); storage.file_storage.save(&expected)?; @@ -477,7 +455,6 @@ fn auto_auth_storage_save_falls_back_when_keyring_errors() -> anyhow::Result<()> }; let storage = AutoAuthStorage::new(codex_home.path().to_path_buf(), Arc::new(failing_keyring)); let key = compute_store_key(codex_home.path())?; - let active_key = keyring_layout_key(&key, KEYRING_ACTIVE_REVISION_ENTRY); let auth = auth_with_prefix("fallback"); storage.save(&auth)?; @@ -493,8 +470,8 @@ fn auto_auth_storage_save_falls_back_when_keyring_errors() -> anyhow::Result<()> .context("fallback auth should exist")?; assert_eq!(saved, auth); assert!( - mock_keyring.saved_secret_utf8(&active_key).is_none(), - "keyring should not point to a saved auth revision when save fails" + load_split_json_from_keyring(&mock_keyring, KEYRING_SERVICE, &key)?.is_none(), + "keyring should not point to saved split auth when save fails" ); Ok(()) } @@ -508,7 +485,7 @@ fn auto_auth_storage_delete_removes_keyring_and_file() -> anyhow::Result<()> { Arc::new(mock_keyring.clone()), ); let auth = auth_with_prefix("auto-delete"); - let (base_key, revision, auth_file) = seed_keyring_and_fallback_auth_file_for_delete( + let (base_key, auth_file) = seed_keyring_and_fallback_auth_file_for_delete( storage.keyring_storage.as_ref(), codex_home.path(), &auth, @@ -518,26 +495,8 @@ fn auto_auth_storage_delete_removes_keyring_and_file() -> anyhow::Result<()> { assert!(removed, "delete should report removal"); assert!( - !mock_keyring.contains(&keyring_layout_key( - &base_key, - KEYRING_ACTIVE_REVISION_ENTRY - )), - "active revision should be removed" - ); - for entry in KEYRING_RECORD_ENTRIES { - let key = keyring_revision_key(&base_key, &revision, entry); - assert!( - !mock_keyring.contains(&key), - "keyring entry should be removed" - ); - } - assert!( - !mock_keyring.contains(&keyring_revision_key( - &base_key, - &revision, - KEYRING_ACCOUNT_ID_ENTRY - )), - "account id entry should be removed" + load_split_json_from_keyring(&mock_keyring, KEYRING_SERVICE, &base_key)?.is_none(), + "split auth should be removed" ); assert!( !auth_file.exists(), diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index cdb64ff151..938fc0d716 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -45,6 +45,9 @@ use tracing::warn; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; +use codex_keyring_store::delete_split_json_from_keyring; +use codex_keyring_store::load_split_json_from_keyring; +use codex_keyring_store::save_split_json_to_keyring; use rmcp::transport::auth::AuthorizationManager; use tokio::sync::Mutex; @@ -155,6 +158,14 @@ fn load_oauth_tokens_from_keyring( url: &str, ) -> Result> { let key = compute_store_key(server_name, url)?; + if let Some(value) = load_split_json_from_keyring(keyring_store, KEYRING_SERVICE, &key) + .map_err(|err| Error::msg(err.to_string()))? + { + let mut tokens: StoredOAuthTokens = serde_json::from_value(value) + .context("failed to deserialize OAuth tokens from keyring")?; + refresh_expires_in_from_timestamp(&mut tokens); + return Ok(Some(tokens)); + } match keyring_store.load(KEYRING_SERVICE, &key) { Ok(Some(serialized)) => { let mut tokens: StoredOAuthTokens = serde_json::from_str(&serialized) @@ -191,23 +202,25 @@ fn save_oauth_tokens_with_keyring( server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { - let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?; - + let value = serde_json::to_value(tokens).context("failed to serialize OAuth tokens")?; let key = compute_store_key(server_name, &tokens.url)?; - match keyring_store.save(KEYRING_SERVICE, &key, &serialized) { + match save_split_json_to_keyring(keyring_store, KEYRING_SERVICE, &key, &value) { Ok(()) => { + if let Err(error) = keyring_store.delete(KEYRING_SERVICE, &key) { + warn!( + "failed to remove legacy OAuth tokens from keyring: {}", + error.message() + ); + } if let Err(error) = delete_oauth_tokens_from_file(&key) { warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); } Ok(()) } Err(error) => { - let message = format!( - "failed to write OAuth tokens to keyring: {}", - error.message() - ); + let message = format!("failed to write OAuth tokens to keyring: {error}"); warn!("{message}"); - Err(Error::new(error.into_error()).context(message)) + Err(Error::msg(message)) } } } @@ -244,8 +257,21 @@ fn delete_oauth_tokens_from_keyring_and_file( url: &str, ) -> Result { let key = compute_store_key(server_name, url)?; - let keyring_result = keyring_store.delete(KEYRING_SERVICE, &key); - let keyring_removed = match keyring_result { + let split_removed = match delete_split_json_from_keyring(keyring_store, KEYRING_SERVICE, &key) { + Ok(removed) => removed, + Err(error) => { + let message = error.to_string(); + warn!("failed to delete OAuth tokens from keyring: {message}"); + match store_mode { + OAuthCredentialsStoreMode::Auto | OAuthCredentialsStoreMode::Keyring => { + return Err(Error::msg(message)) + .context("failed to delete OAuth tokens from keyring"); + } + OAuthCredentialsStoreMode::File => false, + } + } + }; + let legacy_removed = match keyring_store.delete(KEYRING_SERVICE, &key) { Ok(removed) => removed, Err(error) => { let message = error.message(); @@ -261,7 +287,7 @@ fn delete_oauth_tokens_from_keyring_and_file( }; let file_removed = delete_oauth_tokens_from_file(&key)?; - Ok(keyring_removed || file_removed) + Ok(split_removed || legacy_removed || file_removed) } #[derive(Clone)] @@ -604,6 +630,8 @@ fn sha_256_prefix(value: &Value) -> Result { mod tests { use super::*; use anyhow::Result; + use codex_keyring_store::CredentialStoreError; + use codex_keyring_store::save_split_json_to_keyring; use keyring::Error as KeyringError; use pretty_assertions::assert_eq; use std::sync::Mutex; @@ -614,6 +642,101 @@ mod tests { use codex_keyring_store::tests::MockKeyringStore; + #[derive(Clone, Debug)] + struct KeyringStoreWithError { + inner: MockKeyringStore, + fail_delete: bool, + fail_load_secret: bool, + fail_save_secret: bool, + } + + impl KeyringStoreWithError { + fn fail_delete(inner: MockKeyringStore) -> Self { + Self { + inner, + fail_delete: true, + fail_load_secret: false, + fail_save_secret: false, + } + } + + fn fail_load_secret(inner: MockKeyringStore) -> Self { + Self { + inner, + fail_delete: false, + fail_load_secret: true, + fail_save_secret: false, + } + } + + fn fail_save_secret(inner: MockKeyringStore) -> Self { + Self { + inner, + fail_delete: false, + fail_load_secret: false, + fail_save_secret: true, + } + } + } + + impl KeyringStore for KeyringStoreWithError { + fn load( + &self, + service: &str, + account: &str, + ) -> Result, CredentialStoreError> { + self.inner.load(service, account) + } + + fn load_secret( + &self, + service: &str, + account: &str, + ) -> Result>, CredentialStoreError> { + if self.fail_load_secret { + return Err(CredentialStoreError::new(KeyringError::Invalid( + "error".into(), + "load".into(), + ))); + } + self.inner.load_secret(service, account) + } + + fn save( + &self, + service: &str, + account: &str, + value: &str, + ) -> Result<(), CredentialStoreError> { + self.inner.save(service, account, value) + } + + fn save_secret( + &self, + service: &str, + account: &str, + value: &[u8], + ) -> Result<(), CredentialStoreError> { + if self.fail_save_secret { + return Err(CredentialStoreError::new(KeyringError::Invalid( + "error".into(), + "save".into(), + ))); + } + self.inner.save_secret(service, account, value) + } + + fn delete(&self, service: &str, account: &str) -> Result { + if self.fail_delete { + return Err(CredentialStoreError::new(KeyringError::Invalid( + "error".into(), + "delete".into(), + ))); + } + self.inner.delete(service, account) + } + } + struct TempCodexHome { _guard: MutexGuard<'static, ()>, _dir: tempfile::TempDir, @@ -651,6 +774,22 @@ mod tests { let store = MockKeyringStore::default(); let tokens = sample_tokens(); let expected = tokens.clone(); + let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; + let value = serde_json::to_value(&tokens)?; + save_split_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?; + + let loaded = + super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)? + .expect("tokens should load from keyring"); + assert_tokens_match_without_expiry(&loaded, &expected); + Ok(()) + } + + #[test] + fn load_oauth_tokens_supports_legacy_single_entry() -> Result<()> { + let _env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let tokens = sample_tokens(); let serialized = serde_json::to_string(&tokens)?; let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; store.save(KEYRING_SERVICE, &key, &serialized)?; @@ -658,7 +797,7 @@ mod tests { let loaded = super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)? .expect("tokens should load from keyring"); - assert_tokens_match_without_expiry(&loaded, &expected); + assert_tokens_match_without_expiry(&loaded, &tokens); Ok(()) } @@ -684,11 +823,9 @@ mod tests { #[test] fn load_oauth_tokens_falls_back_when_keyring_errors() -> Result<()> { let _env = TempCodexHome::new(); - let store = MockKeyringStore::default(); + let store = KeyringStoreWithError::fail_load_secret(MockKeyringStore::default()); let tokens = sample_tokens(); let expected = tokens.clone(); - let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; - store.set_error(&key, KeyringError::Invalid("error".into(), "load".into())); super::save_oauth_tokens_to_file(&tokens)?; @@ -719,18 +856,20 @@ mod tests { let fallback_path = super::fallback_file_path()?; assert!(!fallback_path.exists(), "fallback file should be removed"); - let stored = store.saved_value(&key).expect("value saved to keyring"); - assert_eq!(serde_json::from_str::(&stored)?, tokens); + assert!(store.saved_value(&key).is_none()); + let stored = + super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)? + .expect("value saved to keyring"); + assert_tokens_match_without_expiry(&stored, &tokens); Ok(()) } #[test] fn save_oauth_tokens_writes_fallback_when_keyring_fails() -> Result<()> { let _env = TempCodexHome::new(); - let store = MockKeyringStore::default(); + let mock_keyring = MockKeyringStore::default(); + let store = KeyringStoreWithError::fail_save_secret(mock_keyring.clone()); let tokens = sample_tokens(); - let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; - store.set_error(&key, KeyringError::Invalid("error".into(), "save".into())); super::save_oauth_tokens_with_keyring_with_fallback_to_file( &store, @@ -750,7 +889,7 @@ mod tests { entry.access_token, tokens.token_response.0.access_token().secret().as_str() ); - assert!(store.saved_value(&key).is_none()); + assert!(mock_keyring.saved_value(&key).is_none()); Ok(()) } @@ -759,8 +898,10 @@ mod tests { let _env = TempCodexHome::new(); let store = MockKeyringStore::default(); let tokens = sample_tokens(); - let serialized = serde_json::to_string(&tokens)?; let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; + let value = serde_json::to_value(&tokens)?; + save_split_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?; + let serialized = serde_json::to_string(&tokens)?; store.save(KEYRING_SERVICE, &key, &serialized)?; super::save_oauth_tokens_to_file(&tokens)?; @@ -781,10 +922,13 @@ mod tests { let _env = TempCodexHome::new(); let store = MockKeyringStore::default(); let tokens = sample_tokens(); - let serialized = serde_json::to_string(&tokens)?; let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; - store.save(KEYRING_SERVICE, &key, &serialized)?; - assert!(store.contains(&key)); + let value = serde_json::to_value(&tokens)?; + save_split_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?; + assert!( + super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)? + .is_some() + ); let removed = super::delete_oauth_tokens_from_keyring_and_file( &store, @@ -794,6 +938,10 @@ mod tests { )?; assert!(removed); assert!(!store.contains(&key)); + assert!( + super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)? + .is_none() + ); assert!(!super::fallback_file_path()?.exists()); Ok(()) } @@ -801,10 +949,8 @@ mod tests { #[test] fn delete_oauth_tokens_propagates_keyring_errors() -> Result<()> { let _env = TempCodexHome::new(); - let store = MockKeyringStore::default(); + let store = KeyringStoreWithError::fail_delete(MockKeyringStore::default()); let tokens = sample_tokens(); - let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; - store.set_error(&key, KeyringError::Invalid("error".into(), "delete".into())); super::save_oauth_tokens_to_file(&tokens).unwrap(); let result = super::delete_oauth_tokens_from_keyring_and_file(