1use axum::Json;
4use axum::extract::Request;
5use axum::http::header::{
6 CONTENT_SECURITY_POLICY, HOST, LOCATION, STRICT_TRANSPORT_SECURITY, X_CONTENT_TYPE_OPTIONS,
7 X_FRAME_OPTIONS, X_XSS_PROTECTION,
8};
9use axum::http::{HeaderValue, StatusCode};
10use axum::middleware::Next;
11use axum::response::{IntoResponse, Response};
12use serde_json::json;
13use subtle::ConstantTimeEq;
14
15pub async fn worker_token_auth(req: Request, next: Next) -> Response {
20 let expected = req.extensions().get::<WorkerToken>().map(|t| t.0.clone());
21
22 let provided = req
23 .headers()
24 .get("authorization")
25 .and_then(|v| v.to_str().ok())
26 .and_then(|v| v.strip_prefix("Bearer "))
27 .map(|t| t.to_string());
28
29 match (expected, provided) {
30 (Some(expected), Some(provided))
31 if expected.as_bytes().ct_eq(provided.as_bytes()).into() =>
32 {
33 next.run(req).await
34 }
35 _ => (
36 StatusCode::UNAUTHORIZED,
37 Json(json!({
38 "error": {
39 "code": "INVALID_WORKER_TOKEN",
40 "message": "Invalid or missing worker token",
41 }
42 })),
43 )
44 .into_response(),
45 }
46}
47
48#[derive(Clone)]
50pub struct WorkerToken(pub String);
51
52#[cfg(feature = "prometheus")]
57pub const UNMATCHED_PATH_LABEL: &str = "unmatched";
58
59#[cfg(feature = "prometheus")]
81pub async fn request_metrics(req: Request, next: Next) -> Response {
82 use std::time::Instant;
83
84 use axum::extract::MatchedPath;
85 use ironflow_core::metric_names::{API_REQUEST_DURATION_SECONDS, API_REQUESTS_TOTAL};
86 use metrics::{counter, histogram};
87
88 let method = req.method().to_string();
89 let path = req
90 .extensions()
91 .get::<MatchedPath>()
92 .map_or(UNMATCHED_PATH_LABEL, MatchedPath::as_str)
93 .to_string();
94 let start = Instant::now();
95
96 let resp = next.run(req).await;
97
98 let status = resp.status().as_u16().to_string();
99 let duration = start.elapsed().as_secs_f64();
100
101 counter!(API_REQUESTS_TOTAL, "method" => method.clone(), "path" => path.clone(), "status" => status).increment(1);
102 histogram!(API_REQUEST_DURATION_SECONDS, "method" => method, "path" => path).record(duration);
103
104 resp
105}
106
107pub async fn https_redirect(req: Request, next: Next) -> Response {
130 let is_http = req
131 .headers()
132 .get("x-forwarded-proto")
133 .and_then(|v| v.to_str().ok())
134 .and_then(|v| v.split(',').next())
135 .is_some_and(|v| v.trim().eq_ignore_ascii_case("http"));
136
137 if is_http {
138 let host = req
139 .headers()
140 .get("x-forwarded-host")
141 .or_else(|| req.headers().get(HOST))
142 .and_then(|v| v.to_str().ok())
143 .map(|v| v.split(',').next().unwrap_or(v).trim())
144 .filter(|v| !v.is_empty());
145 let path = req.uri().path_and_query().map_or("/", |pq| pq.as_str());
146
147 if let Some(host) = host
148 && let Ok(location) = HeaderValue::from_str(&format!("https://{host}{path}"))
149 {
150 let mut resp = StatusCode::PERMANENT_REDIRECT.into_response();
151 resp.headers_mut().insert(LOCATION, location);
152 return resp;
153 }
154 }
155
156 next.run(req).await
157}
158
159pub async fn security_headers(req: Request, next: Next) -> Response {
168 let mut resp = next.run(req).await;
169 let headers = resp.headers_mut();
170
171 headers.insert(X_CONTENT_TYPE_OPTIONS, HeaderValue::from_static("nosniff"));
172 headers.insert(X_FRAME_OPTIONS, HeaderValue::from_static("DENY"));
173 headers.insert(X_XSS_PROTECTION, HeaderValue::from_static("1; mode=block"));
174 headers.insert(
175 STRICT_TRANSPORT_SECURITY,
176 HeaderValue::from_static("max-age=63072000; includeSubDomains"),
177 );
178 headers.insert(
179 CONTENT_SECURITY_POLICY,
180 HeaderValue::from_static(
181 "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; font-src 'self' data:; connect-src 'self'",
182 ),
183 );
184
185 resp
186}
187
188#[cfg(test)]
189mod tests {
190
191 use axum::Router;
192 use axum::body::Body;
193 use axum::http::{Request, StatusCode};
194 use http_body_util::BodyExt;
195 use ironflow_core::providers::claude::ClaudeCodeProvider;
196 use ironflow_engine::engine::Engine;
197 use ironflow_engine::notify::Event;
198 use ironflow_store::memory::InMemoryStore;
199 use serde_json::Value as JsonValue;
200 use std::fs::write;
201 use std::sync::Arc;
202 use tempfile::TempDir;
203 use tokio::sync::broadcast;
204 use tower::ServiceExt;
205
206 use crate::routes::{RouterConfig, create_router};
207 use crate::state::AppState;
208
209 fn test_state() -> AppState {
210 let store = Arc::new(InMemoryStore::new());
211 let provider = Arc::new(ClaudeCodeProvider::new());
212 let engine = Arc::new(Engine::new(store.clone(), provider));
213 let jwt_config = Arc::new(ironflow_auth::jwt::JwtConfig {
214 secret: "test-secret".to_string(),
215 access_token_ttl_secs: 900,
216 refresh_token_ttl_secs: 604800,
217 cookie_domain: None,
218 cookie_secure: false,
219 });
220 let (event_sender, _) = broadcast::channel::<Event>(1);
221 AppState::new(
222 store,
223 engine,
224 jwt_config,
225 "test-worker-token".to_string(),
226 event_sender,
227 )
228 }
229
230 #[tokio::test]
231 async fn worker_token_valid() {
232 let state = test_state();
233 let app = create_router(state.clone(), RouterConfig::default());
234
235 let req = Request::builder()
236 .uri("/api/v1/internal/runs/next")
237 .header("authorization", "Bearer test-worker-token")
238 .body(Body::empty())
239 .unwrap();
240
241 let resp = app.oneshot(req).await.unwrap();
242 assert_eq!(resp.status(), StatusCode::OK);
243 }
244
245 #[tokio::test]
246 async fn worker_token_missing() {
247 let state = test_state();
248 let app = create_router(state, RouterConfig::default());
249
250 let req = Request::builder()
251 .uri("/api/v1/internal/runs/next")
252 .body(Body::empty())
253 .unwrap();
254
255 let resp = app.oneshot(req).await.unwrap();
256 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
257
258 let body = resp.into_body().collect().await.unwrap().to_bytes();
259 let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
260 assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
261 }
262
263 #[tokio::test]
264 async fn worker_token_invalid() {
265 let state = test_state();
266 let app = create_router(state, RouterConfig::default());
267
268 let req = Request::builder()
269 .uri("/api/v1/internal/runs/next")
270 .header("authorization", "Bearer wrong-token")
271 .body(Body::empty())
272 .unwrap();
273
274 let resp = app.oneshot(req).await.unwrap();
275 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
276
277 let body = resp.into_body().collect().await.unwrap().to_bytes();
278 let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
279 assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
280 }
281
282 fn https_request(method: &str, proto: Option<&str>) -> Request<Body> {
283 let mut builder = Request::builder()
284 .method(method)
285 .uri("/api/v1/health-check?x=1")
286 .header("host", "example.com");
287 if let Some(proto) = proto {
288 builder = builder.header("x-forwarded-proto", proto);
289 }
290 builder.body(Body::empty()).unwrap()
291 }
292
293 fn enforcing_router() -> Router {
294 let config = RouterConfig {
295 enforce_https: true,
296 ..RouterConfig::default()
297 };
298 create_router(test_state(), config)
299 }
300
301 #[tokio::test]
302 async fn https_redirect_redirects_plain_http() {
303 let resp = enforcing_router()
304 .oneshot(https_request("GET", Some("http")))
305 .await
306 .unwrap();
307
308 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
309 assert_eq!(
310 resp.headers().get("location").unwrap(),
311 "https://example.com/api/v1/health-check?x=1"
312 );
313 assert!(resp.headers().get("strict-transport-security").is_some());
314 }
315
316 #[tokio::test]
317 async fn https_redirect_uses_first_value_and_ignores_case() {
318 let resp = enforcing_router()
319 .oneshot(https_request("GET", Some("HTTP, https")))
320 .await
321 .unwrap();
322
323 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
324 }
325
326 #[tokio::test]
327 async fn https_redirect_prefers_forwarded_host() {
328 let req = Request::builder()
329 .uri("/api/v1/health-check")
330 .header("host", "internal:3000")
331 .header("x-forwarded-host", "public.example.com")
332 .header("x-forwarded-proto", "http")
333 .body(Body::empty())
334 .unwrap();
335
336 let resp = enforcing_router().oneshot(req).await.unwrap();
337
338 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
339 assert_eq!(
340 resp.headers().get("location").unwrap(),
341 "https://public.example.com/api/v1/health-check"
342 );
343 }
344
345 #[tokio::test]
346 async fn https_redirect_passes_https_through() {
347 let resp = enforcing_router()
348 .oneshot(https_request("GET", Some("https")))
349 .await
350 .unwrap();
351
352 assert_eq!(resp.status(), StatusCode::OK);
353 }
354
355 #[tokio::test]
356 async fn https_redirect_passes_without_forwarded_proto() {
357 let resp = enforcing_router()
358 .oneshot(https_request("GET", None))
359 .await
360 .unwrap();
361
362 assert_eq!(resp.status(), StatusCode::OK);
363 }
364
365 #[tokio::test]
366 async fn https_redirect_without_host_passes_through() {
367 let req = Request::builder()
368 .uri("/api/v1/health-check")
369 .header("x-forwarded-proto", "http")
370 .body(Body::empty())
371 .unwrap();
372
373 let resp = enforcing_router().oneshot(req).await.unwrap();
374
375 assert_eq!(resp.status(), StatusCode::OK);
376 }
377
378 #[tokio::test]
379 async fn https_redirect_disabled_by_default() {
380 let app = create_router(test_state(), RouterConfig::default());
381
382 let resp = app
383 .oneshot(https_request("GET", Some("http")))
384 .await
385 .unwrap();
386
387 assert_eq!(resp.status(), StatusCode::OK);
388 }
389
390 #[tokio::test]
391 async fn https_redirect_post_gets_308() {
392 let resp = enforcing_router()
393 .oneshot(https_request("POST", Some("http")))
394 .await
395 .unwrap();
396
397 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
398 }
399
400 fn dashboard_router(enforce_https: bool) -> (Router, TempDir) {
403 let dir = TempDir::new().unwrap();
404 write(dir.path().join("index.html"), "<!doctype html>").unwrap();
405 let config = RouterConfig {
406 enforce_https,
407 dashboard_dir: Some(dir.path().to_path_buf()),
408 ..RouterConfig::default()
409 };
410 (create_router(test_state(), config), dir)
411 }
412
413 #[tokio::test]
414 async fn https_redirect_covers_dashboard() {
415 let (app, _dir) = dashboard_router(true);
418 let req = Request::builder()
419 .uri("/runs?page=2")
420 .header("host", "example.com")
421 .header("x-forwarded-proto", "http")
422 .body(Body::empty())
423 .unwrap();
424
425 let resp = app.oneshot(req).await.unwrap();
426
427 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
428 assert_eq!(
429 resp.headers().get("location").unwrap(),
430 "https://example.com/runs?page=2"
431 );
432 }
433
434 #[tokio::test]
435 async fn security_headers_cover_dashboard() {
436 let (app, _dir) = dashboard_router(false);
437 let req = Request::builder().uri("/").body(Body::empty()).unwrap();
438
439 let resp = app.oneshot(req).await.unwrap();
440
441 assert_eq!(resp.status(), StatusCode::OK);
442 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
443 assert!(resp.headers().get("strict-transport-security").is_some());
444 assert!(resp.headers().get("content-security-policy").is_some());
445 }
446
447 #[tokio::test]
448 async fn security_headers_present() {
449 let state = test_state();
450 let app = create_router(state, RouterConfig::default());
451
452 let req = Request::builder()
453 .uri("/api/v1/health-check")
454 .body(Body::empty())
455 .unwrap();
456
457 let resp = app.oneshot(req).await.unwrap();
458
459 assert_eq!(
460 resp.headers().get("x-content-type-options").unwrap(),
461 "nosniff"
462 );
463 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
464 assert_eq!(
465 resp.headers().get("x-xss-protection").unwrap(),
466 "1; mode=block"
467 );
468 assert_eq!(
469 resp.headers().get("strict-transport-security").unwrap(),
470 "max-age=63072000; includeSubDomains"
471 );
472 assert!(
473 resp.headers()
474 .get("content-security-policy")
475 .unwrap()
476 .to_str()
477 .unwrap()
478 .contains("default-src 'self'")
479 );
480 }
481}