mcp_proxy/
token.rs

1//! Token passthrough middleware for forwarding client credentials to backends.
2//!
3//! When a backend has `forward_auth = true`, the client's inbound bearer token
4//! is extracted from `RouterRequest.extensions` and stored as a [`ClientToken`]
5//! for downstream middleware and backend services to consume.
6//!
7//! # Configuration
8//!
9//! ```toml
10//! [[backends]]
11//! name = "github"
12//! transport = "http"
13//! url = "http://github-mcp.internal:8080"
14//! forward_auth = true  # forward client's token to this backend
15//!
16//! [[backends]]
17//! name = "db"
18//! transport = "http"
19//! url = "http://db-mcp.internal:8080"
20//! bearer_token = "${DB_API_KEY}"  # static token for this backend
21//! ```
22//!
23//! # How it works
24//!
25//! 1. The proxy's auth layer (JWT/bearer) validates the inbound token and
26//!    stores [`TokenClaims`](tower_mcp::oauth::token::TokenClaims) in request extensions.
27//! 2. This middleware reads the `TokenClaims` and stores the subject (`sub` claim)
28//!    and any available identity info as a [`ClientToken`] in extensions.
29//! 3. Backend-specific middleware or future transport enhancements can read
30//!    `ClientToken` to forward credentials.
31
32use std::collections::HashSet;
33use std::convert::Infallible;
34use std::future::Future;
35use std::pin::Pin;
36use std::sync::Arc;
37use std::task::{Context, Poll};
38
39use tower::Service;
40
41use tower_mcp::router::{RouterRequest, RouterResponse};
42use tower_mcp_types::protocol::McpRequest;
43
44/// A client's identity token extracted from inbound authentication.
45///
46/// Stored in `RouterRequest.extensions` by the [`TokenPassthroughService`]
47/// for downstream consumption.
48#[derive(Debug, Clone)]
49pub struct ClientToken {
50    /// The subject (user/client identifier) from the token.
51    pub subject: Option<String>,
52    /// Space-delimited scopes from the token.
53    pub scope: Option<String>,
54    /// The raw bearer token string, if available.
55    pub raw_token: Option<String>,
56}
57
58/// Middleware that extracts client identity from auth claims and makes it
59/// available to backends configured with `forward_auth = true`.
60#[derive(Clone)]
61pub struct TokenPassthroughService<S> {
62    inner: S,
63    forward_namespaces: Arc<HashSet<String>>,
64}
65
66impl<S> TokenPassthroughService<S> {
67    /// Create a new token passthrough service.
68    ///
69    /// `forward_namespaces` is the set of backend namespace prefixes (e.g. `"github/"`)
70    /// that should receive forwarded tokens.
71    pub fn new(inner: S, forward_namespaces: HashSet<String>) -> Self {
72        Self {
73            inner,
74            forward_namespaces: Arc::new(forward_namespaces),
75        }
76    }
77}
78
79/// Check if a request targets a namespace that wants token forwarding.
80fn request_targets_namespace(req: &McpRequest, namespaces: &HashSet<String>) -> bool {
81    let name = match req {
82        McpRequest::CallTool(params) => Some(params.name.as_str()),
83        McpRequest::ReadResource(params) => Some(params.uri.as_str()),
84        McpRequest::GetPrompt(params) => Some(params.name.as_str()),
85        _ => None,
86    };
87    if let Some(name) = name {
88        namespaces.iter().any(|ns| name.starts_with(ns))
89    } else {
90        false
91    }
92}
93
94impl<S> Service<RouterRequest> for TokenPassthroughService<S>
95where
96    S: Service<RouterRequest, Response = RouterResponse, Error = Infallible>
97        + Clone
98        + Send
99        + 'static,
100    S::Future: Send,
101{
102    type Response = RouterResponse;
103    type Error = Infallible;
104    type Future = Pin<Box<dyn Future<Output = Result<RouterResponse, Infallible>> + Send>>;
105
106    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
107        self.inner.poll_ready(cx)
108    }
109
110    fn call(&mut self, mut req: RouterRequest) -> Self::Future {
111        // Only inject ClientToken for requests targeting forward_auth backends
112        if !self.forward_namespaces.is_empty()
113            && request_targets_namespace(&req.inner, &self.forward_namespaces)
114        {
115            let client_token = req
116                .extensions
117                .get::<tower_mcp::oauth::token::TokenClaims>()
118                .map(|claims| ClientToken {
119                    subject: claims.sub.clone(),
120                    scope: claims.scope.clone(),
121                    raw_token: None, // Raw token not available from TokenClaims
122                });
123            if let Some(token) = client_token {
124                tracing::debug!(
125                    subject = ?token.subject,
126                    "Injected ClientToken for forward_auth backend"
127                );
128                req.extensions.insert(token);
129            }
130        }
131
132        let fut = self.inner.call(req);
133        Box::pin(fut)
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use std::collections::HashSet;
140
141    use tower::Service;
142    use tower_mcp::protocol::{CallToolParams, McpRequest, RequestId};
143    use tower_mcp::router::{Extensions, RouterRequest};
144
145    use super::{TokenPassthroughService, request_targets_namespace};
146    use crate::test_util::{MockService, call_service};
147
148    #[test]
149    fn test_request_targets_namespace_match() {
150        let namespaces: HashSet<String> = ["github/".to_string()].into();
151        let req = McpRequest::CallTool(CallToolParams {
152            name: "github/search".to_string(),
153            arguments: serde_json::json!({}),
154            input_responses: None,
155            request_state: None,
156            meta: None,
157            task: None,
158        });
159        assert!(request_targets_namespace(&req, &namespaces));
160    }
161
162    #[test]
163    fn test_request_targets_namespace_no_match() {
164        let namespaces: HashSet<String> = ["github/".to_string()].into();
165        let req = McpRequest::CallTool(CallToolParams {
166            name: "db/query".to_string(),
167            arguments: serde_json::json!({}),
168            input_responses: None,
169            request_state: None,
170            meta: None,
171            task: None,
172        });
173        assert!(!request_targets_namespace(&req, &namespaces));
174    }
175
176    #[test]
177    fn test_request_targets_namespace_list_tools() {
178        let namespaces: HashSet<String> = ["github/".to_string()].into();
179        let req = McpRequest::ListTools(Default::default());
180        assert!(!request_targets_namespace(&req, &namespaces));
181    }
182
183    #[tokio::test]
184    async fn test_passthrough_injects_client_token() {
185        let mock = MockService::with_tools(&["github/search"]);
186        let namespaces: HashSet<String> = ["github/".to_string()].into();
187        let mut svc = TokenPassthroughService::new(mock, namespaces);
188
189        // Create request with TokenClaims in extensions
190        let mut extensions = Extensions::new();
191        extensions.insert(tower_mcp::oauth::token::TokenClaims {
192            sub: Some("user-123".to_string()),
193            scope: Some("mcp:read".to_string()),
194            iss: None,
195            aud: None,
196            exp: None,
197            client_id: None,
198            extra: Default::default(),
199        });
200
201        let req = RouterRequest {
202            id: RequestId::Number(1),
203            inner: McpRequest::CallTool(CallToolParams {
204                name: "github/search".to_string(),
205                arguments: serde_json::json!({}),
206                input_responses: None,
207                request_state: None,
208                meta: None,
209                task: None,
210            }),
211            extensions,
212        };
213
214        let resp = svc.call(req).await.unwrap();
215        assert!(resp.inner.is_ok());
216    }
217
218    #[tokio::test]
219    async fn test_passthrough_skips_non_forward_backends() {
220        let mock = MockService::with_tools(&["db/query"]);
221        let namespaces: HashSet<String> = ["github/".to_string()].into();
222        let mut svc = TokenPassthroughService::new(mock, namespaces);
223
224        let resp = call_service(
225            &mut svc,
226            McpRequest::CallTool(CallToolParams {
227                name: "db/query".to_string(),
228                arguments: serde_json::json!({}),
229                input_responses: None,
230                request_state: None,
231                meta: None,
232                task: None,
233            }),
234        )
235        .await;
236
237        assert!(resp.inner.is_ok());
238    }
239
240    #[tokio::test]
241    async fn test_passthrough_no_claims_passes_through() {
242        let mock = MockService::with_tools(&["github/search"]);
243        let namespaces: HashSet<String> = ["github/".to_string()].into();
244        let mut svc = TokenPassthroughService::new(mock, namespaces);
245
246        // No TokenClaims in extensions -- should still pass through
247        let resp = call_service(
248            &mut svc,
249            McpRequest::CallTool(CallToolParams {
250                name: "github/search".to_string(),
251                arguments: serde_json::json!({}),
252                input_responses: None,
253                request_state: None,
254                meta: None,
255                task: None,
256            }),
257        )
258        .await;
259
260        assert!(resp.inner.is_ok());
261    }
262}
mcp_proxy/token.rs

mcp_proxy/
token.rs
