rust_mcp_extra/auth_provider/
work_os.rs

1//! # WorkOS AuthKit OAuth2 Provider for MCP Servers
2//!
3//! This module implements an OAuth2 specifically designed to integrate
4//! [WorkOS AuthKit](https://workos.com/docs/authkit) as the identity
5//! provider (IdP) in an MCP (Model Context Protocol) server ecosystem.
6//!
7//! It enables your MCP server to:
8//! - Expose standard OAuth2/.well-known endpoints
9//! - Serve authorization server metadata (`/.well-known/oauth-authorization-server`)
10//! - Serve protected resource metadata (custom per MCP)
11//! - Verify incoming access tokens using JWKs + UserInfo endpoint validation
12//!
13//! ## Features
14//!
15//! - Zero-downtime token verification with cached JWKs
16//! - Automatic construction of OAuth2 discovery documents
17//! - Built-in CORS support for metadata endpoints
18//! - Pluggable into `rust-mcp-sdk`'s authentication system via the `AuthProvider` trait
19//!
20//! ## Example
21//!
22//! ```rust,ignore
23//!
24//! let auth_provider = WorkOsAuthProvider::new(WorkOSAuthOptions {
25//!     // Your AuthKit app domain (found in WorkOS dashboard)
26//!     authkit_domain: "https://your-app.authkit.app".to_string(),
27//!     // Base URL of your MCP server (used to build protected resource metadata URL)
28//!     mcp_server_url: "http://localhost:3000/mcp".to_string(),
29//! })?;
30//!
31//! // Register in your MCP server
32//! let server = rust_mcp_axum::create_axum_server(
33//! server_details,
34//! handler,
35//! rust_mcp_axum::AxumServerOptions {
36//!     host: "localhost".to_string(),
37//!     port: 3000,
38//!     auth: Some(Arc::new(auth_provider)),
39//!     ..Default::default()
40//! });
41//! ```
42use super::resolve_audience;
43use crate::token_verifier::{
44    GenericOauthTokenVerifier, TokenVerifierOptions, VerificationStrategies,
45};
46use async_trait::async_trait;
47use bytes::Bytes;
48use http::{header::CONTENT_TYPE, StatusCode};
49use http_body_util::{BodyExt, Full};
50use rust_mcp_sdk::{
51    auth::{
52        create_discovery_endpoints, Audience, AuthInfo, AuthMetadataBuilder, AuthProvider,
53        AuthenticationError, AuthorizationServerMetadata, OauthEndpoint,
54        OauthProtectedResourceMetadata, OauthTokenVerifier,
55    },
56    error::McpSdkError,
57    mcp_http::{
58        middleware::CorsMiddleware, GenericBody, GenericBodyExt, McpAppState, McpHttpError,
59        McpHttpResult, Middleware,
60    },
61    mcp_server::join_url,
62};
63use std::{collections::HashMap, sync::Arc, vec};
64
65static SCOPES_SUPPORTED: &[&str] = &["email", "offline_access", "openid", "profile"];
66
67/// Configuration options for the WorkOS AuthKit OAuth provider.
68pub struct WorkOSAuthOptions<'a> {
69    pub authkit_domain: String,
70    pub mcp_server_url: String,
71    pub required_scopes: Option<Vec<&'a str>>,
72    pub token_verifier: Option<Box<dyn OauthTokenVerifier>>,
73    pub resource_name: Option<String>,
74    pub resource_documentation: Option<String>,
75    /// Audience to validate the token's `aud` claim against.
76    /// When `None`, the audience defaults to `mcp_server_url` (the resource
77    /// identifier), unless `disable_audience_validation` is set.
78    pub validate_audience: Option<Audience>,
79    /// Disables audience validation entirely. Strongly discouraged: without it a
80    /// token issued for another resource can be replayed against this server.
81    pub disable_audience_validation: bool,
82}
83
84/// WorkOS AuthKit integration implementing `AuthProvider` for MCP servers.
85///
86/// This provider makes your MCP server compatible with clients that expect standard
87/// OAuth2 authorization server and protected resource discovery endpoints when using
88/// WorkOS AuthKit as the identity provider.
89pub struct WorkOsAuthProvider {
90    auth_server_meta: AuthorizationServerMetadata,
91    protected_resource_meta: OauthProtectedResourceMetadata,
92    endpoint_map: HashMap<String, OauthEndpoint>,
93    protected_resource_metadata_url: String,
94    token_verifier: Box<dyn OauthTokenVerifier>,
95}
96
97impl WorkOsAuthProvider {
98    /// Creates a new `WorkOsAuthProvider` instance.
99    ///
100    /// This performs:
101    /// - Validation and parsing of URLs
102    /// - Construction of OAuth2 metadata documents
103    /// - Setup of token verification using JWKs and UserInfo endpoint
104    ///
105    /// /// # Example
106    ///
107    /// ```rust,ignore
108    /// use rust_mcp_extra::auth_provider::work_os::{WorkOSAuthOptions, WorkOsAuthProvider};
109    ///
110    /// let auth_provider = WorkOsAuthProvider::new(WorkOSAuthOptions {
111    ///    authkit_domain: "https://your-app.authkit.app".to_string(),
112    ///    mcp_server_url: "http://localhost:3000/mcp".to_string(),
113    /// })?;
114    ///
115    pub fn new(mut options: WorkOSAuthOptions) -> Result<Self, McpSdkError> {
116        let (endpoint_map, protected_resource_metadata_url) =
117            create_discovery_endpoints(&options.mcp_server_url)?;
118
119        let required_scopes = options.required_scopes.take();
120        let scopes_supported = required_scopes.clone().unwrap_or(SCOPES_SUPPORTED.to_vec());
121
122        let mut builder = AuthMetadataBuilder::new(&options.mcp_server_url)
123            .issuer(&options.authkit_domain)
124            .authorization_servers(vec![&options.authkit_domain])
125            .authorization_endpoint("/oauth2/authorize")
126            .introspection_endpoint("/oauth2/introspection")
127            .registration_endpoint("/oauth2/register")
128            .token_endpoint("/oauth2/token")
129            .jwks_uri("/oauth2/jwks")
130            .scopes_supported(scopes_supported);
131
132        if let Some(scopes) = required_scopes {
133            builder = builder.reqquired_scopes(scopes)
134        }
135        if let Some(resource_name) = options.resource_name.as_ref() {
136            builder = builder.resource_name(resource_name)
137        }
138        if let Some(resource_documentation) = options.resource_documentation.as_ref() {
139            builder = builder.service_documentation(resource_documentation)
140        }
141
142        let (auth_server_meta, protected_resource_meta) = builder.build()?;
143
144        let Some(jwks_uri) = auth_server_meta.jwks_uri.as_ref().map(|s| s.to_string()) else {
145            return Err(McpSdkError::Internal {
146                description: "jwks_uri is not defined!".to_string(),
147            });
148        };
149
150        let userinfo_uri = join_url(&auth_server_meta.issuer, "oauth2/userinfo")
151            .map_err(|err| McpSdkError::Internal {
152                description: format!("invalid userinfo url :{err}"),
153            })?
154            .to_string();
155
156        let validate_audience = resolve_audience(
157            options.disable_audience_validation,
158            options.validate_audience.take(),
159            &options.mcp_server_url,
160        );
161
162        let token_verifier: Box<dyn OauthTokenVerifier> = match options.token_verifier {
163            Some(verifier) => verifier,
164            None => Box::new(GenericOauthTokenVerifier::new(TokenVerifierOptions {
165                strategies: vec![
166                    VerificationStrategies::JWKs { jwks_uri },
167                    VerificationStrategies::UserInfo { userinfo_uri },
168                ],
169                validate_audience,
170                validate_issuer: Some(options.authkit_domain.clone()),
171                cache_capacity: None,
172            })?),
173        };
174
175        Ok(Self {
176            endpoint_map,
177            protected_resource_metadata_url,
178            token_verifier,
179            auth_server_meta,
180            protected_resource_meta,
181        })
182    }
183
184    /// Helper to build JSON response for authorization server metadata with CORS.
185    fn handle_authorization_server_metadata(
186        response_str: String,
187    ) -> McpHttpResult<http::Response<GenericBody>> {
188        let body = Full::new(Bytes::from(response_str))
189            .map_err(|err| McpHttpError::HttpError(err.to_string()))
190            .boxed();
191        http::Response::builder()
192            .status(StatusCode::OK)
193            .header(CONTENT_TYPE, "application/json")
194            .body(body)
195            .map_err(|err| McpHttpError::HttpError(err.to_string()))
196    }
197
198    /// Helper to build JSON response for protected resource metadata with permissive CORS.
199    fn handle_protected_resource_metadata(
200        response_str: String,
201    ) -> McpHttpResult<http::Response<GenericBody>> {
202        use http_body_util::BodyExt;
203
204        let body = Full::new(Bytes::from(response_str))
205            .map_err(|err| McpHttpError::HttpError(err.to_string()))
206            .boxed();
207        http::Response::builder()
208            .status(StatusCode::OK)
209            .header(CONTENT_TYPE, "application/json")
210            .body(body)
211            .map_err(|err| McpHttpError::HttpError(err.to_string()))
212    }
213}
214
215#[async_trait]
216impl AuthProvider for WorkOsAuthProvider {
217    /// Returns the map of supported OAuth discovery endpoints.
218    fn auth_endpoints(&self) -> Option<&HashMap<String, OauthEndpoint>> {
219        Some(&self.endpoint_map)
220    }
221
222    /// Handles incoming requests to OAuth metadata endpoints.
223    async fn handle_request(
224        &self,
225        request: http::Request<&str>,
226        state: Arc<McpAppState>,
227    ) -> Result<http::Response<GenericBody>, McpHttpError> {
228        let Some(endpoint) = self.endpoint_type(&request) else {
229            return http::Response::builder()
230                .status(StatusCode::NOT_FOUND)
231                .body(GenericBody::empty())
232                .map_err(|err| McpHttpError::HttpError(err.to_string()));
233        };
234
235        // return early if method is not allowed
236        if let Some(response) = self.validate_allowed_methods(endpoint, request.method()) {
237            return Ok(response);
238        }
239
240        match endpoint {
241            OauthEndpoint::AuthorizationServerMetadata => {
242                let json_payload = serde_json::to_string(&self.auth_server_meta)
243                    .map_err(|err| McpHttpError::HttpError(err.to_string()))?;
244                let cors = &CorsMiddleware::default();
245                cors.handle(
246                    request,
247                    state,
248                    Box::new(move |_req, _state| {
249                        Box::pin(
250                            async move { Self::handle_authorization_server_metadata(json_payload) },
251                        )
252                    }),
253                )
254                .await
255            }
256            OauthEndpoint::ProtectedResourceMetadata => {
257                let json_payload = serde_json::to_string(&self.protected_resource_meta)
258                    .map_err(|err| McpHttpError::HttpError(err.to_string()))?;
259                let cors = &CorsMiddleware::default();
260                cors.handle(
261                    request,
262                    state,
263                    Box::new(move |_req, _state| {
264                        Box::pin(
265                            async move { Self::handle_protected_resource_metadata(json_payload) },
266                        )
267                    }),
268                )
269                .await
270            }
271            _ => Ok(GenericBody::create_404_response()),
272        }
273    }
274
275    /// Verifies an access token using JWKs and optional UserInfo validation.
276    ///
277    /// Returns authenticated `AuthInfo` on success.
278    async fn verify_token(&self, access_token: String) -> Result<AuthInfo, AuthenticationError> {
279        self.token_verifier.verify_token(access_token).await
280    }
281
282    /// Returns the full URL to the protected resource metadata document.
283    fn protected_resource_metadata_url(&self) -> Option<&str> {
284        Some(self.protected_resource_metadata_url.as_str())
285    }
286}
rust_mcp_extra/auth_provider/work_os.rs

rust_mcp_extra/auth_provider/
work_os.rs
