Skip to main content

ironflow_api/
middleware.rs

1//! Middleware for internal route protection and HTTP security hardening.
2
3use 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
15/// Axum middleware that validates a static worker token.
16///
17/// Extracts `Authorization: Bearer {token}` and compares against the
18/// expected token. Returns 401 if missing or invalid.
19pub 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/// Newtype wrapper for the static worker token, stored in request extensions.
49#[derive(Clone)]
50pub struct WorkerToken(pub String);
51
52/// Middleware that records API request metrics (counter + duration histogram).
53///
54/// Emits `ironflow_api_requests_total` and `ironflow_api_request_duration_seconds`
55/// for every request. Only compiled when the `prometheus` feature is enabled.
56#[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
78/// Middleware that redirects plain-HTTP requests to HTTPS.
79///
80/// TLS terminates at a reverse proxy, so plain HTTP is detected through the
81/// `X-Forwarded-Proto` header (first value of a comma-separated list,
82/// case-insensitive). When it is `http`, the response is a
83/// `308 Permanent Redirect` to `https://{host}{path_and_query}`, which keeps
84/// the method and body. The host comes from `X-Forwarded-Host`, falling back
85/// to `Host`. Requests without `X-Forwarded-Proto` (probes, worker traffic) and
86/// requests with no usable host pass through untouched.
87///
88/// # Examples
89///
90/// ```
91/// use axum::Router;
92/// use axum::middleware::from_fn;
93/// use axum::routing::get;
94/// use ironflow_api::middleware::https_redirect;
95///
96/// let app: Router = Router::new()
97///     .route("/", get(|| async { "ok" }))
98///     .layer(from_fn(https_redirect));
99/// ```
100pub 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
130/// Middleware that injects standard HTTP security headers on every response.
131///
132/// Headers set:
133/// - `X-Content-Type-Options: nosniff` — prevents MIME-type sniffing
134/// - `X-Frame-Options: DENY` — blocks clickjacking via iframes
135/// - `X-XSS-Protection: 1; mode=block` — legacy XSS filter hint
136/// - `Strict-Transport-Security: max-age=63072000; includeSubDomains` — enforces HTTPS for 2 years
137/// - `Content-Security-Policy: default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; font-src 'self' data:; connect-src 'self'`
138pub 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    /// Router serving a dashboard from a temporary directory. The directory
372    /// is returned so it outlives the router.
373    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        // Non-regression: the dashboard fallback used to be attached after the
387        // layers, so plain HTTP on the dashboard was served without redirect.
388        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}