fix: Implement a CredentialStore for auth (#5741)
This commit is contained in:
Generated
+44
-8
@@ -281,9 +281,9 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "async-trait"
|
name = "async-trait"
|
||||||
version = "0.1.87"
|
version = "0.1.89"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "d556ec1359574147ec0c4fc5eb525f3f23263a592b1a9c07e0a75b427de55c97"
|
checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"proc-macro2",
|
"proc-macro2",
|
||||||
"quote",
|
"quote",
|
||||||
@@ -2661,7 +2661,7 @@ dependencies = [
|
|||||||
"rand 0.8.5",
|
"rand 0.8.5",
|
||||||
"regex",
|
"regex",
|
||||||
"reqwest 0.12.12",
|
"reqwest 0.12.12",
|
||||||
"rmcp",
|
"rmcp 0.9.0",
|
||||||
"schemars",
|
"schemars",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
@@ -2711,7 +2711,7 @@ dependencies = [
|
|||||||
"once_cell",
|
"once_cell",
|
||||||
"paste",
|
"paste",
|
||||||
"regex",
|
"regex",
|
||||||
"rmcp",
|
"rmcp 0.9.0",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"tokio",
|
"tokio",
|
||||||
@@ -2750,7 +2750,7 @@ dependencies = [
|
|||||||
"open",
|
"open",
|
||||||
"rand 0.8.5",
|
"rand 0.8.5",
|
||||||
"regex",
|
"regex",
|
||||||
"rmcp",
|
"rmcp 0.9.0",
|
||||||
"rustyline",
|
"rustyline",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
@@ -2805,7 +2805,7 @@ dependencies = [
|
|||||||
"rayon",
|
"rayon",
|
||||||
"regex",
|
"regex",
|
||||||
"reqwest 0.11.27",
|
"reqwest 0.11.27",
|
||||||
"rmcp",
|
"rmcp 0.8.5",
|
||||||
"schemars",
|
"schemars",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
@@ -2856,7 +2856,7 @@ dependencies = [
|
|||||||
"goose-mcp",
|
"goose-mcp",
|
||||||
"http 1.2.0",
|
"http 1.2.0",
|
||||||
"reqwest 0.12.12",
|
"reqwest 0.12.12",
|
||||||
"rmcp",
|
"rmcp 0.9.0",
|
||||||
"schemars",
|
"schemars",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
@@ -5574,6 +5574,29 @@ version = "0.8.5"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "e5947688160b56fb6c827e3c20a72c90392a1d7e9dec74749197aa1780ac42ca"
|
checksum = "e5947688160b56fb6c827e3c20a72c90392a1d7e9dec74749197aa1780ac42ca"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
"base64 0.22.1",
|
||||||
|
"chrono",
|
||||||
|
"futures",
|
||||||
|
"paste",
|
||||||
|
"pin-project-lite",
|
||||||
|
"rmcp-macros 0.8.5",
|
||||||
|
"schemars",
|
||||||
|
"serde",
|
||||||
|
"serde_json",
|
||||||
|
"thiserror 2.0.12",
|
||||||
|
"tokio",
|
||||||
|
"tokio-stream",
|
||||||
|
"tokio-util",
|
||||||
|
"tracing",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "rmcp"
|
||||||
|
version = "0.9.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "acc36ea743d4bbc97e9f3c33bf0b97765a5cf338de3d9c3d2f321a6e38095615"
|
||||||
|
dependencies = [
|
||||||
|
"async-trait",
|
||||||
"base64 0.22.1",
|
"base64 0.22.1",
|
||||||
"chrono",
|
"chrono",
|
||||||
"futures",
|
"futures",
|
||||||
@@ -5583,7 +5606,7 @@ dependencies = [
|
|||||||
"pin-project-lite",
|
"pin-project-lite",
|
||||||
"process-wrap",
|
"process-wrap",
|
||||||
"reqwest 0.12.12",
|
"reqwest 0.12.12",
|
||||||
"rmcp-macros",
|
"rmcp-macros 0.9.0",
|
||||||
"schemars",
|
"schemars",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
@@ -5609,6 +5632,19 @@ dependencies = [
|
|||||||
"syn 2.0.99",
|
"syn 2.0.99",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "rmcp-macros"
|
||||||
|
version = "0.9.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "263caba1c96f2941efca0fdcd97b03f42bcde52d2347d05e5d77c93ab18c5b58"
|
||||||
|
dependencies = [
|
||||||
|
"darling 0.21.0",
|
||||||
|
"proc-macro2",
|
||||||
|
"quote",
|
||||||
|
"serde_json",
|
||||||
|
"syn 2.0.99",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ron"
|
name = "ron"
|
||||||
version = "0.8.1"
|
version = "0.8.1"
|
||||||
|
|||||||
+1
-1
@@ -15,7 +15,7 @@ uninlined_format_args = "allow"
|
|||||||
string_slice = "warn"
|
string_slice = "warn"
|
||||||
|
|
||||||
[workspace.dependencies]
|
[workspace.dependencies]
|
||||||
rmcp = { version = "0.8.5", features = ["schemars", "auth"] }
|
rmcp = { version = "0.9.0", features = ["schemars", "auth"] }
|
||||||
|
|
||||||
# Patch for Windows cross-compilation issue with crunchy
|
# Patch for Windows cross-compilation issue with crunchy
|
||||||
[patch.crates-io]
|
[patch.crates-io]
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ paste = "1.0"
|
|||||||
ctor = "0.2.7"
|
ctor = "0.2.7"
|
||||||
goose = { path = "../goose" }
|
goose = { path = "../goose" }
|
||||||
rmcp = { workspace = true }
|
rmcp = { workspace = true }
|
||||||
async-trait = "0.1.86"
|
async-trait = "0.1.89"
|
||||||
chrono = { version = "0.4", features = ["serde"] }
|
chrono = { version = "0.4", features = ["serde"] }
|
||||||
serde_json = "1.0"
|
serde_json = "1.0"
|
||||||
serde = { version = "1.0", features = ["derive"] }
|
serde = { version = "1.0", features = ["derive"] }
|
||||||
|
|||||||
@@ -43,7 +43,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json",
|
|||||||
tracing-appender = "0.2"
|
tracing-appender = "0.2"
|
||||||
once_cell = "1.20.2"
|
once_cell = "1.20.2"
|
||||||
shlex = "1.3.0"
|
shlex = "1.3.0"
|
||||||
async-trait = "0.1.86"
|
async-trait = "0.1.89"
|
||||||
base64 = "0.22.1"
|
base64 = "0.22.1"
|
||||||
regex = "1.11.1"
|
regex = "1.11.1"
|
||||||
nix = { version = "0.30.1", features = ["process", "signal"] }
|
nix = { version = "0.30.1", features = ["process", "signal"] }
|
||||||
|
|||||||
@@ -33,7 +33,7 @@ reqwest = { version = "0.11", features = [
|
|||||||
"json",
|
"json",
|
||||||
"rustls-tls-native-roots",
|
"rustls-tls-native-roots",
|
||||||
], default-features = false }
|
], default-features = false }
|
||||||
async-trait = "0.1"
|
async-trait = "0.1.89"
|
||||||
chrono = { version = "0.4.38", features = ["serde"] }
|
chrono = { version = "0.4.38", features = ["serde"] }
|
||||||
etcetera = "0.8.0"
|
etcetera = "0.8.0"
|
||||||
tempfile = "3.8"
|
tempfile = "3.8"
|
||||||
|
|||||||
@@ -50,5 +50,5 @@ path = "src/bin/generate_schema.rs"
|
|||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
tower = "0.5"
|
tower = "0.5"
|
||||||
async-trait = "0.1"
|
async-trait = "0.1.89"
|
||||||
tempfile = "3.15.0"
|
tempfile = "3.15.0"
|
||||||
|
|||||||
@@ -49,7 +49,7 @@ serde_urlencoded = "0.7"
|
|||||||
jsonschema = "0.30.0"
|
jsonschema = "0.30.0"
|
||||||
uuid = { version = "1.0", features = ["v4"] }
|
uuid = { version = "1.0", features = ["v4"] }
|
||||||
regex = "1.11.1"
|
regex = "1.11.1"
|
||||||
async-trait = "0.1"
|
async-trait = "0.1.89"
|
||||||
async-stream = "0.3"
|
async-stream = "0.3"
|
||||||
minijinja = { version = "2.10.2", features = ["loader"] }
|
minijinja = { version = "2.10.2", features = ["loader"] }
|
||||||
include_dir = "0.7.4"
|
include_dir = "0.7.4"
|
||||||
|
|||||||
@@ -674,6 +674,7 @@ impl ExtensionManager {
|
|||||||
output_schema: tool.output_schema,
|
output_schema: tool.output_schema,
|
||||||
icons: None,
|
icons: None,
|
||||||
title: None,
|
title: None,
|
||||||
|
meta: None,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,9 +1,11 @@
|
|||||||
|
mod persist;
|
||||||
|
|
||||||
use axum::extract::{Query, State};
|
use axum::extract::{Query, State};
|
||||||
use axum::response::Html;
|
use axum::response::Html;
|
||||||
use axum::routing::get;
|
use axum::routing::get;
|
||||||
use axum::Router;
|
use axum::Router;
|
||||||
use minijinja::render;
|
use minijinja::render;
|
||||||
use rmcp::transport::auth::OAuthState;
|
use rmcp::transport::auth::{CredentialStore, OAuthState, StoredCredentials};
|
||||||
use rmcp::transport::AuthorizationManager;
|
use rmcp::transport::AuthorizationManager;
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
@@ -11,9 +13,7 @@ use std::sync::Arc;
|
|||||||
use tokio::sync::{oneshot, Mutex};
|
use tokio::sync::{oneshot, Mutex};
|
||||||
use tracing::warn;
|
use tracing::warn;
|
||||||
|
|
||||||
use crate::oauth::persist::{clear_credentials, load_cached_state, save_credentials};
|
use crate::oauth::persist::GooseCredentialStore;
|
||||||
|
|
||||||
mod persist;
|
|
||||||
|
|
||||||
const CALLBACK_TEMPLATE: &str = include_str!("oauth_callback.html");
|
const CALLBACK_TEMPLATE: &str = include_str!("oauth_callback.html");
|
||||||
|
|
||||||
@@ -32,18 +32,21 @@ pub async fn oauth_flow(
|
|||||||
mcp_server_url: &String,
|
mcp_server_url: &String,
|
||||||
name: &String,
|
name: &String,
|
||||||
) -> Result<AuthorizationManager, anyhow::Error> {
|
) -> Result<AuthorizationManager, anyhow::Error> {
|
||||||
if let Ok(oauth_state) = load_cached_state(mcp_server_url, name).await {
|
let credential_store = GooseCredentialStore::new(name.clone());
|
||||||
if let Some(authorization_manager) = oauth_state.into_authorization_manager() {
|
let mut auth_manager = AuthorizationManager::new(mcp_server_url).await?;
|
||||||
if authorization_manager.refresh_token().await.is_ok() {
|
auth_manager.set_credential_store(credential_store.clone());
|
||||||
return Ok(authorization_manager);
|
|
||||||
}
|
if auth_manager.initialize_from_store().await? {
|
||||||
|
if auth_manager.refresh_token().await.is_ok() {
|
||||||
|
return Ok(auth_manager);
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Err(e) = clear_credentials(name) {
|
if let Err(e) = credential_store.clear().await {
|
||||||
warn!("error clearing bad credentials: {}", e);
|
warn!("error clearing bad credentials: {}", e);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// No existing credentials or they were invalid - need to do the full oauth flow
|
||||||
let (code_sender, code_receiver) = oneshot::channel::<CallbackParams>();
|
let (code_sender, code_receiver) = oneshot::channel::<CallbackParams>();
|
||||||
let app_state = AppState {
|
let app_state = AppState {
|
||||||
code_receiver: Arc::new(Mutex::new(Some(code_sender))),
|
code_receiver: Arc::new(Mutex::new(Some(code_sender))),
|
||||||
@@ -74,6 +77,7 @@ pub async fn oauth_flow(
|
|||||||
});
|
});
|
||||||
|
|
||||||
let mut oauth_state = OAuthState::new(mcp_server_url, None).await?;
|
let mut oauth_state = OAuthState::new(mcp_server_url, None).await?;
|
||||||
|
|
||||||
let redirect_uri = format!("http://localhost:{}/oauth_callback", used_addr.port());
|
let redirect_uri = format!("http://localhost:{}/oauth_callback", used_addr.port());
|
||||||
oauth_state
|
oauth_state
|
||||||
.start_authorization(&[], redirect_uri.as_str(), Some("goose"))
|
.start_authorization(&[], redirect_uri.as_str(), Some("goose"))
|
||||||
@@ -91,13 +95,20 @@ pub async fn oauth_flow(
|
|||||||
} = code_receiver.await?;
|
} = code_receiver.await?;
|
||||||
oauth_state.handle_callback(&auth_code, &csrf_token).await?;
|
oauth_state.handle_callback(&auth_code, &csrf_token).await?;
|
||||||
|
|
||||||
if let Err(e) = save_credentials(name, &oauth_state).await {
|
let (client_id, token_response) = oauth_state.get_credentials().await?;
|
||||||
warn!("Failed to save credentials: {}", e);
|
|
||||||
}
|
|
||||||
|
|
||||||
let auth_manager = oauth_state
|
let mut auth_manager = oauth_state
|
||||||
.into_authorization_manager()
|
.into_authorization_manager()
|
||||||
.ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?;
|
.ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?;
|
||||||
|
|
||||||
|
credential_store
|
||||||
|
.save(StoredCredentials {
|
||||||
|
client_id,
|
||||||
|
token_response,
|
||||||
|
})
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
auth_manager.set_credential_store(credential_store);
|
||||||
|
|
||||||
Ok(auth_manager)
|
Ok(auth_manager)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,71 +1,54 @@
|
|||||||
use oauth2::{basic::BasicTokenType, EmptyExtraTokenFields, StandardTokenResponse};
|
use rmcp::transport::auth::{AuthError, CredentialStore, StoredCredentials};
|
||||||
use reqwest::IntoUrl;
|
|
||||||
use rmcp::transport::{auth::OAuthState, AuthError};
|
|
||||||
use serde::{Deserialize, Serialize};
|
|
||||||
|
|
||||||
use crate::config::Config;
|
use crate::config::Config;
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
/// Goose-specific credential store that uses the Config system
|
||||||
pub struct SerializableCredentials {
|
///
|
||||||
pub client_id: String,
|
/// This implementation stores OAuth credentials in the goose configuration
|
||||||
pub token_response: Option<StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>>,
|
/// system, which handles secure storage (e.g., keychain integration).
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct GooseCredentialStore {
|
||||||
|
name: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
fn secret_key(name: &str) -> String {
|
impl GooseCredentialStore {
|
||||||
format!("oauth_creds_{name}")
|
pub fn new(name: String) -> Self {
|
||||||
}
|
Self { name }
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn save_credentials(
|
fn secret_key(&self) -> String {
|
||||||
name: &str,
|
format!("oauth_creds_{}", self.name)
|
||||||
oauth_state: &OAuthState,
|
}
|
||||||
) -> Result<(), Box<dyn std::error::Error>> {
|
}
|
||||||
let config = Config::global();
|
|
||||||
let (client_id, token_response) = oauth_state.get_credentials().await?;
|
#[async_trait::async_trait]
|
||||||
|
impl CredentialStore for GooseCredentialStore {
|
||||||
let credentials = SerializableCredentials {
|
async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
|
||||||
client_id,
|
let config = Config::global();
|
||||||
token_response,
|
let key = self.secret_key();
|
||||||
};
|
|
||||||
|
match config.get_secret::<StoredCredentials>(&key) {
|
||||||
let key = secret_key(name);
|
Ok(credentials) => Ok(Some(credentials)),
|
||||||
config.set_secret(&key, &credentials)?;
|
Err(_) => Ok(None), // No credentials found
|
||||||
|
}
|
||||||
Ok(())
|
}
|
||||||
}
|
|
||||||
|
async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> {
|
||||||
async fn load_credentials(
|
let config = Config::global();
|
||||||
name: &str,
|
let key = self.secret_key();
|
||||||
) -> Result<SerializableCredentials, Box<dyn std::error::Error>> {
|
|
||||||
let config = Config::global();
|
config
|
||||||
let key = secret_key(name);
|
.set_secret(&key, &credentials)
|
||||||
let credentials: SerializableCredentials = config.get_secret(&key)?;
|
.map_err(|e| AuthError::InternalError(format!("Failed to save credentials: {}", e)))
|
||||||
|
}
|
||||||
Ok(credentials)
|
|
||||||
}
|
async fn clear(&self) -> Result<(), AuthError> {
|
||||||
|
let config = Config::global();
|
||||||
pub fn clear_credentials(name: &str) -> Result<(), Box<dyn std::error::Error>> {
|
let key = self.secret_key();
|
||||||
let config = Config::global();
|
|
||||||
|
config
|
||||||
Ok(config.delete_secret(&secret_key(name))?)
|
.delete_secret(&key)
|
||||||
}
|
.map_err(|e| AuthError::InternalError(format!("Failed to clear credentials: {}", e)))
|
||||||
|
|
||||||
pub async fn load_cached_state<U: IntoUrl>(
|
|
||||||
base_url: U,
|
|
||||||
name: &str,
|
|
||||||
) -> Result<OAuthState, AuthError> {
|
|
||||||
let credentials = load_credentials(name)
|
|
||||||
.await
|
|
||||||
.map_err(|e| AuthError::InternalError(format!("Failed to load credentials: {}", e)))?;
|
|
||||||
|
|
||||||
if let Some(token_response) = credentials.token_response {
|
|
||||||
let mut oauth_state = OAuthState::new(base_url, None).await?;
|
|
||||||
oauth_state
|
|
||||||
.set_credentials(&credentials.client_id, token_response)
|
|
||||||
.await?;
|
|
||||||
Ok(oauth_state)
|
|
||||||
} else {
|
|
||||||
Err(AuthError::InternalError(
|
|
||||||
"No token response in cached credentials".to_string(),
|
|
||||||
))
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4643,6 +4643,10 @@
|
|||||||
"inputSchema"
|
"inputSchema"
|
||||||
],
|
],
|
||||||
"properties": {
|
"properties": {
|
||||||
|
"_meta": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": true
|
||||||
|
},
|
||||||
"annotations": {
|
"annotations": {
|
||||||
"anyOf": [
|
"anyOf": [
|
||||||
{
|
{
|
||||||
|
|||||||
@@ -810,6 +810,9 @@ export type TokenState = {
|
|||||||
};
|
};
|
||||||
|
|
||||||
export type Tool = {
|
export type Tool = {
|
||||||
|
_meta?: {
|
||||||
|
[key: string]: unknown;
|
||||||
|
};
|
||||||
annotations?: ToolAnnotations | {
|
annotations?: ToolAnnotations | {
|
||||||
[key: string]: unknown;
|
[key: string]: unknown;
|
||||||
};
|
};
|
||||||
|
|||||||
Reference in New Issue
Block a user