public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first pass of merging in warp (doesn't build)

Browse Source
This commit is contained in:
Ryan Ward
2026-07-01 16:08:58 -05:00
parent 2f64909469
commit 4770ac06b5
3662 changed files with 414574 additions and 89772 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/mcp/src/oauth.rs
+402
View File
@@ -0,0 +1,402 @@
use std::collections::HashMap;
use futures::future::BoxFuture;
use oauth2::{RefreshToken, TokenResponse as _};
use rmcp::transport::auth::{
AuthClient, AuthorizationManager, CredentialStore, InMemoryCredentialStore, OAuthClientConfig,
OAuthState, StoredCredentials,
};
use rmcp::transport::{AuthError, AuthorizationSession};
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use url::Url;
use uuid::Uuid;
use warp_core::channel::ChannelState;
use warpui_extras::secure_storage::AppContextExt as _;
pub const TEMPLATABLE_MCP_CREDENTIALS_KEY: &str = "TemplatableMcpCredentials";
pub const FILE_BASED_MCP_CREDENTIALS_KEY: &str = "FileBasedMcpCredentials";
/// The issuer URL for GitHub's OAuth provider.
const GITHUB_ISSUER: &str = "https://github.com/login/oauth";
static GITHUB_OAUTH_SCOPES: [&str; 7] = [
"repo",
"read:org",
"gist",
"notifications",
"user",
"project",
"workflow",
];
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PersistedCredentials {
/// The credential information that `rmcp` wants us to store and retrieve.
#[serde(flatten)]
credentials: StoredCredentials,
/// The client secret for the OAuth application.
///
/// This is needed to properly refresh tokens when using DCR (Dynamic Client Registration),
/// as the server expects the client to provide the secret when refreshing.
client_secret: Option<String>,
}
/// Maps cloud MCP installation UUID to its OAuth credentials in secure storage.
pub type PersistedCredentialsMap = HashMap<Uuid, PersistedCredentials>;
// Maps a consistent hash of the installation to its persisted credentials
pub type FileBasedPersistedCredentialsMap = HashMap<u64, PersistedCredentials>;
pub type PersistCredentialsCallback =
Box<dyn Fn(Uuid, PersistedCredentials) -> BoxFuture<'static, anyhow::Result<()>> + Send>;
pub type RequiresAuthenticationCallback =
Box<dyn Fn(Uuid, String, String) -> BoxFuture<'static, anyhow::Result<()>> + Send>;
pub type AuthenticatedCallback =
Box<dyn Fn(String) -> BoxFuture<'static, anyhow::Result<()>> + Send>;
/// A credential store that wraps [`InMemoryCredentialStore`] and persists token
/// updates to Warp's secure storage via a channel.
///
/// When rmcp auto-refreshes an expired access token at runtime, the rotated
/// tokens are only saved to the in-memory store by default. This wrapper
/// ensures they also get written back to secure storage so they survive app
/// restarts.
struct PersistingCredentialStore {
inner: InMemoryCredentialStore,
client_secret: Option<String>,
persist_tx: async_channel::Sender<PersistedCredentials>,
}
impl PersistingCredentialStore {
/// Per RFC 6749 §6, the authorization server MAY issue a new refresh token on
/// refresh, but is not required to. Many OAuth providers (e.g. Figma) only
/// issue a refresh token on the initial authorization grant and omit it from
/// subsequent refresh responses. If we blindly persist the new token response,
/// the refresh token is lost and the next session (or next in-process refresh)
/// requires a full re-auth.
///
/// When the new response omits a refresh token, carry forward the one already
/// in the store. See: <https://datatracker.ietf.org/doc/html/rfc6749#section-6>
async fn apply_refresh_token_carry_forward(&self, credentials: &mut StoredCredentials) {
if credentials
.token_response
.as_ref()
.is_none_or(|tr| tr.refresh_token().is_some())
{
return;
}
if let Some(prev_rt) = self
.inner
.load()
.await
.ok()
.and_then(|opt| opt)
.and_then(|prev| prev.token_response)
.and_then(|prev_tr| prev_tr.refresh_token().cloned())
{
if let Some(tr) = credentials.token_response.as_mut() {
// Carry forward the existing/previous refresh token, constructing new if needed
tr.set_refresh_token(Some(RefreshToken::new(prev_rt.secret().to_string())));
}
}
}
}
#[async_trait::async_trait]
impl CredentialStore for PersistingCredentialStore {
async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
self.inner.load().await
}
async fn save(&self, mut credentials: StoredCredentials) -> Result<(), AuthError> {
self.apply_refresh_token_carry_forward(&mut credentials)
.await;
self.inner.save(credentials.clone()).await?;
// Only persist credentials if we actually have any.
if credentials.token_response.is_some() {
let _ = self.persist_tx.try_send(PersistedCredentials {
credentials,
client_secret: self.client_secret.clone(),
});
}
Ok(())
}
async fn clear(&self) -> Result<(), AuthError> {
self.inner.clear().await
}
}
/// Installs a [`PersistingCredentialStore`] on the given auth manager so that
/// runtime token auto-refreshes are written back to Warp's secure storage.
///
/// A background tokio task is spawned to receive credential updates and persist
/// them via the provided callback. The task terminates when the auth manager (and
/// thus the credential store's sender) is dropped.
///
/// Note: this store is not responsible for the initial population of credentials.
/// Instead, the caller seeds the inner store with any existing credentials prior
/// to installation (see [`install_persisting_credential_store`]). This store's
/// sole role is to write token updates back to secure storage as they occur.
async fn install_persisting_credential_store(
auth_manager: &mut AuthorizationManager,
persisted_credentials: Option<PersistedCredentials>,
persist_credentials: PersistCredentialsCallback,
installation_uuid: Uuid,
) {
let client_secret = persisted_credentials
.as_ref()
.and_then(|c| c.client_secret.clone());
let in_memory_store = InMemoryCredentialStore::new();
// If we have persisted credentials, populate the backing in-memory store with them.
if let Some(credentials) = persisted_credentials {
let _ = in_memory_store.save(credentials.credentials).await;
}
let (persist_tx, persist_rx) = async_channel::unbounded();
let store = PersistingCredentialStore {
inner: in_memory_store,
client_secret,
persist_tx,
};
auth_manager.set_credential_store(store);
tokio::spawn(async move {
while let Ok(credentials) = persist_rx.recv().await {
if let Err(err) = persist_credentials(installation_uuid, credentials).await {
log::warn!("Failed to persist auto-refreshed MCP credentials: {err:?}");
}
}
});
}
/// Context for OAuth authentication flows.
pub struct AuthContext {
pub oauth_result_rx: async_channel::Receiver<CallbackResult>,
pub uuid: Uuid,
pub persisted_credentials: Option<PersistedCredentials>,
/// Whether the client is running in headless/CLI mode.
pub is_headless: bool,
/// Whether this server was auto-discovered from a repo MCP configuration file.
pub is_file_based: bool,
pub persist_credentials: PersistCredentialsCallback,
pub requires_authentication: RequiresAuthenticationCallback,
pub authenticated: Option<AuthenticatedCallback>,
}
/// Result of OAuth callback.
#[derive(Debug, Clone)]
pub enum CallbackResult {
Success { code: String, csrf_token: String },
Error { error: Option<String> },
}
/// Makes an authenticated client for the given authorization server.
///
/// This takes in the URL of the resource to authenticate for, and uses that
/// to determine the authorization server.
///
/// Upon success, returns the client and a boolean indicating whether the user was required to
/// re-authenticate (e.g. re-log in).
pub async fn make_authenticated_client(
resource_url: &str,
auth_context: AuthContext,
) -> Result<(AuthClient<reqwest::Client>, bool), AuthError> {
let AuthContext {
oauth_result_rx,
uuid,
persisted_credentials,
is_headless,
is_file_based,
persist_credentials,
requires_authentication,
..
} = auth_context;
// Build the redirect URI using the channel's URL scheme.
// Routing data (the server UUID) is passed via the OAuth `state` parameter instead
// of the redirect URI so that the URI exactly matches what is registered during
// Dynamic Client Registration, satisfying RFC 6749 §3.1.2.2 exact-match validation.
let redirect_uri = format!("{}://mcp/oauth2callback", ChannelState::url_scheme());
// Create the auth manager and initialize it with a backing credential store that persists
// new credentials to secure storage.
let client_id = persisted_credentials
.as_ref()
.map(|c| c.credentials.client_id.clone());
let client_secret = persisted_credentials
.as_ref()
.and_then(|c| c.client_secret.clone());
let mut auth_manager = AuthorizationManager::new(resource_url).await?;
install_persisting_credential_store(
&mut auth_manager,
persisted_credentials,
persist_credentials,
uuid,
)
.await;
// If we loaded persisted credentials from the store, and we have a valid access token
// (or successfully refreshed a valid refresh token), we're already authorized and good
// to go.
if auth_manager.initialize_from_store().await? && auth_manager.get_access_token().await.is_ok()
{
if let (Some(client_id), Some(client_secret)) = (client_id, client_secret) {
auth_manager.configure_client(
OAuthClientConfig::new(client_id, redirect_uri.clone())
.with_client_secret(client_secret),
)?;
}
return Ok((AuthClient::new(reqwest::Client::new(), auth_manager), false));
}
// If we're in headless mode and we reach here, it means we either have no credentials
// or the cached credentials failed to refresh. Block interactive OAuth in headless mode.
if is_headless {
if is_file_based {
log::warn!(
"File-based MCP server {uuid} requires OAuth authentication; \
skipping in headless mode. To use this server, authenticate it \
in the Warp desktop app first."
);
}
return Err(AuthError::AuthorizationFailed(
"MCP server requires OAuth authentication. Please authenticate this server in the \
Warp desktop app first, then try again."
.to_string(),
));
}
let metadata = auth_manager.discover_metadata().await?;
// Configure the auth manager's OAuth client using dynamic or static client registration.
let mut oauth_state = if let Some(provider) = metadata
.issuer
.as_deref()
.and_then(ChannelState::mcp_oauth_provider_by_issuer)
{
// Configure the auth manager based on the static MCP configuration for this
// issuer.
auth_manager.set_metadata(metadata);
let scopes = if provider.issuer == GITHUB_ISSUER {
GITHUB_OAUTH_SCOPES
.into_iter()
.map(ToString::to_string)
.collect()
} else {
vec![]
};
auth_manager.configure_client(
OAuthClientConfig::new(provider.client_id, redirect_uri.clone())
.with_client_secret(provider.client_secret)
.with_scopes(scopes),
)?;
// We do a scope "upgrade" with no additional scopes here as it's the easiest way
// to construct an auth URL.
let auth_url = auth_manager.request_scope_upgrade("").await?;
OAuthState::Session(AuthorizationSession::for_scope_upgrade(
auth_manager,
auth_url,
&redirect_uri,
))
} else {
// Try dynamic client registration.
let mut oauth_state = OAuthState::Unauthorized(auth_manager);
oauth_state
.start_authorization(&[], &redirect_uri, Some("Warp"))
.await?;
oauth_state
};
let auth_url = oauth_state.get_authorization_url().await?;
// Extract the CSRF token that rmcp embedded as the `state` query parameter in the
// authorization URL. We register a csrf→uuid mapping on the manager so that
// `handle_oauth_callback` can route the callback to the right server without
// relying on `server_id` being present in the redirect URI.
let csrf_state = Url::parse(&auth_url)
.ok()
.and_then(|u| {
u.query_pairs()
.find(|(k, _)| k == "state")
.map(|(_, v)| v.into_owned())
})
.unwrap_or_default();
if let Err(err) = requires_authentication(uuid, csrf_state, auth_url).await {
log::warn!("Failed to emit RequiresAuthentication state: {err:?}");
}
// Wait for the authorization code from the OAuth callback channel.
let oauth_result = oauth_result_rx
.recv()
.await
.map_err(|e| AuthError::InternalError(e.to_string()))?;
let (code, csrf_token) = match &oauth_result {
CallbackResult::Success { code, csrf_token } => (code, csrf_token),
CallbackResult::Error { error } => {
return Err(AuthError::AuthorizationFailed(
error.as_deref().unwrap_or("unknown error").to_string(),
));
}
};
// Handle the callback with the received authorization code and CSRF token.
oauth_state.handle_callback(code, csrf_token).await?;
let auth_manager = oauth_state.into_authorization_manager().ok_or_else(|| {
AuthError::InternalError("Failed to create authorization manager".to_string())
})?;
Ok((AuthClient::new(reqwest::Client::new(), auth_manager), true))
}
/// Loads credentials from secure storage at the provided key.
pub fn load_credentials_from_secure_storage<T: DeserializeOwned + Default>(
app: &mut warpui::AppContext,
key: &str,
) -> T {
app.secure_storage()
.read_value(key)
.inspect_err(|err| {
if !matches!(err, warpui_extras::secure_storage::Error::NotFound) {
log::warn!("Failed to read MCP credentials from secure storage: {err:#}");
}
})
.ok()
.and_then(|value| serde_json::from_str(&value).ok())
.unwrap_or_default()
}
/// Writes credentials to secure storage at the provided key.
pub fn write_to_secure_storage<T: Serialize>(
app: &mut warpui::AppContext,
key: &str,
credentials: &T,
) {
match serde_json::to_string(credentials) {
Ok(json) => {
app.secure_storage()
.write_value(key, &json)
.inspect_err(|err| {
log::error!("Failed to write MCP credentials to secure storage: {err:#}")
})
.ok();
}
Err(err) => {
log::error!("Failed to serialize MCP credentials for secure storage: {err:#}");
}
}
}
#[cfg(test)]
#[path = "oauth_tests.rs"]
mod tests;