Skip to main content

turul_mcp_aws_lambda/
cors.rs

1//! CORS (Cross-Origin Resource Sharing) support for Lambda MCP servers
2//!
3//! This module provides CORS header injection for Lambda responses, since Tower
4//! middleware cannot be used in the Lambda execution environment.
5
6use std::collections::HashSet;
7
8use http::{HeaderValue, Method};
9use lambda_http::{Body as LambdaBody, Response as LambdaResponse};
10use tracing::debug;
11
12use crate::error::{LambdaError, Result};
13
14/// CORS configuration for Lambda MCP servers
15#[derive(Debug, Clone)]
16pub struct CorsConfig {
17    /// Allowed origins for CORS requests
18    /// Use "*" to allow all origins (not recommended for production)
19    pub allowed_origins: Vec<String>,
20
21    /// Allowed HTTP methods
22    pub allowed_methods: Vec<Method>,
23
24    /// Allowed request headers
25    pub allowed_headers: Vec<String>,
26
27    /// Whether to allow credentials (cookies, authorization headers)
28    pub allow_credentials: bool,
29
30    /// Maximum age for preflight cache (in seconds)
31    pub max_age: Option<u32>,
32
33    /// Headers to expose to the client
34    pub expose_headers: Vec<String>,
35}
36
37impl Default for CorsConfig {
38    fn default() -> Self {
39        Self {
40            allowed_origins: vec!["*".to_string()],
41            allowed_methods: vec![Method::GET, Method::POST, Method::DELETE, Method::OPTIONS],
42            allowed_headers: vec![
43                "Content-Type".to_string(),
44                "Accept".to_string(),
45                "Authorization".to_string(),
46                "Mcp-Session-Id".to_string(),
47                "Mcp-Protocol-Version".to_string(),
48                "Last-Event-ID".to_string(),
49            ],
50            allow_credentials: false,
51            max_age: Some(86400), // 24 hours
52            expose_headers: vec![
53                "Mcp-Session-Id".to_string(),
54                "Mcp-Protocol-Version".to_string(),
55                // Exposed so browser OAuth clients can read the RFC 9728
56                // challenge on 401 responses (non-safelisted CORS header).
57                "WWW-Authenticate".to_string(),
58            ],
59        }
60    }
61}
62
63impl CorsConfig {
64    /// Create a CORS config that allows all origins (for development)
65    pub fn allow_all() -> Self {
66        Self::default()
67    }
68
69    /// Create a CORS config for specific origins
70    pub fn for_origins(origins: Vec<String>) -> Self {
71        Self {
72            allowed_origins: origins,
73            ..Default::default()
74        }
75    }
76
77    /// Create a CORS config from environment variables
78    pub fn from_env() -> Self {
79        let allowed_origins = std::env::var("MCP_CORS_ORIGINS")
80            .map(|s| s.split(',').map(|s| s.trim().to_string()).collect())
81            .unwrap_or_else(|_| vec!["*".to_string()]);
82
83        let allow_credentials = std::env::var("MCP_CORS_CREDENTIALS")
84            .map(|s| s.parse().unwrap_or(false))
85            .unwrap_or(false);
86
87        let max_age = std::env::var("MCP_CORS_MAX_AGE")
88            .ok()
89            .and_then(|s| s.parse().ok());
90
91        Self {
92            allowed_origins,
93            allow_credentials,
94            max_age,
95            ..Default::default()
96        }
97    }
98}
99
100/// Inject CORS headers into a Lambda response (generic over body type)
101///
102/// This function adds the appropriate CORS headers based on the configuration
103/// and the incoming request's Origin header.
104pub fn inject_cors_headers<B>(
105    response: &mut lambda_http::Response<B>,
106    config: &CorsConfig,
107    request_origin: Option<&str>,
108) -> Result<()> {
109    debug!("Injecting CORS headers for origin: {:?}", request_origin);
110
111    // Determine allowed origin
112    let allowed_origin = determine_allowed_origin(config, request_origin);
113
114    if let Some(origin) = allowed_origin {
115        response.headers_mut().insert(
116            "Access-Control-Allow-Origin",
117            HeaderValue::from_str(&origin)
118                .map_err(|e| LambdaError::Cors(format!("Invalid origin: {}", e)))?,
119        );
120    }
121
122    // Add allowed methods
123    let methods_str = config
124        .allowed_methods
125        .iter()
126        .map(|m| m.as_str())
127        .collect::<Vec<_>>()
128        .join(", ");
129    response.headers_mut().insert(
130        "Access-Control-Allow-Methods",
131        HeaderValue::from_str(&methods_str)
132            .map_err(|e| LambdaError::Cors(format!("Invalid methods: {}", e)))?,
133    );
134
135    // Add allowed headers
136    if !config.allowed_headers.is_empty() {
137        let headers_str = config.allowed_headers.join(", ");
138        response.headers_mut().insert(
139            "Access-Control-Allow-Headers",
140            HeaderValue::from_str(&headers_str)
141                .map_err(|e| LambdaError::Cors(format!("Invalid headers: {}", e)))?,
142        );
143    }
144
145    // Add exposed headers
146    if !config.expose_headers.is_empty() {
147        let expose_str = config.expose_headers.join(", ");
148        response.headers_mut().insert(
149            "Access-Control-Expose-Headers",
150            HeaderValue::from_str(&expose_str)
151                .map_err(|e| LambdaError::Cors(format!("Invalid expose headers: {}", e)))?,
152        );
153    }
154
155    // Add credentials if allowed
156    if config.allow_credentials {
157        response.headers_mut().insert(
158            "Access-Control-Allow-Credentials",
159            HeaderValue::from_static("true"),
160        );
161    }
162
163    // Add max age for preflight requests
164    if let Some(max_age) = config.max_age {
165        response.headers_mut().insert(
166            "Access-Control-Max-Age",
167            HeaderValue::from_str(&max_age.to_string())
168                .map_err(|e| LambdaError::Cors(format!("Invalid max age: {}", e)))?,
169        );
170    }
171
172    debug!("CORS headers injected successfully");
173    Ok(())
174}
175
176/// Create a CORS preflight response
177///
178/// Handles OPTIONS requests that browsers send before making actual CORS requests.
179pub fn create_preflight_response(
180    config: &CorsConfig,
181    request_origin: Option<&str>,
182) -> Result<LambdaResponse<LambdaBody>> {
183    debug!("Creating CORS preflight response");
184
185    let mut response = LambdaResponse::builder()
186        .status(200)
187        .body(LambdaBody::Empty)
188        .map_err(LambdaError::Http)?;
189
190    inject_cors_headers(&mut response, config, request_origin)?;
191
192    Ok(response)
193}
194
195/// Determine the allowed origin based on configuration and request
196fn determine_allowed_origin(config: &CorsConfig, request_origin: Option<&str>) -> Option<String> {
197    // If wildcard is configured, return it
198    if config.allowed_origins.contains(&"*".to_string()) {
199        return Some("*".to_string());
200    }
201
202    // If no origin in request, no CORS header needed
203    let request_origin = request_origin?;
204
205    // Check if the request origin is in the allowed list
206    if config.allowed_origins.contains(&request_origin.to_string()) {
207        Some(request_origin.to_string())
208    } else {
209        // Origin not allowed, don't set CORS header
210        None
211    }
212}
213
214/// Validate CORS configuration
215pub fn validate_config(config: &CorsConfig) -> Result<()> {
216    // Check for wildcard with credentials (security issue)
217    if config.allow_credentials && config.allowed_origins.contains(&"*".to_string()) {
218        return Err(LambdaError::Cors(
219            "Cannot use wildcard origin (*) with credentials enabled".to_string(),
220        ));
221    }
222
223    // Validate origins are proper URLs or wildcards
224    for origin in &config.allowed_origins {
225        if origin != "*" && !origin.starts_with("http://") && !origin.starts_with("https://") {
226            return Err(LambdaError::Cors(format!(
227                "Invalid origin format: {}",
228                origin
229            )));
230        }
231    }
232
233    // Check for duplicate headers
234    let headers_set: HashSet<_> = config.allowed_headers.iter().collect();
235    if headers_set.len() != config.allowed_headers.len() {
236        return Err(LambdaError::Cors(
237            "Duplicate headers in allowed_headers".to_string(),
238        ));
239    }
240
241    Ok(())
242}
243
244#[cfg(test)]
245mod tests {
246    use super::*;
247    use lambda_http::Body;
248
249    #[test]
250    fn test_default_config() {
251        let config = CorsConfig::default();
252        assert!(config.allowed_origins.contains(&"*".to_string()));
253        assert!(config.allowed_methods.contains(&Method::GET));
254        assert!(config.allowed_methods.contains(&Method::POST));
255        assert!(config.allowed_headers.contains(&"Content-Type".to_string()));
256    }
257
258    #[test]
259    fn test_config_validation() {
260        let mut config = CorsConfig::default();
261        assert!(validate_config(&config).is_ok());
262
263        // Test invalid wildcard with credentials
264        config.allow_credentials = true;
265        assert!(validate_config(&config).is_err());
266
267        // Test invalid origin format
268        config.allow_credentials = false;
269        config.allowed_origins = vec!["invalid-origin".to_string()];
270        assert!(validate_config(&config).is_err());
271    }
272
273    #[tokio::test]
274    async fn test_cors_headers_injection() {
275        let config = CorsConfig::default();
276        let mut response = LambdaResponse::builder()
277            .status(200)
278            .body(Body::Empty)
279            .unwrap();
280
281        inject_cors_headers(&mut response, &config, Some("https://example.com")).unwrap();
282
283        assert_eq!(
284            response.headers().get("access-control-allow-origin"),
285            Some(&HeaderValue::from_static("*"))
286        );
287
288        assert!(
289            response
290                .headers()
291                .contains_key("access-control-allow-methods")
292        );
293        assert!(
294            response
295                .headers()
296                .contains_key("access-control-allow-headers")
297        );
298    }
299
300    #[test]
301    fn test_default_expose_headers_contains_www_authenticate() {
302        // Browsers can only read non-safelisted response headers when listed
303        // in Access-Control-Expose-Headers. RFC 9728 OAuth discovery requires
304        // clients to parse the WWW-Authenticate challenge on 401 responses,
305        // so it must be exposed by default for any browser-fronted MCP server.
306        let config = CorsConfig::default();
307        assert!(
308            config
309                .expose_headers
310                .iter()
311                .any(|h| h.eq_ignore_ascii_case("WWW-Authenticate")),
312            "default expose_headers must include WWW-Authenticate; got {:?}",
313            config.expose_headers,
314        );
315    }
316
317    #[tokio::test]
318    async fn test_custom_expose_headers_not_mutated_by_injection() {
319        // Caller-supplied expose_headers wins as-is. Injection must not
320        // augment or rewrite the list — consumers control the surface.
321        let config = CorsConfig {
322            expose_headers: vec!["X-Custom".to_string()],
323            ..Default::default()
324        };
325        let original = config.expose_headers.clone();
326
327        let mut response = LambdaResponse::builder()
328            .status(200)
329            .body(Body::Empty)
330            .unwrap();
331        inject_cors_headers(&mut response, &config, Some("https://example.com")).unwrap();
332
333        assert_eq!(config.expose_headers, original);
334        assert_eq!(
335            response.headers().get("access-control-expose-headers"),
336            Some(&HeaderValue::from_static("X-Custom"))
337        );
338    }
339
340    #[tokio::test]
341    async fn test_preflight_response() {
342        let config = CorsConfig::default();
343        let response = create_preflight_response(&config, Some("https://example.com")).unwrap();
344
345        assert_eq!(response.status(), 200);
346        assert!(
347            response
348                .headers()
349                .contains_key("access-control-allow-origin")
350        );
351        assert!(
352            response
353                .headers()
354                .contains_key("access-control-allow-methods")
355        );
356    }
357}