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 async fn request_metrics(req: Request, next: Next) -> Response {
58 use std::time::Instant;
59
60 use ironflow_core::metric_names::{API_REQUEST_DURATION_SECONDS, API_REQUESTS_TOTAL};
61 use metrics::{counter, histogram};
62
63 let method = req.method().to_string();
64 let path = req.uri().path().to_string();
65 let start = Instant::now();
66
67 let resp = next.run(req).await;
68
69 let status = resp.status().as_u16().to_string();
70 let duration = start.elapsed().as_secs_f64();
71
72 counter!(API_REQUESTS_TOTAL, "method" => method.clone(), "path" => path.clone(), "status" => status).increment(1);
73 histogram!(API_REQUEST_DURATION_SECONDS, "method" => method, "path" => path).record(duration);
74
75 resp
76}
77
78pub async fn https_redirect(req: Request, next: Next) -> Response {
101 let is_http = req
102 .headers()
103 .get("x-forwarded-proto")
104 .and_then(|v| v.to_str().ok())
105 .and_then(|v| v.split(',').next())
106 .is_some_and(|v| v.trim().eq_ignore_ascii_case("http"));
107
108 if is_http {
109 let host = req
110 .headers()
111 .get("x-forwarded-host")
112 .or_else(|| req.headers().get(HOST))
113 .and_then(|v| v.to_str().ok())
114 .map(|v| v.split(',').next().unwrap_or(v).trim())
115 .filter(|v| !v.is_empty());
116 let path = req.uri().path_and_query().map_or("/", |pq| pq.as_str());
117
118 if let Some(host) = host
119 && let Ok(location) = HeaderValue::from_str(&format!("https://{host}{path}"))
120 {
121 let mut resp = StatusCode::PERMANENT_REDIRECT.into_response();
122 resp.headers_mut().insert(LOCATION, location);
123 return resp;
124 }
125 }
126
127 next.run(req).await
128}
129
130pub async fn security_headers(req: Request, next: Next) -> Response {
139 let mut resp = next.run(req).await;
140 let headers = resp.headers_mut();
141
142 headers.insert(X_CONTENT_TYPE_OPTIONS, HeaderValue::from_static("nosniff"));
143 headers.insert(X_FRAME_OPTIONS, HeaderValue::from_static("DENY"));
144 headers.insert(X_XSS_PROTECTION, HeaderValue::from_static("1; mode=block"));
145 headers.insert(
146 STRICT_TRANSPORT_SECURITY,
147 HeaderValue::from_static("max-age=63072000; includeSubDomains"),
148 );
149 headers.insert(
150 CONTENT_SECURITY_POLICY,
151 HeaderValue::from_static(
152 "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; font-src 'self' data:; connect-src 'self'",
153 ),
154 );
155
156 resp
157}
158
159#[cfg(test)]
160mod tests {
161
162 use axum::Router;
163 use axum::body::Body;
164 use axum::http::{Request, StatusCode};
165 use http_body_util::BodyExt;
166 use ironflow_core::providers::claude::ClaudeCodeProvider;
167 use ironflow_engine::engine::Engine;
168 use ironflow_engine::notify::Event;
169 use ironflow_store::memory::InMemoryStore;
170 use serde_json::Value as JsonValue;
171 use std::fs::write;
172 use std::sync::Arc;
173 use tempfile::TempDir;
174 use tokio::sync::broadcast;
175 use tower::ServiceExt;
176
177 use crate::routes::{RouterConfig, create_router};
178 use crate::state::AppState;
179
180 fn test_state() -> AppState {
181 let store = Arc::new(InMemoryStore::new());
182 let provider = Arc::new(ClaudeCodeProvider::new());
183 let engine = Arc::new(Engine::new(store.clone(), provider));
184 let jwt_config = Arc::new(ironflow_auth::jwt::JwtConfig {
185 secret: "test-secret".to_string(),
186 access_token_ttl_secs: 900,
187 refresh_token_ttl_secs: 604800,
188 cookie_domain: None,
189 cookie_secure: false,
190 });
191 let (event_sender, _) = broadcast::channel::<Event>(1);
192 AppState::new(
193 store,
194 engine,
195 jwt_config,
196 "test-worker-token".to_string(),
197 event_sender,
198 )
199 }
200
201 #[tokio::test]
202 async fn worker_token_valid() {
203 let state = test_state();
204 let app = create_router(state.clone(), RouterConfig::default());
205
206 let req = Request::builder()
207 .uri("/api/v1/internal/runs/next")
208 .header("authorization", "Bearer test-worker-token")
209 .body(Body::empty())
210 .unwrap();
211
212 let resp = app.oneshot(req).await.unwrap();
213 assert_eq!(resp.status(), StatusCode::OK);
214 }
215
216 #[tokio::test]
217 async fn worker_token_missing() {
218 let state = test_state();
219 let app = create_router(state, RouterConfig::default());
220
221 let req = Request::builder()
222 .uri("/api/v1/internal/runs/next")
223 .body(Body::empty())
224 .unwrap();
225
226 let resp = app.oneshot(req).await.unwrap();
227 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
228
229 let body = resp.into_body().collect().await.unwrap().to_bytes();
230 let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
231 assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
232 }
233
234 #[tokio::test]
235 async fn worker_token_invalid() {
236 let state = test_state();
237 let app = create_router(state, RouterConfig::default());
238
239 let req = Request::builder()
240 .uri("/api/v1/internal/runs/next")
241 .header("authorization", "Bearer wrong-token")
242 .body(Body::empty())
243 .unwrap();
244
245 let resp = app.oneshot(req).await.unwrap();
246 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
247
248 let body = resp.into_body().collect().await.unwrap().to_bytes();
249 let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
250 assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
251 }
252
253 fn https_request(method: &str, proto: Option<&str>) -> Request<Body> {
254 let mut builder = Request::builder()
255 .method(method)
256 .uri("/api/v1/health-check?x=1")
257 .header("host", "example.com");
258 if let Some(proto) = proto {
259 builder = builder.header("x-forwarded-proto", proto);
260 }
261 builder.body(Body::empty()).unwrap()
262 }
263
264 fn enforcing_router() -> Router {
265 let config = RouterConfig {
266 enforce_https: true,
267 ..RouterConfig::default()
268 };
269 create_router(test_state(), config)
270 }
271
272 #[tokio::test]
273 async fn https_redirect_redirects_plain_http() {
274 let resp = enforcing_router()
275 .oneshot(https_request("GET", Some("http")))
276 .await
277 .unwrap();
278
279 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
280 assert_eq!(
281 resp.headers().get("location").unwrap(),
282 "https://example.com/api/v1/health-check?x=1"
283 );
284 assert!(resp.headers().get("strict-transport-security").is_some());
285 }
286
287 #[tokio::test]
288 async fn https_redirect_uses_first_value_and_ignores_case() {
289 let resp = enforcing_router()
290 .oneshot(https_request("GET", Some("HTTP, https")))
291 .await
292 .unwrap();
293
294 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
295 }
296
297 #[tokio::test]
298 async fn https_redirect_prefers_forwarded_host() {
299 let req = Request::builder()
300 .uri("/api/v1/health-check")
301 .header("host", "internal:3000")
302 .header("x-forwarded-host", "public.example.com")
303 .header("x-forwarded-proto", "http")
304 .body(Body::empty())
305 .unwrap();
306
307 let resp = enforcing_router().oneshot(req).await.unwrap();
308
309 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
310 assert_eq!(
311 resp.headers().get("location").unwrap(),
312 "https://public.example.com/api/v1/health-check"
313 );
314 }
315
316 #[tokio::test]
317 async fn https_redirect_passes_https_through() {
318 let resp = enforcing_router()
319 .oneshot(https_request("GET", Some("https")))
320 .await
321 .unwrap();
322
323 assert_eq!(resp.status(), StatusCode::OK);
324 }
325
326 #[tokio::test]
327 async fn https_redirect_passes_without_forwarded_proto() {
328 let resp = enforcing_router()
329 .oneshot(https_request("GET", None))
330 .await
331 .unwrap();
332
333 assert_eq!(resp.status(), StatusCode::OK);
334 }
335
336 #[tokio::test]
337 async fn https_redirect_without_host_passes_through() {
338 let req = Request::builder()
339 .uri("/api/v1/health-check")
340 .header("x-forwarded-proto", "http")
341 .body(Body::empty())
342 .unwrap();
343
344 let resp = enforcing_router().oneshot(req).await.unwrap();
345
346 assert_eq!(resp.status(), StatusCode::OK);
347 }
348
349 #[tokio::test]
350 async fn https_redirect_disabled_by_default() {
351 let app = create_router(test_state(), RouterConfig::default());
352
353 let resp = app
354 .oneshot(https_request("GET", Some("http")))
355 .await
356 .unwrap();
357
358 assert_eq!(resp.status(), StatusCode::OK);
359 }
360
361 #[tokio::test]
362 async fn https_redirect_post_gets_308() {
363 let resp = enforcing_router()
364 .oneshot(https_request("POST", Some("http")))
365 .await
366 .unwrap();
367
368 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
369 }
370
371 fn dashboard_router(enforce_https: bool) -> (Router, TempDir) {
374 let dir = TempDir::new().unwrap();
375 write(dir.path().join("index.html"), "<!doctype html>").unwrap();
376 let config = RouterConfig {
377 enforce_https,
378 dashboard_dir: Some(dir.path().to_path_buf()),
379 ..RouterConfig::default()
380 };
381 (create_router(test_state(), config), dir)
382 }
383
384 #[tokio::test]
385 async fn https_redirect_covers_dashboard() {
386 let (app, _dir) = dashboard_router(true);
389 let req = Request::builder()
390 .uri("/runs?page=2")
391 .header("host", "example.com")
392 .header("x-forwarded-proto", "http")
393 .body(Body::empty())
394 .unwrap();
395
396 let resp = app.oneshot(req).await.unwrap();
397
398 assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
399 assert_eq!(
400 resp.headers().get("location").unwrap(),
401 "https://example.com/runs?page=2"
402 );
403 }
404
405 #[tokio::test]
406 async fn security_headers_cover_dashboard() {
407 let (app, _dir) = dashboard_router(false);
408 let req = Request::builder().uri("/").body(Body::empty()).unwrap();
409
410 let resp = app.oneshot(req).await.unwrap();
411
412 assert_eq!(resp.status(), StatusCode::OK);
413 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
414 assert!(resp.headers().get("strict-transport-security").is_some());
415 assert!(resp.headers().get("content-security-policy").is_some());
416 }
417
418 #[tokio::test]
419 async fn security_headers_present() {
420 let state = test_state();
421 let app = create_router(state, RouterConfig::default());
422
423 let req = Request::builder()
424 .uri("/api/v1/health-check")
425 .body(Body::empty())
426 .unwrap();
427
428 let resp = app.oneshot(req).await.unwrap();
429
430 assert_eq!(
431 resp.headers().get("x-content-type-options").unwrap(),
432 "nosniff"
433 );
434 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
435 assert_eq!(
436 resp.headers().get("x-xss-protection").unwrap(),
437 "1; mode=block"
438 );
439 assert_eq!(
440 resp.headers().get("strict-transport-security").unwrap(),
441 "max-age=63072000; includeSubDomains"
442 );
443 assert!(
444 resp.headers()
445 .get("content-security-policy")
446 .unwrap()
447 .to_str()
448 .unwrap()
449 .contains("default-src 'self'")
450 );
451 }
452}