1use axum::Json;
4use axum::extract::Request;
5#[cfg(feature = "prometheus")]
6use axum::http::Method;
7use axum::http::header::{
8 CONTENT_SECURITY_POLICY, HOST, LOCATION, STRICT_TRANSPORT_SECURITY, X_CONTENT_TYPE_OPTIONS,
9 X_FRAME_OPTIONS, X_XSS_PROTECTION,
10};
11use axum::http::{HeaderValue, StatusCode};
12use axum::middleware::Next;
13use axum::response::{IntoResponse, Response};
14use serde_json::json;
15use subtle::ConstantTimeEq;
16
17pub async fn worker_token_auth(req: Request, next: Next) -> Response {
22 let expected = req.extensions().get::<WorkerToken>().map(|t| t.0.clone());
23
24 let provided = req
25 .headers()
26 .get("authorization")
27 .and_then(|v| v.to_str().ok())
28 .and_then(|v| v.strip_prefix("Bearer "))
29 .map(|t| t.to_string());
30
31 match (expected, provided) {
32 (Some(expected), Some(provided))
33 if expected.as_bytes().ct_eq(provided.as_bytes()).into() =>
34 {
35 next.run(req).await
36 }
37 _ => (
38 StatusCode::UNAUTHORIZED,
39 Json(json!({
40 "error": {
41 "code": "INVALID_WORKER_TOKEN",
42 "message": "Invalid or missing worker token",
43 }
44 })),
45 )
46 .into_response(),
47 }
48}
49
50#[derive(Clone)]
52pub struct WorkerToken(pub String);
53
54#[cfg(feature = "prometheus")]
59pub const UNMATCHED_PATH_LABEL: &str = "unmatched";
60
61#[cfg(feature = "prometheus")]
67pub const OTHER_METHOD_LABEL: &str = "other";
68
69#[cfg(feature = "prometheus")]
72fn method_label(method: &Method) -> &'static str {
73 match *method {
74 Method::GET => "GET",
75 Method::POST => "POST",
76 Method::PUT => "PUT",
77 Method::PATCH => "PATCH",
78 Method::DELETE => "DELETE",
79 Method::HEAD => "HEAD",
80 Method::OPTIONS => "OPTIONS",
81 _ => OTHER_METHOD_LABEL,
82 }
83}
84
85#[cfg(feature = "prometheus")]
109pub async fn request_metrics(req: Request, next: Next) -> Response {
110 use std::time::Instant;
111
112 use axum::extract::MatchedPath;
113 use ironflow_core::metric_names::{API_REQUEST_DURATION_SECONDS, API_REQUESTS_TOTAL};
114 use metrics::{counter, histogram};
115
116 let method = method_label(req.method());
117 let path = req
118 .extensions()
119 .get::<MatchedPath>()
120 .map_or(UNMATCHED_PATH_LABEL, MatchedPath::as_str)
121 .to_string();
122 let start = Instant::now();
123
124 let resp = next.run(req).await;
125
126 let status = resp.status().as_u16().to_string();
127 let duration = start.elapsed().as_secs_f64();
128
129 counter!(API_REQUESTS_TOTAL, "method" => method, "path" => path.clone(), "status" => status)
130 .increment(1);
131 histogram!(API_REQUEST_DURATION_SECONDS, "method" => method, "path" => path).record(duration);
132
133 resp
134}
135
136pub async fn https_redirect(req: Request, next: Next) -> Response {
159 let is_http = req
160 .headers()
161 .get("x-forwarded-proto")
162 .and_then(|v| v.to_str().ok())
163 .and_then(|v| v.split(',').next())
164 .is_some_and(|v| v.trim().eq_ignore_ascii_case("http"));
165
166 if is_http {
167 let host = req
168 .headers()
169 .get("x-forwarded-host")
170 .or_else(|| req.headers().get(HOST))
171 .and_then(|v| v.to_str().ok())
172 .map(|v| v.split(',').next().unwrap_or(v).trim())
173 .filter(|v| !v.is_empty());
174 let path = req.uri().path_and_query().map_or("/", |pq| pq.as_str());
175
176 if let Some(host) = host
177 && let Ok(location) = HeaderValue::from_str(&format!("https://{host}{path}"))
178 {
179 let mut resp = StatusCode::PERMANENT_REDIRECT.into_response();
180 resp.headers_mut().insert(LOCATION, location);
181 return resp;
182 }
183 }
184
185 next.run(req).await
186}
187
188pub async fn security_headers(req: Request, next: Next) -> Response {
197 let mut resp = next.run(req).await;
198 let headers = resp.headers_mut();
199
200 headers.insert(X_CONTENT_TYPE_OPTIONS, HeaderValue::from_static("nosniff"));
201 headers.insert(X_FRAME_OPTIONS, HeaderValue::from_static("DENY"));
202 headers.insert(X_XSS_PROTECTION, HeaderValue::from_static("1; mode=block"));
203 headers.insert(
204 STRICT_TRANSPORT_SECURITY,
205 HeaderValue::from_static("max-age=63072000; includeSubDomains"),
206 );
207 headers.insert(
208 CONTENT_SECURITY_POLICY,
209 HeaderValue::from_static(
210 "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; font-src 'self' data:; connect-src 'self'",
211 ),
212 );
213
214 resp
215}
216
217#[cfg(test)]
218mod tests {
219
220 use axum::Router;
221 use axum::body::Body;
222 use axum::http::{Request, StatusCode};
223 use http_body_util::BodyExt;
224 use ironflow_core::providers::claude::ClaudeCodeProvider;
225 use ironflow_engine::engine::Engine;
226 use ironflow_engine::notify::Event;
227 use ironflow_store::memory::InMemoryStore;
228 use serde_json::Value as JsonValue;
229 use std::fs::write;
230 use std::sync::Arc;
231 use tempfile::TempDir;
232 use tokio::sync::broadcast;
233 use tower::ServiceExt;
234
235 use crate::routes::{RouterConfig, create_router};
236 use crate::state::AppState;
237
238 fn test_state() -> AppState {
239 let store = Arc::new(InMemoryStore::new());
240 let provider = Arc::new(ClaudeCodeProvider::new());
241 let engine = Arc::new(Engine::new(store.clone(), provider));
242 let jwt_config = Arc::new(ironflow_auth::jwt::JwtConfig {
243 secret: "test-secret".to_string(),
244 access_token_ttl_secs: 900,
245 refresh_token_ttl_secs: 604800,
246 cookie_domain: None,
247 cookie_secure: false,
248 });
249 let (event_sender, _) = broadcast::channel::<Event>(1);
250 AppState::new(
251 store,
252 engine,
253 jwt_config,
254 "test-worker-token".to_string(),
255 event_sender,
256 )
257 }
258
259 #[tokio::test]
260 async fn worker_token_valid() {
261 let state = test_state();
262 let app = create_router(state.clone(), RouterConfig::default());
263
264 let req = Request::builder()
265 .uri("/api/v1/internal/runs/next")
266 .header("authorization", "Bearer test-worker-token")
267 .body(Body::empty())
268 .unwrap();
269
270 let resp = app.oneshot(req).await.unwrap();
271 assert_eq!(resp.status(), StatusCode::OK);
272 }
273
274 #[tokio::test]
275 async fn worker_token_missing() {
276 let state = test_state();
277 let app = create_router(state, RouterConfig::default());
278
279 let req = Request::builder()
280 .uri("/api/v1/internal/runs/next")
281 .body(Body::empty())
282 .unwrap();
283
284 let resp = app.oneshot(req).await.unwrap();
285 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
286
287 let body = resp.into_body().collect().await.unwrap().to_bytes();
288 let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
289 assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
290 }
291
292 #[tokio::test]
293 async fn worker_token_invalid() {
294 let state = test_state();
295 let app = create_router(state, RouterConfig::default());
296
297 let req = Request::builder()
298 .uri("/api/v1/internal/runs/next")
299 .header("authorization", "Bearer wrong-token")
300 .body(Body::empty())
301 .unwrap();
302
303 let resp = app.oneshot(req).await.unwrap();
304 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
305
306 let body = resp.into_body().collect().await.unwrap().to_bytes();
307 let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
308 assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
309 }
310
311 fn https_request(method: &str, proto: Option<&str>) -> Request<Body> {
312 let mut builder = Request::builder()
313 .method(method)
314 .uri("/api/v1/health-check?x=1")
315 .header("host", "example.com");
316 if let Some(proto) = proto {
317 builder = builder.header("x-forwarded-proto", proto);
318 }
319 builder.body(Body::empty()).unwrap()
320 }
321
322 fn enforcing_router() -> Router {
323 let config = RouterConfig {
324 enforce_https: true,
325 ..RouterConfig::default()
326 };
327 create_router(test_state(), config)
328 }
329
330 #[tokio::test]
331 async fn https_redirect_redirects_plain_http() {
332 let resp = enforcing_router()
333 .oneshot(https_request("GET", Some("http")))
334 .await
335 .unwrap();
336
337 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
338 assert_eq!(
339 resp.headers().get("location").unwrap(),
340 "https://example.com/api/v1/health-check?x=1"
341 );
342 assert!(resp.headers().get("strict-transport-security").is_some());
343 }
344
345 #[tokio::test]
346 async fn https_redirect_uses_first_value_and_ignores_case() {
347 let resp = enforcing_router()
348 .oneshot(https_request("GET", Some("HTTP, https")))
349 .await
350 .unwrap();
351
352 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
353 }
354
355 #[tokio::test]
356 async fn https_redirect_prefers_forwarded_host() {
357 let req = Request::builder()
358 .uri("/api/v1/health-check")
359 .header("host", "internal:3000")
360 .header("x-forwarded-host", "public.example.com")
361 .header("x-forwarded-proto", "http")
362 .body(Body::empty())
363 .unwrap();
364
365 let resp = enforcing_router().oneshot(req).await.unwrap();
366
367 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
368 assert_eq!(
369 resp.headers().get("location").unwrap(),
370 "https://public.example.com/api/v1/health-check"
371 );
372 }
373
374 #[tokio::test]
375 async fn https_redirect_passes_https_through() {
376 let resp = enforcing_router()
377 .oneshot(https_request("GET", Some("https")))
378 .await
379 .unwrap();
380
381 assert_eq!(resp.status(), StatusCode::OK);
382 }
383
384 #[tokio::test]
385 async fn https_redirect_passes_without_forwarded_proto() {
386 let resp = enforcing_router()
387 .oneshot(https_request("GET", None))
388 .await
389 .unwrap();
390
391 assert_eq!(resp.status(), StatusCode::OK);
392 }
393
394 #[tokio::test]
395 async fn https_redirect_without_host_passes_through() {
396 let req = Request::builder()
397 .uri("/api/v1/health-check")
398 .header("x-forwarded-proto", "http")
399 .body(Body::empty())
400 .unwrap();
401
402 let resp = enforcing_router().oneshot(req).await.unwrap();
403
404 assert_eq!(resp.status(), StatusCode::OK);
405 }
406
407 #[tokio::test]
408 async fn https_redirect_disabled_by_default() {
409 let app = create_router(test_state(), RouterConfig::default());
410
411 let resp = app
412 .oneshot(https_request("GET", Some("http")))
413 .await
414 .unwrap();
415
416 assert_eq!(resp.status(), StatusCode::OK);
417 }
418
419 #[tokio::test]
420 async fn https_redirect_post_gets_308() {
421 let resp = enforcing_router()
422 .oneshot(https_request("POST", Some("http")))
423 .await
424 .unwrap();
425
426 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
427 }
428
429 fn dashboard_router(enforce_https: bool) -> (Router, TempDir) {
432 let dir = TempDir::new().unwrap();
433 write(dir.path().join("index.html"), "<!doctype html>").unwrap();
434 let config = RouterConfig {
435 enforce_https,
436 dashboard_dir: Some(dir.path().to_path_buf()),
437 ..RouterConfig::default()
438 };
439 (create_router(test_state(), config), dir)
440 }
441
442 #[tokio::test]
443 async fn https_redirect_covers_dashboard() {
444 let (app, _dir) = dashboard_router(true);
447 let req = Request::builder()
448 .uri("/runs?page=2")
449 .header("host", "example.com")
450 .header("x-forwarded-proto", "http")
451 .body(Body::empty())
452 .unwrap();
453
454 let resp = app.oneshot(req).await.unwrap();
455
456 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
457 assert_eq!(
458 resp.headers().get("location").unwrap(),
459 "https://example.com/runs?page=2"
460 );
461 }
462
463 #[tokio::test]
464 async fn security_headers_cover_dashboard() {
465 let (app, _dir) = dashboard_router(false);
466 let req = Request::builder().uri("/").body(Body::empty()).unwrap();
467
468 let resp = app.oneshot(req).await.unwrap();
469
470 assert_eq!(resp.status(), StatusCode::OK);
471 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
472 assert!(resp.headers().get("strict-transport-security").is_some());
473 assert!(resp.headers().get("content-security-policy").is_some());
474 }
475
476 #[tokio::test]
477 async fn security_headers_present() {
478 let state = test_state();
479 let app = create_router(state, RouterConfig::default());
480
481 let req = Request::builder()
482 .uri("/api/v1/health-check")
483 .body(Body::empty())
484 .unwrap();
485
486 let resp = app.oneshot(req).await.unwrap();
487
488 assert_eq!(
489 resp.headers().get("x-content-type-options").unwrap(),
490 "nosniff"
491 );
492 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
493 assert_eq!(
494 resp.headers().get("x-xss-protection").unwrap(),
495 "1; mode=block"
496 );
497 assert_eq!(
498 resp.headers().get("strict-transport-security").unwrap(),
499 "max-age=63072000; includeSubDomains"
500 );
501 assert!(
502 resp.headers()
503 .get("content-security-policy")
504 .unwrap()
505 .to_str()
506 .unwrap()
507 .contains("default-src 'self'")
508 );
509 }
510
511 #[cfg(feature = "prometheus")]
512 #[test]
513 fn method_label_bounds_methods() {
514 use axum::http::Method;
515
516 use super::{OTHER_METHOD_LABEL, method_label};
517
518 assert_eq!(method_label(&Method::GET), "GET");
519 assert_eq!(method_label(&Method::POST), "POST");
520 assert_eq!(method_label(&Method::PUT), "PUT");
521 assert_eq!(method_label(&Method::PATCH), "PATCH");
522 assert_eq!(method_label(&Method::DELETE), "DELETE");
523 assert_eq!(method_label(&Method::HEAD), "HEAD");
524 assert_eq!(method_label(&Method::OPTIONS), "OPTIONS");
525 assert_eq!(method_label(&Method::TRACE), OTHER_METHOD_LABEL);
526 assert_eq!(method_label(&Method::CONNECT), OTHER_METHOD_LABEL);
527 let custom = Method::from_bytes(b"FOOBAR").expect("valid method token");
528 assert_eq!(method_label(&custom), OTHER_METHOD_LABEL);
529 }
530}