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/// `path` label of a request that matched no route.
53///
54/// Labelling it with its raw URL would create one time series per URL, so
55/// every unmatched request shares this value.
56#[cfg(feature = "prometheus")]
57pub const UNMATCHED_PATH_LABEL: &str = "unmatched";
58
59/// Middleware that records API request metrics (counter + duration histogram).
60///
61/// Emits `ironflow_api_requests_total` and `ironflow_api_request_duration_seconds`
62/// for every request. The `path` label is the matched route pattern (e.g.
63/// `/api/v1/internal/steps/{id}`), never the raw URL, so ids do not create a
64/// time series each; a request that matched no route is labelled
65/// [`UNMATCHED_PATH_LABEL`]. Only compiled when the `prometheus` feature is
66/// enabled.
67///
68/// # Examples
69///
70/// ```
71/// use axum::Router;
72/// use axum::middleware::from_fn;
73/// use axum::routing::get;
74/// use ironflow_api::middleware::request_metrics;
75///
76/// let app: Router = Router::new()
77///     .route("/runs/{id}", get(|| async { "ok" }))
78///     .layer(from_fn(request_metrics));
79/// ```
80#[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
107/// Middleware that redirects plain-HTTP requests to HTTPS.
108///
109/// TLS terminates at a reverse proxy, so plain HTTP is detected through the
110/// `X-Forwarded-Proto` header (first value of a comma-separated list,
111/// case-insensitive). When it is `http`, the response is a
112/// `308 Permanent Redirect` to `https://{host}{path_and_query}`, which keeps
113/// the method and body. The host comes from `X-Forwarded-Host`, falling back
114/// to `Host`. Requests without `X-Forwarded-Proto` (probes, worker traffic) and
115/// requests with no usable host pass through untouched.
116///
117/// # Examples
118///
119/// ```
120/// use axum::Router;
121/// use axum::middleware::from_fn;
122/// use axum::routing::get;
123/// use ironflow_api::middleware::https_redirect;
124///
125/// let app: Router = Router::new()
126///     .route("/", get(|| async { "ok" }))
127///     .layer(from_fn(https_redirect));
128/// ```
129pub 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
159/// Middleware that injects standard HTTP security headers on every response.
160///
161/// Headers set:
162/// - `X-Content-Type-Options: nosniff` — prevents MIME-type sniffing
163/// - `X-Frame-Options: DENY` — blocks clickjacking via iframes
164/// - `X-XSS-Protection: 1; mode=block` — legacy XSS filter hint
165/// - `Strict-Transport-Security: max-age=63072000; includeSubDomains` — enforces HTTPS for 2 years
166/// - `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'`
167pub 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    /// Router serving a dashboard from a temporary directory. The directory
401    /// is returned so it outlives the router.
402    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        // Non-regression: the dashboard fallback used to be attached after the
416        // layers, so plain HTTP on the dashboard was served without redirect.
417        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}