turul_mcp_aws_lambda/
cors.rs1use 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#[derive(Debug, Clone)]
16pub struct CorsConfig {
17 pub allowed_origins: Vec<String>,
20
21 pub allowed_methods: Vec<Method>,
23
24 pub allowed_headers: Vec<String>,
26
27 pub allow_credentials: bool,
29
30 pub max_age: Option<u32>,
32
33 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), expose_headers: vec![
53 "Mcp-Session-Id".to_string(),
54 "Mcp-Protocol-Version".to_string(),
55 "WWW-Authenticate".to_string(),
58 ],
59 }
60 }
61}
62
63impl CorsConfig {
64 pub fn allow_all() -> Self {
66 Self::default()
67 }
68
69 pub fn for_origins(origins: Vec<String>) -> Self {
71 Self {
72 allowed_origins: origins,
73 ..Default::default()
74 }
75 }
76
77 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
100pub 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 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 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 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 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 if config.allow_credentials {
157 response.headers_mut().insert(
158 "Access-Control-Allow-Credentials",
159 HeaderValue::from_static("true"),
160 );
161 }
162
163 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
176pub 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
195fn determine_allowed_origin(config: &CorsConfig, request_origin: Option<&str>) -> Option<String> {
197 if config.allowed_origins.contains(&"*".to_string()) {
199 return Some("*".to_string());
200 }
201
202 let request_origin = request_origin?;
204
205 if config.allowed_origins.contains(&request_origin.to_string()) {
207 Some(request_origin.to_string())
208 } else {
209 None
211 }
212}
213
214pub fn validate_config(config: &CorsConfig) -> Result<()> {
216 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 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 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 config.allow_credentials = true;
265 assert!(validate_config(&config).is_err());
266
267 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 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 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}