tower_mcp/
auth.rs

1//! Authentication middleware helpers for MCP servers
2//!
3//! This module provides helper types and layers for common authentication patterns.
4//! Since tower-mcp is built on Tower, standard tower middleware can be used directly.
5//!
6//! # Patterns
7//!
8//! ## API Key Authentication
9//!
10//! ```rust,ignore
11//! // Requires the `http` feature
12//! use tower_mcp::auth::{AuthConfig, ApiKeyValidator};
13//! use tower_mcp::{McpRouter, HttpTransport};
14//! use std::sync::Arc;
15//!
16//! // Simple in-memory API key validator
17//! let valid_keys = vec!["sk-test-key-123".to_string()];
18//! let validator = ApiKeyValidator::new(valid_keys);
19//!
20//! let router = McpRouter::new().server_info("my-server", "1.0.0");
21//! let transport = HttpTransport::new(router);
22//!
23//! // The auth layer extracts the key from the Authorization header
24//! // and validates it using the provided validator
25//! ```
26//!
27//! ## Bearer Token Authentication
28//!
29//! For OAuth2/JWT tokens, use the `BearerTokenValidator` trait to implement
30//! custom validation logic (e.g., JWT verification, token introspection).
31//!
32//! ## Custom Authentication
33//!
34//! You can implement custom auth by creating a Tower layer. See the examples
35//! directory for a complete example.
36
37use std::collections::HashSet;
38use std::future::Future;
39use std::sync::Arc;
40
41use tower::Layer;
42#[cfg(feature = "http")]
43use tower::ServiceExt;
44
45/// Result of an authentication attempt
46#[derive(Debug, Clone)]
47#[non_exhaustive]
48pub enum AuthResult {
49    /// Authentication succeeded with optional user/client info
50    Authenticated(Option<AuthInfo>),
51    /// Authentication failed with a reason
52    Failed(AuthError),
53}
54
55/// Information about an authenticated client
56#[derive(Debug, Clone)]
57pub struct AuthInfo {
58    /// Client/user identifier
59    pub client_id: String,
60    /// Optional additional claims or metadata
61    pub claims: Option<serde_json::Value>,
62}
63
64/// Authentication error
65#[derive(Debug, Clone)]
66pub struct AuthError {
67    /// Error code (e.g., "invalid_token", "expired_token")
68    pub code: String,
69    /// Human-readable error message
70    pub message: String,
71}
72
73impl std::fmt::Display for AuthError {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        write!(f, "{}: {}", self.code, self.message)
76    }
77}
78
79impl std::error::Error for AuthError {}
80
81// =============================================================================
82// Validation Trait
83// =============================================================================
84
85/// Trait for validating authentication credentials.
86///
87/// Implement this trait to provide custom authentication logic for use
88/// with [`AuthLayer`] and [`AuthService`].
89///
90/// The credential string passed to [`validate`](Validate::validate) is the
91/// value extracted from the configured request header after parsing
92/// (e.g., the token portion of `"Bearer sk-123"`).
93///
94/// # Example
95///
96/// ```rust
97/// use tower_mcp::auth::{Validate, AuthResult, AuthInfo, AuthError};
98///
99/// #[derive(Clone)]
100/// struct MyValidator;
101///
102/// impl Validate for MyValidator {
103///     async fn validate(&self, credential: &str) -> AuthResult {
104///         if credential.starts_with("sk-") {
105///             AuthResult::Authenticated(Some(AuthInfo {
106///                 client_id: credential.to_string(),
107///                 claims: None,
108///             }))
109///         } else {
110///             AuthResult::Failed(AuthError {
111///                 code: "invalid_credential".to_string(),
112///                 message: "Credential must start with sk-".to_string(),
113///             })
114///         }
115///     }
116/// }
117/// ```
118pub trait Validate: Clone + Send + Sync + 'static {
119    /// Validate a credential and return the authentication result.
120    fn validate(&self, credential: &str) -> impl Future<Output = AuthResult> + Send;
121}
122
123// =============================================================================
124// API Key Authentication
125// =============================================================================
126
127/// Simple in-memory API key validator
128///
129/// For production use, consider:
130/// - Database-backed validation
131/// - Caching with TTL
132/// - Rate limiting per key
133#[derive(Debug, Clone)]
134pub struct ApiKeyValidator {
135    valid_keys: Arc<HashSet<String>>,
136}
137
138impl ApiKeyValidator {
139    /// Create a new validator with a list of valid API keys
140    pub fn new(keys: impl IntoIterator<Item = String>) -> Self {
141        Self {
142            valid_keys: Arc::new(keys.into_iter().collect()),
143        }
144    }
145
146    /// Add a key to the valid set
147    pub fn add_key(&mut self, key: String) {
148        Arc::make_mut(&mut self.valid_keys).insert(key);
149    }
150
151    /// Check if a key is valid
152    pub fn is_valid(&self, key: &str) -> bool {
153        self.valid_keys.contains(key)
154    }
155}
156
157impl Validate for ApiKeyValidator {
158    async fn validate(&self, key: &str) -> AuthResult {
159        if self.valid_keys.contains(key) {
160            AuthResult::Authenticated(Some(AuthInfo {
161                client_id: format!("api_key:{}", &key[..8.min(key.len())]),
162                claims: None,
163            }))
164        } else {
165            AuthResult::Failed(AuthError {
166                code: "invalid_api_key".to_string(),
167                message: "The provided API key is not valid".to_string(),
168            })
169        }
170    }
171}
172
173// =============================================================================
174// Bearer Token Authentication
175// =============================================================================
176
177/// Simple bearer token validator that checks against a static set of tokens.
178///
179/// For production, implement [`Validate`] with:
180/// - JWT verification using a signing key
181/// - OAuth2 token introspection
182/// - OIDC ID token validation
183#[derive(Debug, Clone)]
184pub struct StaticBearerValidator {
185    valid_tokens: Arc<HashSet<String>>,
186}
187
188impl StaticBearerValidator {
189    /// Create a new validator with a list of valid tokens
190    pub fn new(tokens: impl IntoIterator<Item = String>) -> Self {
191        Self {
192            valid_tokens: Arc::new(tokens.into_iter().collect()),
193        }
194    }
195}
196
197impl Validate for StaticBearerValidator {
198    async fn validate(&self, token: &str) -> AuthResult {
199        if self.valid_tokens.contains(token) {
200            AuthResult::Authenticated(Some(AuthInfo {
201                client_id: format!("bearer:{}", &token[..8.min(token.len())]),
202                claims: None,
203            }))
204        } else {
205            AuthResult::Failed(AuthError {
206                code: "invalid_token".to_string(),
207                message: "The provided bearer token is not valid".to_string(),
208            })
209        }
210    }
211}
212
213// =============================================================================
214// Authorization Header Parsing
215// =============================================================================
216
217/// Extract an API key from an Authorization header
218///
219/// Supports formats:
220/// - `Bearer <key>` (standard)
221/// - `ApiKey <key>`
222/// - `<key>` (raw key)
223pub fn extract_api_key(auth_header: &str) -> Option<&str> {
224    let auth_header = auth_header.trim();
225
226    if let Some(key) = auth_header.strip_prefix("Bearer ") {
227        Some(key.trim())
228    } else if let Some(key) = auth_header.strip_prefix("ApiKey ") {
229        Some(key.trim())
230    } else if !auth_header.contains(' ') {
231        // Raw key without prefix
232        Some(auth_header)
233    } else {
234        None
235    }
236}
237
238/// Extract a bearer token from an Authorization header
239pub fn extract_bearer_token(auth_header: &str) -> Option<&str> {
240    auth_header.trim().strip_prefix("Bearer ").map(|t| t.trim())
241}
242
243// =============================================================================
244// Generic Auth Layer
245// =============================================================================
246
247/// A Tower layer that performs authentication using a provided validator
248///
249/// This is a generic auth layer that can be used with any validator that
250/// implements the appropriate validation trait.
251#[derive(Clone)]
252pub struct AuthLayer<V> {
253    validator: V,
254    header_name: String,
255}
256
257impl<V> AuthLayer<V> {
258    /// Create a new auth layer with the given validator
259    ///
260    /// By default, looks for the `Authorization` header
261    pub fn new(validator: V) -> Self {
262        Self {
263            validator,
264            header_name: "Authorization".to_string(),
265        }
266    }
267
268    /// Use a custom header name for the auth token
269    pub fn header_name(mut self, name: impl Into<String>) -> Self {
270        self.header_name = name.into();
271        self
272    }
273}
274
275impl<S, V: Clone> Layer<S> for AuthLayer<V> {
276    type Service = AuthService<S, V>;
277
278    fn layer(&self, inner: S) -> Self::Service {
279        AuthService {
280            inner,
281            validator: self.validator.clone(),
282            header_name: self.header_name.clone(),
283        }
284    }
285}
286
287/// Tower service that performs authentication on incoming requests.
288///
289/// Created by [`AuthLayer`]. Extracts credentials from the configured HTTP
290/// header, validates them using the provided [`Validate`] implementation,
291/// and either forwards the request (injecting [`AuthInfo`] into request
292/// extensions) or returns an HTTP 401 response.
293///
294/// # Example
295///
296/// ```rust,ignore
297/// // Requires the `http` feature
298/// use tower::ServiceBuilder;
299/// use tower_mcp::auth::{AuthLayer, ApiKeyValidator};
300///
301/// let validator = ApiKeyValidator::new(vec!["sk-test-key-123".to_string()]);
302///
303/// let service = ServiceBuilder::new()
304///     .layer(AuthLayer::new(validator))
305///     .service(inner_service);
306/// ```
307#[derive(Clone)]
308#[cfg_attr(not(feature = "http"), allow(dead_code))]
309pub struct AuthService<S, V> {
310    inner: S,
311    validator: V,
312    header_name: String,
313}
314
315#[cfg(feature = "http")]
316impl<S, V> tower_service::Service<axum::http::Request<axum::body::Body>> for AuthService<S, V>
317where
318    S: tower_service::Service<
319            axum::http::Request<axum::body::Body>,
320            Response = axum::response::Response,
321        > + Clone
322        + Send
323        + 'static,
324    S::Future: Send,
325    S::Error: Into<crate::BoxError> + Send,
326    V: Validate,
327{
328    type Response = axum::response::Response;
329    type Error = S::Error;
330    type Future =
331        std::pin::Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
332
333    fn poll_ready(
334        &mut self,
335        cx: &mut std::task::Context<'_>,
336    ) -> std::task::Poll<Result<(), Self::Error>> {
337        self.inner.poll_ready(cx)
338    }
339
340    fn call(&mut self, req: axum::http::Request<axum::body::Body>) -> Self::Future {
341        let credential = req
342            .headers()
343            .get(&self.header_name)
344            .and_then(|v| v.to_str().ok())
345            .and_then(extract_api_key)
346            .map(|s| s.to_owned());
347
348        let inner = self.inner.clone();
349        let validator = self.validator.clone();
350
351        Box::pin(async move {
352            let Some(credential) = credential else {
353                return Ok(unauthorized_response(
354                    "Missing authentication credentials. Provide via Authorization header.",
355                ));
356            };
357
358            match validator.validate(&credential).await {
359                AuthResult::Authenticated(info) => {
360                    let mut req = req;
361                    if let Some(info) = info {
362                        req.extensions_mut().insert(info);
363                    }
364                    inner.oneshot(req).await
365                }
366                AuthResult::Failed(err) => Ok(unauthorized_response(&err.message)),
367            }
368        })
369    }
370}
371
372/// Construct an HTTP 401 Unauthorized response with a JSON-RPC error body.
373///
374/// Uses the MCP `Forbidden` code (-32007). The previous code (-32001) was
375/// reclaimed by SEP-2243 for `HeaderMismatch`; emitting that here would
376/// confuse clients that route on the JSON-RPC error code.
377#[cfg(feature = "http")]
378fn unauthorized_response(message: &str) -> axum::response::Response {
379    use axum::http::StatusCode;
380    use axum::response::IntoResponse;
381
382    let body = serde_json::json!({
383        "jsonrpc": "2.0",
384        "error": {
385            "code": tower_mcp_types::McpErrorCode::Forbidden.code(),
386            "message": message
387        },
388        "id": null
389    });
390
391    (StatusCode::UNAUTHORIZED, axum::Json(body)).into_response()
392}
393
394// =============================================================================
395// Helper for building auth middleware
396// =============================================================================
397
398/// Builder for creating auth middleware configurations
399#[derive(Clone)]
400pub struct AuthConfig {
401    /// Whether to allow unauthenticated requests to pass through
402    pub allow_anonymous: bool,
403    /// Paths that don't require authentication
404    pub public_paths: Vec<String>,
405    /// Custom header name for auth token
406    pub header_name: String,
407}
408
409impl Default for AuthConfig {
410    fn default() -> Self {
411        Self {
412            allow_anonymous: false,
413            public_paths: Vec::new(),
414            header_name: "Authorization".to_string(),
415        }
416    }
417}
418
419impl AuthConfig {
420    /// Create a new auth config
421    pub fn new() -> Self {
422        Self::default()
423    }
424
425    /// Allow anonymous requests (no auth required)
426    pub fn allow_anonymous(mut self, allow: bool) -> Self {
427        self.allow_anonymous = allow;
428        self
429    }
430
431    /// Add paths that don't require authentication
432    pub fn public_path(mut self, path: impl Into<String>) -> Self {
433        self.public_paths.push(path.into());
434        self
435    }
436
437    /// Set the header name for auth tokens
438    pub fn header_name(mut self, name: impl Into<String>) -> Self {
439        self.header_name = name.into();
440        self
441    }
442
443    /// Check if a path is public (doesn't require auth)
444    pub fn is_public(&self, path: &str) -> bool {
445        self.public_paths.iter().any(|p| path.starts_with(p))
446    }
447}
448
449#[cfg(test)]
450mod tests {
451    use super::*;
452
453    #[test]
454    fn test_extract_api_key_bearer() {
455        assert_eq!(extract_api_key("Bearer sk-123"), Some("sk-123"));
456        assert_eq!(extract_api_key("Bearer  sk-123 "), Some("sk-123"));
457    }
458
459    #[test]
460    fn test_extract_api_key_apikey_prefix() {
461        assert_eq!(extract_api_key("ApiKey sk-123"), Some("sk-123"));
462    }
463
464    #[test]
465    fn test_extract_api_key_raw() {
466        assert_eq!(extract_api_key("sk-123"), Some("sk-123"));
467    }
468
469    #[test]
470    fn test_extract_api_key_invalid() {
471        assert_eq!(extract_api_key("Basic user:pass"), None);
472    }
473
474    #[test]
475    fn test_extract_bearer_token() {
476        assert_eq!(extract_bearer_token("Bearer abc123"), Some("abc123"));
477        assert_eq!(extract_bearer_token("bearer abc123"), None); // case sensitive
478        assert_eq!(extract_bearer_token("abc123"), None);
479    }
480
481    #[tokio::test]
482    async fn test_api_key_validator() {
483        let validator = ApiKeyValidator::new(vec!["valid-key".to_string()]);
484
485        match validator.validate("valid-key").await {
486            AuthResult::Authenticated(info) => {
487                assert!(info.is_some());
488            }
489            AuthResult::Failed(_) => panic!("Expected authentication to succeed"),
490        }
491
492        match validator.validate("invalid-key").await {
493            AuthResult::Authenticated(_) => panic!("Expected authentication to fail"),
494            AuthResult::Failed(err) => {
495                assert_eq!(err.code, "invalid_api_key");
496            }
497        }
498    }
499
500    #[tokio::test]
501    async fn test_bearer_validator() {
502        let validator = StaticBearerValidator::new(vec!["token123".to_string()]);
503
504        match validator.validate("token123").await {
505            AuthResult::Authenticated(info) => {
506                assert!(info.is_some());
507            }
508            AuthResult::Failed(_) => panic!("Expected authentication to succeed"),
509        }
510
511        match validator.validate("bad-token").await {
512            AuthResult::Authenticated(_) => panic!("Expected authentication to fail"),
513            AuthResult::Failed(err) => {
514                assert_eq!(err.code, "invalid_token");
515            }
516        }
517    }
518
519    #[test]
520    fn test_auth_config() {
521        let config = AuthConfig::new()
522            .allow_anonymous(false)
523            .public_path("/health")
524            .public_path("/metrics")
525            .header_name("X-API-Key");
526
527        assert!(!config.allow_anonymous);
528        assert!(config.is_public("/health"));
529        assert!(config.is_public("/metrics/cpu"));
530        assert!(!config.is_public("/api/tools"));
531        assert_eq!(config.header_name, "X-API-Key");
532    }
533
534    #[test]
535    fn test_auth_layer_creates_service() {
536        let validator = ApiKeyValidator::new(vec!["key".to_string()]);
537        let layer = AuthLayer::new(validator);
538        // Wrap a no-op service to verify the Layer impl works
539        let _service: AuthService<(), ApiKeyValidator> = layer.layer(());
540    }
541
542    #[cfg(feature = "http")]
543    mod http_tests {
544        use super::*;
545        use std::pin::Pin;
546        use std::task::{Context, Poll};
547
548        use axum::body::Body;
549        use axum::http::{Request, StatusCode};
550        use tower::ServiceExt;
551        use tower_service::Service;
552
553        /// A minimal inner service that returns 200 OK for any request
554        #[derive(Clone)]
555        struct OkService;
556
557        impl Service<Request<Body>> for OkService {
558            type Response = axum::response::Response;
559            type Error = std::convert::Infallible;
560            type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
561
562            fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
563                Poll::Ready(Ok(()))
564            }
565
566            fn call(&mut self, _req: Request<Body>) -> Self::Future {
567                Box::pin(async {
568                    Ok(axum::response::Response::builder()
569                        .status(StatusCode::OK)
570                        .body(Body::empty())
571                        .unwrap())
572                })
573            }
574        }
575
576        #[tokio::test]
577        async fn test_auth_service_rejects_missing_credentials() {
578            let validator = ApiKeyValidator::new(vec!["sk-test-123".to_string()]);
579            let layer = AuthLayer::new(validator);
580            let mut service = layer.layer(OkService);
581
582            let req = Request::builder().uri("/").body(Body::empty()).unwrap();
583
584            let resp = service.ready().await.unwrap().call(req).await.unwrap();
585            assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
586        }
587
588        #[tokio::test]
589        async fn test_auth_service_rejects_invalid_key() {
590            let validator = ApiKeyValidator::new(vec!["sk-test-123".to_string()]);
591            let layer = AuthLayer::new(validator);
592            let mut service = layer.layer(OkService);
593
594            let req = Request::builder()
595                .uri("/")
596                .header("Authorization", "Bearer sk-wrong-key")
597                .body(Body::empty())
598                .unwrap();
599
600            let resp = service.ready().await.unwrap().call(req).await.unwrap();
601            assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
602        }
603
604        #[tokio::test]
605        async fn test_auth_service_accepts_valid_key() {
606            let validator = ApiKeyValidator::new(vec!["sk-test-123".to_string()]);
607            let layer = AuthLayer::new(validator);
608            let mut service = layer.layer(OkService);
609
610            let req = Request::builder()
611                .uri("/")
612                .header("Authorization", "Bearer sk-test-123")
613                .body(Body::empty())
614                .unwrap();
615
616            let resp = service.ready().await.unwrap().call(req).await.unwrap();
617            assert_eq!(resp.status(), StatusCode::OK);
618        }
619
620        #[tokio::test]
621        async fn test_auth_service_injects_auth_info() {
622            let validator = ApiKeyValidator::new(vec!["sk-test-123".to_string()]);
623            let layer = AuthLayer::new(validator);
624
625            // Inner service that checks for AuthInfo in extensions
626            #[derive(Clone)]
627            struct CheckAuthInfo;
628
629            impl Service<Request<Body>> for CheckAuthInfo {
630                type Response = axum::response::Response;
631                type Error = std::convert::Infallible;
632                type Future =
633                    Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
634
635                fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
636                    Poll::Ready(Ok(()))
637                }
638
639                fn call(&mut self, req: Request<Body>) -> Self::Future {
640                    let has_auth = req.extensions().get::<AuthInfo>().is_some();
641                    Box::pin(async move {
642                        let status = if has_auth {
643                            StatusCode::OK
644                        } else {
645                            StatusCode::INTERNAL_SERVER_ERROR
646                        };
647                        Ok(axum::response::Response::builder()
648                            .status(status)
649                            .body(Body::empty())
650                            .unwrap())
651                    })
652                }
653            }
654
655            let mut service = layer.layer(CheckAuthInfo);
656
657            let req = Request::builder()
658                .uri("/")
659                .header("Authorization", "Bearer sk-test-123")
660                .body(Body::empty())
661                .unwrap();
662
663            let resp = service.ready().await.unwrap().call(req).await.unwrap();
664            assert_eq!(resp.status(), StatusCode::OK);
665        }
666
667        #[tokio::test]
668        async fn test_auth_service_custom_header() {
669            let validator = ApiKeyValidator::new(vec!["my-key".to_string()]);
670            let layer = AuthLayer::new(validator).header_name("X-API-Key");
671            let mut service = layer.layer(OkService);
672
673            // Standard Authorization header should not work
674            let req = Request::builder()
675                .uri("/")
676                .header("Authorization", "Bearer my-key")
677                .body(Body::empty())
678                .unwrap();
679            let resp = service.ready().await.unwrap().call(req).await.unwrap();
680            assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
681
682            // Custom header should work
683            let req = Request::builder()
684                .uri("/")
685                .header("X-API-Key", "my-key")
686                .body(Body::empty())
687                .unwrap();
688            let resp = service.ready().await.unwrap().call(req).await.unwrap();
689            assert_eq!(resp.status(), StatusCode::OK);
690        }
691    }
692}
tower_mcp/auth.rs

tower_mcp/
auth.rs
