1use std::net::SocketAddr;
4
5use axum::Json;
6use axum::extract::{ConnectInfo, Request, State};
7#[cfg(feature = "prometheus")]
8use axum::http::Method;
9use axum::http::header::{
10 CONTENT_SECURITY_POLICY, HOST, LOCATION, STRICT_TRANSPORT_SECURITY, X_CONTENT_TYPE_OPTIONS,
11 X_FRAME_OPTIONS, X_XSS_PROTECTION,
12};
13use axum::http::{HeaderValue, StatusCode};
14use axum::middleware::Next;
15use axum::response::{IntoResponse, Response};
16use serde_json::json;
17use subtle::ConstantTimeEq;
18
19use crate::rate_limit::TrustedProxies;
20
21pub async fn worker_token_auth(req: Request, next: Next) -> Response {
26 let expected = req.extensions().get::<WorkerToken>().map(|t| t.0.clone());
27
28 let provided = req
29 .headers()
30 .get("authorization")
31 .and_then(|v| v.to_str().ok())
32 .and_then(|v| v.strip_prefix("Bearer "))
33 .map(|t| t.to_string());
34
35 match (expected, provided) {
36 (Some(expected), Some(provided))
37 if expected.as_bytes().ct_eq(provided.as_bytes()).into() =>
38 {
39 next.run(req).await
40 }
41 _ => (
42 StatusCode::UNAUTHORIZED,
43 Json(json!({
44 "error": {
45 "code": "INVALID_WORKER_TOKEN",
46 "message": "Invalid or missing worker token",
47 }
48 })),
49 )
50 .into_response(),
51 }
52}
53
54#[derive(Clone)]
56pub struct WorkerToken(pub String);
57
58#[cfg(feature = "prometheus")]
63pub const UNMATCHED_PATH_LABEL: &str = "unmatched";
64
65#[cfg(feature = "prometheus")]
71pub const OTHER_METHOD_LABEL: &str = "other";
72
73#[cfg(feature = "prometheus")]
76fn method_label(method: &Method) -> &'static str {
77 match *method {
78 Method::GET => "GET",
79 Method::POST => "POST",
80 Method::PUT => "PUT",
81 Method::PATCH => "PATCH",
82 Method::DELETE => "DELETE",
83 Method::HEAD => "HEAD",
84 Method::OPTIONS => "OPTIONS",
85 _ => OTHER_METHOD_LABEL,
86 }
87}
88
89#[cfg(feature = "prometheus")]
113pub async fn request_metrics(req: Request, next: Next) -> Response {
114 use std::time::Instant;
115
116 use axum::extract::MatchedPath;
117 use ironflow_core::metric_names::{API_REQUEST_DURATION_SECONDS, API_REQUESTS_TOTAL};
118 use metrics::{counter, histogram};
119
120 let method = method_label(req.method());
121 let path = req
122 .extensions()
123 .get::<MatchedPath>()
124 .map_or(UNMATCHED_PATH_LABEL, MatchedPath::as_str)
125 .to_string();
126 let start = Instant::now();
127
128 let resp = next.run(req).await;
129
130 let status = resp.status().as_u16().to_string();
131 let duration = start.elapsed().as_secs_f64();
132
133 counter!(API_REQUESTS_TOTAL, "method" => method, "path" => path.clone(), "status" => status)
134 .increment(1);
135 histogram!(API_REQUEST_DURATION_SECONDS, "method" => method, "path" => path).record(duration);
136
137 resp
138}
139
140pub async fn https_redirect(
171 State(trusted): State<TrustedProxies>,
172 req: Request,
173 next: Next,
174) -> Response {
175 let from_trusted_proxy = req
177 .extensions()
178 .get::<ConnectInfo<SocketAddr>>()
179 .is_some_and(|ConnectInfo(addr)| trusted.contains(addr.ip()));
180 if !from_trusted_proxy {
181 return next.run(req).await;
182 }
183
184 let is_http = req
185 .headers()
186 .get("x-forwarded-proto")
187 .and_then(|v| v.to_str().ok())
188 .and_then(|v| v.split(',').next())
189 .is_some_and(|v| v.trim().eq_ignore_ascii_case("http"));
190
191 if is_http {
192 let host = req
193 .headers()
194 .get("x-forwarded-host")
195 .or_else(|| req.headers().get(HOST))
196 .and_then(|v| v.to_str().ok())
197 .map(|v| v.split(',').next().unwrap_or(v).trim())
198 .filter(|v| !v.is_empty());
199 let path = req.uri().path_and_query().map_or("/", |pq| pq.as_str());
200
201 if let Some(host) = host
202 && let Ok(location) = HeaderValue::from_str(&format!("https://{host}{path}"))
203 {
204 let mut resp = StatusCode::PERMANENT_REDIRECT.into_response();
205 resp.headers_mut().insert(LOCATION, location);
206 return resp;
207 }
208 }
209
210 next.run(req).await
211}
212
213pub async fn security_headers(req: Request, next: Next) -> Response {
222 let mut resp = next.run(req).await;
223 let headers = resp.headers_mut();
224
225 headers.insert(X_CONTENT_TYPE_OPTIONS, HeaderValue::from_static("nosniff"));
226 headers.insert(X_FRAME_OPTIONS, HeaderValue::from_static("DENY"));
227 headers.insert(X_XSS_PROTECTION, HeaderValue::from_static("1; mode=block"));
228 headers.insert(
229 STRICT_TRANSPORT_SECURITY,
230 HeaderValue::from_static("max-age=63072000; includeSubDomains"),
231 );
232 headers.insert(
233 CONTENT_SECURITY_POLICY,
234 HeaderValue::from_static(
235 "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; font-src 'self' data:; connect-src 'self'",
236 ),
237 );
238
239 resp
240}
241
242#[cfg(test)]
243mod tests {
244
245 use axum::Router;
246 use axum::body::Body;
247 use axum::extract::ConnectInfo;
248 use axum::http::{Request, StatusCode};
249 use http_body_util::BodyExt;
250 use ironflow_core::providers::claude::ClaudeCodeProvider;
251 use ironflow_engine::engine::Engine;
252 use ironflow_engine::notify::Event;
253 use ironflow_store::memory::InMemoryStore;
254 use serde_json::Value as JsonValue;
255 use std::fs::write;
256 use std::net::SocketAddr;
257 use std::sync::Arc;
258 use tempfile::TempDir;
259 use tokio::sync::broadcast;
260 use tower::ServiceExt;
261
262 use crate::routes::{RouterConfig, create_router};
263 use crate::state::AppState;
264
265 fn test_state() -> AppState {
266 let store = Arc::new(InMemoryStore::new());
267 let provider = Arc::new(ClaudeCodeProvider::new());
268 let engine = Arc::new(Engine::new(store.clone(), provider));
269 let jwt_config = Arc::new(ironflow_auth::jwt::JwtConfig {
270 secret: "test-secret".to_string(),
271 access_token_ttl_secs: 900,
272 refresh_token_ttl_secs: 604800,
273 cookie_domain: None,
274 cookie_secure: false,
275 });
276 let (event_sender, _) = broadcast::channel::<Event>(1);
277 AppState::new(
278 store,
279 engine,
280 jwt_config,
281 "test-worker-token".to_string(),
282 event_sender,
283 )
284 }
285
286 #[tokio::test]
287 async fn worker_token_valid() {
288 let state = test_state();
289 let app = create_router(state.clone(), RouterConfig::default());
290
291 let req = Request::builder()
292 .uri("/api/v1/internal/runs/next")
293 .header("authorization", "Bearer test-worker-token")
294 .body(Body::empty())
295 .unwrap();
296
297 let resp = app.oneshot(req).await.unwrap();
298 assert_eq!(resp.status(), StatusCode::OK);
299 }
300
301 #[tokio::test]
302 async fn worker_token_missing() {
303 let state = test_state();
304 let app = create_router(state, RouterConfig::default());
305
306 let req = Request::builder()
307 .uri("/api/v1/internal/runs/next")
308 .body(Body::empty())
309 .unwrap();
310
311 let resp = app.oneshot(req).await.unwrap();
312 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
313
314 let body = resp.into_body().collect().await.unwrap().to_bytes();
315 let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
316 assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
317 }
318
319 #[tokio::test]
320 async fn worker_token_invalid() {
321 let state = test_state();
322 let app = create_router(state, RouterConfig::default());
323
324 let req = Request::builder()
325 .uri("/api/v1/internal/runs/next")
326 .header("authorization", "Bearer wrong-token")
327 .body(Body::empty())
328 .unwrap();
329
330 let resp = app.oneshot(req).await.unwrap();
331 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
332
333 let body = resp.into_body().collect().await.unwrap().to_bytes();
334 let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
335 assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
336 }
337
338 fn https_request(method: &str, proto: Option<&str>) -> Request<Body> {
339 let mut builder = Request::builder()
340 .method(method)
341 .uri("/api/v1/health-check?x=1")
342 .header("host", "example.com");
343 if let Some(proto) = proto {
344 builder = builder.header("x-forwarded-proto", proto);
345 }
346 builder.body(Body::empty()).unwrap()
347 }
348
349 const TRUSTED_NETWORK: &str = "10.0.0.0/8";
351
352 const TRUSTED_PEER: &str = "10.0.0.2";
354
355 fn from_peer(mut req: Request<Body>, peer: &str) -> Request<Body> {
357 let addr: SocketAddr = format!("{peer}:4000").parse().unwrap();
358 req.extensions_mut().insert(ConnectInfo(addr));
359 req
360 }
361
362 fn enforcing_router() -> Router {
363 let config = RouterConfig {
364 enforce_https: true,
365 trusted_proxies: TRUSTED_NETWORK.parse().unwrap(),
366 ..RouterConfig::default()
367 };
368 create_router(test_state(), config)
369 }
370
371 fn spoofed_request() -> Request<Body> {
373 Request::builder()
374 .uri("/api/v1/health-check")
375 .header("host", "example.com")
376 .header("x-forwarded-host", "evil.example")
377 .header("x-forwarded-proto", "http")
378 .body(Body::empty())
379 .unwrap()
380 }
381
382 #[tokio::test]
383 async fn https_redirect_ignores_forwarded_host_from_an_untrusted_peer() {
384 let resp = enforcing_router()
385 .oneshot(from_peer(spoofed_request(), "203.0.113.7"))
386 .await
387 .unwrap();
388
389 assert_eq!(resp.status(), StatusCode::OK);
390 assert!(resp.headers().get("location").is_none());
391 }
392
393 #[tokio::test]
394 async fn https_redirect_ignores_forwarded_headers_without_peer_address() {
395 let req = spoofed_request();
396
397 let resp = enforcing_router().oneshot(req).await.unwrap();
398
399 assert_eq!(resp.status(), StatusCode::OK);
400 assert!(resp.headers().get("location").is_none());
401 }
402
403 #[tokio::test]
404 async fn https_redirect_accepts_an_ipv4_mapped_trusted_peer() {
405 let resp = enforcing_router()
406 .oneshot(from_peer(spoofed_request(), "[::ffff:10.0.0.2]"))
407 .await
408 .unwrap();
409
410 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
411 }
412
413 #[tokio::test]
414 async fn https_redirect_redirects_plain_http() {
415 let resp = enforcing_router()
416 .oneshot(from_peer(https_request("GET", Some("http")), TRUSTED_PEER))
417 .await
418 .unwrap();
419
420 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
421 assert_eq!(
422 resp.headers().get("location").unwrap(),
423 "https://example.com/api/v1/health-check?x=1"
424 );
425 assert!(resp.headers().get("strict-transport-security").is_some());
426 }
427
428 #[tokio::test]
429 async fn https_redirect_uses_first_value_and_ignores_case() {
430 let resp = enforcing_router()
431 .oneshot(from_peer(
432 https_request("GET", Some("HTTP, https")),
433 TRUSTED_PEER,
434 ))
435 .await
436 .unwrap();
437
438 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
439 }
440
441 #[tokio::test]
442 async fn https_redirect_prefers_forwarded_host() {
443 let req = Request::builder()
444 .uri("/api/v1/health-check")
445 .header("host", "internal:3000")
446 .header("x-forwarded-host", "public.example.com")
447 .header("x-forwarded-proto", "http")
448 .body(Body::empty())
449 .unwrap();
450
451 let resp = enforcing_router()
452 .oneshot(from_peer(req, TRUSTED_PEER))
453 .await
454 .unwrap();
455
456 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
457 assert_eq!(
458 resp.headers().get("location").unwrap(),
459 "https://public.example.com/api/v1/health-check"
460 );
461 }
462
463 #[tokio::test]
464 async fn https_redirect_passes_https_through() {
465 let resp = enforcing_router()
466 .oneshot(from_peer(https_request("GET", Some("https")), TRUSTED_PEER))
467 .await
468 .unwrap();
469
470 assert_eq!(resp.status(), StatusCode::OK);
471 }
472
473 #[tokio::test]
474 async fn https_redirect_passes_without_forwarded_proto() {
475 let resp = enforcing_router()
476 .oneshot(from_peer(https_request("GET", None), TRUSTED_PEER))
477 .await
478 .unwrap();
479
480 assert_eq!(resp.status(), StatusCode::OK);
481 }
482
483 #[tokio::test]
484 async fn https_redirect_without_host_passes_through() {
485 let req = Request::builder()
486 .uri("/api/v1/health-check")
487 .header("x-forwarded-proto", "http")
488 .body(Body::empty())
489 .unwrap();
490
491 let resp = enforcing_router()
492 .oneshot(from_peer(req, TRUSTED_PEER))
493 .await
494 .unwrap();
495
496 assert_eq!(resp.status(), StatusCode::OK);
497 }
498
499 #[tokio::test]
500 async fn https_redirect_disabled_by_default() {
501 let app = create_router(test_state(), RouterConfig::default());
502
503 let resp = app
504 .oneshot(from_peer(https_request("GET", Some("http")), TRUSTED_PEER))
505 .await
506 .unwrap();
507
508 assert_eq!(resp.status(), StatusCode::OK);
509 }
510
511 #[tokio::test]
512 async fn https_redirect_post_gets_308() {
513 let resp = enforcing_router()
514 .oneshot(from_peer(https_request("POST", Some("http")), TRUSTED_PEER))
515 .await
516 .unwrap();
517
518 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
519 }
520
521 fn dashboard_router(enforce_https: bool) -> (Router, TempDir) {
524 let dir = TempDir::new().unwrap();
525 write(dir.path().join("index.html"), "<!doctype html>").unwrap();
526 let config = RouterConfig {
527 enforce_https,
528 trusted_proxies: TRUSTED_NETWORK.parse().unwrap(),
529 dashboard_dir: Some(dir.path().to_path_buf()),
530 ..RouterConfig::default()
531 };
532 (create_router(test_state(), config), dir)
533 }
534
535 #[tokio::test]
536 async fn https_redirect_covers_dashboard() {
537 let (app, _dir) = dashboard_router(true);
540 let req = Request::builder()
541 .uri("/runs?page=2")
542 .header("host", "example.com")
543 .header("x-forwarded-proto", "http")
544 .body(Body::empty())
545 .unwrap();
546
547 let resp = app.oneshot(from_peer(req, TRUSTED_PEER)).await.unwrap();
548
549 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
550 assert_eq!(
551 resp.headers().get("location").unwrap(),
552 "https://example.com/runs?page=2"
553 );
554 }
555
556 #[tokio::test]
557 async fn security_headers_cover_dashboard() {
558 let (app, _dir) = dashboard_router(false);
559 let req = Request::builder().uri("/").body(Body::empty()).unwrap();
560
561 let resp = app.oneshot(req).await.unwrap();
562
563 assert_eq!(resp.status(), StatusCode::OK);
564 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
565 assert!(resp.headers().get("strict-transport-security").is_some());
566 assert!(resp.headers().get("content-security-policy").is_some());
567 }
568
569 #[tokio::test]
570 async fn security_headers_present() {
571 let state = test_state();
572 let app = create_router(state, RouterConfig::default());
573
574 let req = Request::builder()
575 .uri("/api/v1/health-check")
576 .body(Body::empty())
577 .unwrap();
578
579 let resp = app.oneshot(req).await.unwrap();
580
581 assert_eq!(
582 resp.headers().get("x-content-type-options").unwrap(),
583 "nosniff"
584 );
585 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
586 assert_eq!(
587 resp.headers().get("x-xss-protection").unwrap(),
588 "1; mode=block"
589 );
590 assert_eq!(
591 resp.headers().get("strict-transport-security").unwrap(),
592 "max-age=63072000; includeSubDomains"
593 );
594 assert!(
595 resp.headers()
596 .get("content-security-policy")
597 .unwrap()
598 .to_str()
599 .unwrap()
600 .contains("default-src 'self'")
601 );
602 }
603
604 #[cfg(feature = "prometheus")]
605 #[test]
606 fn method_label_bounds_methods() {
607 use axum::http::Method;
608
609 use super::{OTHER_METHOD_LABEL, method_label};
610
611 assert_eq!(method_label(&Method::GET), "GET");
612 assert_eq!(method_label(&Method::POST), "POST");
613 assert_eq!(method_label(&Method::PUT), "PUT");
614 assert_eq!(method_label(&Method::PATCH), "PATCH");
615 assert_eq!(method_label(&Method::DELETE), "DELETE");
616 assert_eq!(method_label(&Method::HEAD), "HEAD");
617 assert_eq!(method_label(&Method::OPTIONS), "OPTIONS");
618 assert_eq!(method_label(&Method::TRACE), OTHER_METHOD_LABEL);
619 assert_eq!(method_label(&Method::CONNECT), OTHER_METHOD_LABEL);
620 let custom = Method::from_bytes(b"FOOBAR").expect("valid method token");
621 assert_eq!(method_label(&custom), OTHER_METHOD_LABEL);
622 }
623}