mirror of
https://github.com/openai/codex.git
synced 2026-09-13 11:47:17 +00:00
rmcp-client to use json store api
This commit is contained in:
@@ -45,6 +45,9 @@ use tracing::warn;
|
||||
|
||||
use codex_keyring_store::DefaultKeyringStore;
|
||||
use codex_keyring_store::KeyringStore;
|
||||
use codex_keyring_store::delete_json_from_keyring;
|
||||
use codex_keyring_store::load_json_from_keyring;
|
||||
use codex_keyring_store::save_json_to_keyring;
|
||||
use rmcp::transport::auth::AuthorizationManager;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
@@ -155,16 +158,15 @@ fn load_oauth_tokens_from_keyring<K: KeyringStore>(
|
||||
url: &str,
|
||||
) -> Result<Option<StoredOAuthTokens>> {
|
||||
let key = compute_store_key(server_name, url)?;
|
||||
match keyring_store.load(KEYRING_SERVICE, &key) {
|
||||
Ok(Some(serialized)) => {
|
||||
let mut tokens: StoredOAuthTokens = serde_json::from_str(&serialized)
|
||||
.context("failed to deserialize OAuth tokens from keyring")?;
|
||||
refresh_expires_in_from_timestamp(&mut tokens);
|
||||
Ok(Some(tokens))
|
||||
}
|
||||
Ok(None) => Ok(None),
|
||||
Err(error) => Err(Error::new(error.into_error())),
|
||||
}
|
||||
let Some(value) = load_json_from_keyring(keyring_store, KEYRING_SERVICE, &key)
|
||||
.map_err(|err| Error::msg(err.to_string()))?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut tokens: StoredOAuthTokens =
|
||||
serde_json::from_value(value).context("failed to deserialize OAuth tokens from keyring")?;
|
||||
refresh_expires_in_from_timestamp(&mut tokens);
|
||||
Ok(Some(tokens))
|
||||
}
|
||||
|
||||
pub fn save_oauth_tokens(
|
||||
@@ -191,10 +193,9 @@ fn save_oauth_tokens_with_keyring<K: KeyringStore>(
|
||||
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_json_to_keyring(keyring_store, KEYRING_SERVICE, &key, &value) {
|
||||
Ok(()) => {
|
||||
if let Err(error) = delete_oauth_tokens_from_file(&key) {
|
||||
warn!("failed to remove OAuth tokens from fallback storage: {error:?}");
|
||||
@@ -202,12 +203,9 @@ fn save_oauth_tokens_with_keyring<K: KeyringStore>(
|
||||
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,22 +242,20 @@ fn delete_oauth_tokens_from_keyring_and_file<K: KeyringStore>(
|
||||
url: &str,
|
||||
) -> Result<bool> {
|
||||
let key = compute_store_key(server_name, url)?;
|
||||
let keyring_result = keyring_store.delete(KEYRING_SERVICE, &key);
|
||||
let keyring_removed = match keyring_result {
|
||||
let keyring_removed = match delete_json_from_keyring(keyring_store, KEYRING_SERVICE, &key) {
|
||||
Ok(removed) => removed,
|
||||
Err(error) => {
|
||||
let message = error.message();
|
||||
let message = error.to_string();
|
||||
warn!("failed to delete OAuth tokens from keyring: {message}");
|
||||
match store_mode {
|
||||
OAuthCredentialsStoreMode::Auto | OAuthCredentialsStoreMode::Keyring => {
|
||||
return Err(error.into_error())
|
||||
return Err(Error::msg(message))
|
||||
.context("failed to delete OAuth tokens from keyring");
|
||||
}
|
||||
OAuthCredentialsStoreMode::File => false,
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let file_removed = delete_oauth_tokens_from_file(&key)?;
|
||||
Ok(keyring_removed || file_removed)
|
||||
}
|
||||
@@ -604,6 +600,9 @@ fn sha_256_prefix(value: &Value) -> Result<String> {
|
||||
mod tests {
|
||||
use super::*;
|
||||
use anyhow::Result;
|
||||
use codex_keyring_store::CredentialStoreError;
|
||||
use codex_keyring_store::load_json_from_keyring;
|
||||
use codex_keyring_store::save_json_to_keyring;
|
||||
use keyring::Error as KeyringError;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::sync::Mutex;
|
||||
@@ -614,6 +613,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<Option<String>, CredentialStoreError> {
|
||||
self.inner.load(service, account)
|
||||
}
|
||||
|
||||
fn load_secret(
|
||||
&self,
|
||||
service: &str,
|
||||
account: &str,
|
||||
) -> Result<Option<Vec<u8>>, 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<bool, CredentialStoreError> {
|
||||
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,9 +745,9 @@ mod tests {
|
||||
let store = MockKeyringStore::default();
|
||||
let tokens = sample_tokens();
|
||||
let expected = tokens.clone();
|
||||
let serialized = serde_json::to_string(&tokens)?;
|
||||
let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
|
||||
store.save(KEYRING_SERVICE, &key, &serialized)?;
|
||||
let value = serde_json::to_value(&tokens)?;
|
||||
save_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?;
|
||||
|
||||
let loaded =
|
||||
super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)?
|
||||
@@ -662,6 +756,31 @@ mod tests {
|
||||
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)?;
|
||||
|
||||
let loaded =
|
||||
super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)?;
|
||||
|
||||
#[cfg(not(windows))]
|
||||
{
|
||||
let loaded = loaded.expect("tokens should load from keyring");
|
||||
assert_tokens_match_without_expiry(&loaded, &tokens);
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
{
|
||||
assert!(loaded.is_none());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_oauth_tokens_falls_back_when_missing_in_keyring() -> Result<()> {
|
||||
let _env = TempCodexHome::new();
|
||||
@@ -684,11 +803,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 +836,29 @@ 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::<StoredOAuthTokens>(&stored)?, tokens);
|
||||
#[cfg(windows)]
|
||||
assert!(
|
||||
store.saved_secret(&key).is_none(),
|
||||
"windows should not store the full JSON record under the base key"
|
||||
);
|
||||
#[cfg(not(windows))]
|
||||
assert!(
|
||||
store.saved_secret(&key).is_some(),
|
||||
"non-windows should store the full JSON record as one secret"
|
||||
);
|
||||
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,18 +878,22 @@ 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());
|
||||
assert!(
|
||||
load_json_from_keyring(&mock_keyring, KEYRING_SERVICE, &key)?.is_none(),
|
||||
"keyring should not point at saved OAuth tokens when save fails"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn delete_oauth_tokens_removes_all_storage() -> Result<()> {
|
||||
fn delete_oauth_tokens_removes_active_storage() -> 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)?;
|
||||
let value = serde_json::to_value(&tokens)?;
|
||||
save_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?;
|
||||
super::save_oauth_tokens_to_file(&tokens)?;
|
||||
|
||||
let removed = super::delete_oauth_tokens_from_keyring_and_file(
|
||||
@@ -771,7 +903,10 @@ mod tests {
|
||||
&tokens.url,
|
||||
)?;
|
||||
assert!(removed);
|
||||
assert!(!store.contains(&key));
|
||||
assert!(
|
||||
load_json_from_keyring(&store, KEYRING_SERVICE, &key)?.is_none(),
|
||||
"keyring entry should be removed"
|
||||
);
|
||||
assert!(!super::fallback_file_path()?.exists());
|
||||
Ok(())
|
||||
}
|
||||
@@ -781,10 +916,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_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,
|
||||
@@ -793,7 +931,14 @@ mod tests {
|
||||
&tokens.url,
|
||||
)?;
|
||||
assert!(removed);
|
||||
assert!(!store.contains(&key));
|
||||
assert!(
|
||||
load_json_from_keyring(&store, KEYRING_SERVICE, &key)?.is_none(),
|
||||
"keyring entry should be removed"
|
||||
);
|
||||
assert!(
|
||||
super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)?
|
||||
.is_none()
|
||||
);
|
||||
assert!(!super::fallback_file_path()?.exists());
|
||||
Ok(())
|
||||
}
|
||||
@@ -801,10 +946,11 @@ 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()));
|
||||
let value = serde_json::to_value(&tokens)?;
|
||||
save_json_to_keyring(&store, KEYRING_SERVICE, &key, &value)?;
|
||||
super::save_oauth_tokens_to_file(&tokens).unwrap();
|
||||
|
||||
let result = super::delete_oauth_tokens_from_keyring_and_file(
|
||||
|
||||
Reference in New Issue
Block a user