Skip to main content

ironflow_api/
middleware.rs

1//! Middleware for internal route protection and HTTP security hardening.
2
3use std::net::SocketAddr;
4
5use axum::Json;
6use axum::extract::{ConnectInfo, Request, State};
7#[cfg(feature = "prometheus")]
8use axum::http::Method;
9use axum::http::header::{
10    CONTENT_SECURITY_POLICY, HOST, LOCATION, STRICT_TRANSPORT_SECURITY, X_CONTENT_TYPE_OPTIONS,
11    X_FRAME_OPTIONS, X_XSS_PROTECTION,
12};
13use axum::http::{HeaderValue, StatusCode};
14use axum::middleware::Next;
15use axum::response::{IntoResponse, Response};
16use serde_json::json;
17use subtle::ConstantTimeEq;
18
19use crate::rate_limit::TrustedProxies;
20
21/// Axum middleware that validates a static worker token.
22///
23/// Extracts `Authorization: Bearer {token}` and compares against the
24/// expected token. Returns 401 if missing or invalid.
25pub async fn worker_token_auth(req: Request, next: Next) -> Response {
26    let expected = req.extensions().get::<WorkerToken>().map(|t| t.0.clone());
27
28    let provided = req
29        .headers()
30        .get("authorization")
31        .and_then(|v| v.to_str().ok())
32        .and_then(|v| v.strip_prefix("Bearer "))
33        .map(|t| t.to_string());
34
35    match (expected, provided) {
36        (Some(expected), Some(provided))
37            if expected.as_bytes().ct_eq(provided.as_bytes()).into() =>
38        {
39            next.run(req).await
40        }
41        _ => (
42            StatusCode::UNAUTHORIZED,
43            Json(json!({
44                "error": {
45                    "code": "INVALID_WORKER_TOKEN",
46                    "message": "Invalid or missing worker token",
47                }
48            })),
49        )
50            .into_response(),
51    }
52}
53
54/// Newtype wrapper for the static worker token, stored in request extensions.
55#[derive(Clone)]
56pub struct WorkerToken(pub String);
57
58/// `path` label of a request that matched no route.
59///
60/// Labelling it with its raw URL would create one time series per URL, so
61/// every unmatched request shares this value.
62#[cfg(feature = "prometheus")]
63pub const UNMATCHED_PATH_LABEL: &str = "unmatched";
64
65/// `method` label of a request whose HTTP method is not one of the seven
66/// standard methods the API serves.
67///
68/// Any token is a valid HTTP method, so labelling requests with the raw method
69/// would let a client create one time series per invented method.
70#[cfg(feature = "prometheus")]
71pub const OTHER_METHOD_LABEL: &str = "other";
72
73/// Bounded `method` label of a request: the standard method name, or
74/// [`OTHER_METHOD_LABEL`] for `TRACE`, `CONNECT` and any extension method.
75#[cfg(feature = "prometheus")]
76fn method_label(method: &Method) -> &'static str {
77    match *method {
78        Method::GET => "GET",
79        Method::POST => "POST",
80        Method::PUT => "PUT",
81        Method::PATCH => "PATCH",
82        Method::DELETE => "DELETE",
83        Method::HEAD => "HEAD",
84        Method::OPTIONS => "OPTIONS",
85        _ => OTHER_METHOD_LABEL,
86    }
87}
88
89/// Middleware that records API request metrics (counter + duration histogram).
90///
91/// Emits `ironflow_api_requests_total` and `ironflow_api_request_duration_seconds`
92/// for every request. The `path` label is the matched route pattern (e.g.
93/// `/api/v1/internal/steps/{id}`), never the raw URL, so ids do not create a
94/// time series each; a request that matched no route is labelled
95/// [`UNMATCHED_PATH_LABEL`]. The `method` label is limited to `GET`, `POST`,
96/// `PUT`, `PATCH`, `DELETE`, `HEAD` and `OPTIONS`; any other method is
97/// labelled [`OTHER_METHOD_LABEL`]. Only compiled when the `prometheus` feature
98/// is enabled.
99///
100/// # Examples
101///
102/// ```
103/// use axum::Router;
104/// use axum::middleware::from_fn;
105/// use axum::routing::get;
106/// use ironflow_api::middleware::request_metrics;
107///
108/// let app: Router = Router::new()
109///     .route("/runs/{id}", get(|| async { "ok" }))
110///     .layer(from_fn(request_metrics));
111/// ```
112#[cfg(feature = "prometheus")]
113pub async fn request_metrics(req: Request, next: Next) -> Response {
114    use std::time::Instant;
115
116    use axum::extract::MatchedPath;
117    use ironflow_core::metric_names::{API_REQUEST_DURATION_SECONDS, API_REQUESTS_TOTAL};
118    use metrics::{counter, histogram};
119
120    let method = method_label(req.method());
121    let path = req
122        .extensions()
123        .get::<MatchedPath>()
124        .map_or(UNMATCHED_PATH_LABEL, MatchedPath::as_str)
125        .to_string();
126    let start = Instant::now();
127
128    let resp = next.run(req).await;
129
130    let status = resp.status().as_u16().to_string();
131    let duration = start.elapsed().as_secs_f64();
132
133    counter!(API_REQUESTS_TOTAL, "method" => method, "path" => path.clone(), "status" => status)
134        .increment(1);
135    histogram!(API_REQUEST_DURATION_SECONDS, "method" => method, "path" => path).record(duration);
136
137    resp
138}
139
140/// Middleware that redirects plain-HTTP requests to HTTPS.
141///
142/// TLS terminates at a reverse proxy, so plain HTTP is detected through the
143/// `X-Forwarded-Proto` header (first value of a comma-separated list,
144/// case-insensitive). When it is `http`, the response is a
145/// `308 Permanent Redirect` to `https://{host}{path_and_query}`, which keeps
146/// the method and body. The host comes from `X-Forwarded-Host`, falling back
147/// to `Host`. Requests without `X-Forwarded-Proto` (probes, worker traffic) and
148/// requests with no usable host pass through untouched.
149///
150/// Forwarding headers are client input unless a proxy overwrote them, so they
151/// are read only when the TCP peer is in [`TrustedProxies`]. The peer comes
152/// from `ConnectInfo<SocketAddr>`, so the router must be served with
153/// `into_make_service_with_connect_info::<SocketAddr>()`. A request from any
154/// other peer, or without connect info, passes through untouched: with an
155/// empty [`TrustedProxies`], nothing is ever redirected.
156///
157/// # Examples
158///
159/// ```
160/// use axum::Router;
161/// use axum::middleware::from_fn_with_state;
162/// use axum::routing::get;
163/// use ironflow_api::middleware::https_redirect;
164/// use ironflow_api::rate_limit::TrustedProxies;
165///
166/// let app: Router = Router::new()
167///     .route("/", get(|| async { "ok" }))
168///     .layer(from_fn_with_state(TrustedProxies::default(), https_redirect));
169/// ```
170pub async fn https_redirect(
171    State(trusted): State<TrustedProxies>,
172    req: Request,
173    next: Next,
174) -> Response {
175    // Forwarding headers are client input unless a trusted proxy overwrote them.
176    let from_trusted_proxy = req
177        .extensions()
178        .get::<ConnectInfo<SocketAddr>>()
179        .is_some_and(|ConnectInfo(addr)| trusted.contains(addr.ip()));
180    if !from_trusted_proxy {
181        return next.run(req).await;
182    }
183
184    let is_http = req
185        .headers()
186        .get("x-forwarded-proto")
187        .and_then(|v| v.to_str().ok())
188        .and_then(|v| v.split(',').next())
189        .is_some_and(|v| v.trim().eq_ignore_ascii_case("http"));
190
191    if is_http {
192        let host = req
193            .headers()
194            .get("x-forwarded-host")
195            .or_else(|| req.headers().get(HOST))
196            .and_then(|v| v.to_str().ok())
197            .map(|v| v.split(',').next().unwrap_or(v).trim())
198            .filter(|v| !v.is_empty());
199        let path = req.uri().path_and_query().map_or("/", |pq| pq.as_str());
200
201        if let Some(host) = host
202            && let Ok(location) = HeaderValue::from_str(&format!("https://{host}{path}"))
203        {
204            let mut resp = StatusCode::PERMANENT_REDIRECT.into_response();
205            resp.headers_mut().insert(LOCATION, location);
206            return resp;
207        }
208    }
209
210    next.run(req).await
211}
212
213/// Middleware that injects standard HTTP security headers on every response.
214///
215/// Headers set:
216/// - `X-Content-Type-Options: nosniff` — prevents MIME-type sniffing
217/// - `X-Frame-Options: DENY` — blocks clickjacking via iframes
218/// - `X-XSS-Protection: 1; mode=block` — legacy XSS filter hint
219/// - `Strict-Transport-Security: max-age=63072000; includeSubDomains` — enforces HTTPS for 2 years
220/// - `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'`
221pub async fn security_headers(req: Request, next: Next) -> Response {
222    let mut resp = next.run(req).await;
223    let headers = resp.headers_mut();
224
225    headers.insert(X_CONTENT_TYPE_OPTIONS, HeaderValue::from_static("nosniff"));
226    headers.insert(X_FRAME_OPTIONS, HeaderValue::from_static("DENY"));
227    headers.insert(X_XSS_PROTECTION, HeaderValue::from_static("1; mode=block"));
228    headers.insert(
229        STRICT_TRANSPORT_SECURITY,
230        HeaderValue::from_static("max-age=63072000; includeSubDomains"),
231    );
232    headers.insert(
233        CONTENT_SECURITY_POLICY,
234        HeaderValue::from_static(
235            "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; font-src 'self' data:; connect-src 'self'",
236        ),
237    );
238
239    resp
240}
241
242#[cfg(test)]
243mod tests {
244
245    use axum::Router;
246    use axum::body::Body;
247    use axum::extract::ConnectInfo;
248    use axum::http::{Request, StatusCode};
249    use http_body_util::BodyExt;
250    use ironflow_core::providers::claude::ClaudeCodeProvider;
251    use ironflow_engine::engine::Engine;
252    use ironflow_engine::notify::Event;
253    use ironflow_store::memory::InMemoryStore;
254    use serde_json::Value as JsonValue;
255    use std::fs::write;
256    use std::net::SocketAddr;
257    use std::sync::Arc;
258    use tempfile::TempDir;
259    use tokio::sync::broadcast;
260    use tower::ServiceExt;
261
262    use crate::routes::{RouterConfig, create_router};
263    use crate::state::AppState;
264
265    fn test_state() -> AppState {
266        let store = Arc::new(InMemoryStore::new());
267        let provider = Arc::new(ClaudeCodeProvider::new());
268        let engine = Arc::new(Engine::new(store.clone(), provider));
269        let jwt_config = Arc::new(ironflow_auth::jwt::JwtConfig {
270            secret: "test-secret".to_string(),
271            access_token_ttl_secs: 900,
272            refresh_token_ttl_secs: 604800,
273            cookie_domain: None,
274            cookie_secure: false,
275        });
276        let (event_sender, _) = broadcast::channel::<Event>(1);
277        AppState::new(
278            store,
279            engine,
280            jwt_config,
281            "test-worker-token".to_string(),
282            event_sender,
283        )
284    }
285
286    #[tokio::test]
287    async fn worker_token_valid() {
288        let state = test_state();
289        let app = create_router(state.clone(), RouterConfig::default());
290
291        let req = Request::builder()
292            .uri("/api/v1/internal/runs/next")
293            .header("authorization", "Bearer test-worker-token")
294            .body(Body::empty())
295            .unwrap();
296
297        let resp = app.oneshot(req).await.unwrap();
298        assert_eq!(resp.status(), StatusCode::OK);
299    }
300
301    #[tokio::test]
302    async fn worker_token_missing() {
303        let state = test_state();
304        let app = create_router(state, RouterConfig::default());
305
306        let req = Request::builder()
307            .uri("/api/v1/internal/runs/next")
308            .body(Body::empty())
309            .unwrap();
310
311        let resp = app.oneshot(req).await.unwrap();
312        assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
313
314        let body = resp.into_body().collect().await.unwrap().to_bytes();
315        let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
316        assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
317    }
318
319    #[tokio::test]
320    async fn worker_token_invalid() {
321        let state = test_state();
322        let app = create_router(state, RouterConfig::default());
323
324        let req = Request::builder()
325            .uri("/api/v1/internal/runs/next")
326            .header("authorization", "Bearer wrong-token")
327            .body(Body::empty())
328            .unwrap();
329
330        let resp = app.oneshot(req).await.unwrap();
331        assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
332
333        let body = resp.into_body().collect().await.unwrap().to_bytes();
334        let json_val: JsonValue = serde_json::from_slice(&body).unwrap();
335        assert_eq!(json_val["error"]["code"], "INVALID_WORKER_TOKEN");
336    }
337
338    fn https_request(method: &str, proto: Option<&str>) -> Request<Body> {
339        let mut builder = Request::builder()
340            .method(method)
341            .uri("/api/v1/health-check?x=1")
342            .header("host", "example.com");
343        if let Some(proto) = proto {
344            builder = builder.header("x-forwarded-proto", proto);
345        }
346        builder.body(Body::empty()).unwrap()
347    }
348
349    /// Proxies trusted by the test routers.
350    const TRUSTED_NETWORK: &str = "10.0.0.0/8";
351
352    /// A peer inside [`TRUSTED_NETWORK`].
353    const TRUSTED_PEER: &str = "10.0.0.2";
354
355    /// What a real server attaches from the TCP socket.
356    fn from_peer(mut req: Request<Body>, peer: &str) -> Request<Body> {
357        let addr: SocketAddr = format!("{peer}:4000").parse().unwrap();
358        req.extensions_mut().insert(ConnectInfo(addr));
359        req
360    }
361
362    fn enforcing_router() -> Router {
363        let config = RouterConfig {
364            enforce_https: true,
365            trusted_proxies: TRUSTED_NETWORK.parse().unwrap(),
366            ..RouterConfig::default()
367        };
368        create_router(test_state(), config)
369    }
370
371    /// A plain-HTTP request whose forwarding headers point elsewhere.
372    fn spoofed_request() -> Request<Body> {
373        Request::builder()
374            .uri("/api/v1/health-check")
375            .header("host", "example.com")
376            .header("x-forwarded-host", "evil.example")
377            .header("x-forwarded-proto", "http")
378            .body(Body::empty())
379            .unwrap()
380    }
381
382    #[tokio::test]
383    async fn https_redirect_ignores_forwarded_host_from_an_untrusted_peer() {
384        let resp = enforcing_router()
385            .oneshot(from_peer(spoofed_request(), "203.0.113.7"))
386            .await
387            .unwrap();
388
389        assert_eq!(resp.status(), StatusCode::OK);
390        assert!(resp.headers().get("location").is_none());
391    }
392
393    #[tokio::test]
394    async fn https_redirect_ignores_forwarded_headers_without_peer_address() {
395        let req = spoofed_request();
396
397        let resp = enforcing_router().oneshot(req).await.unwrap();
398
399        assert_eq!(resp.status(), StatusCode::OK);
400        assert!(resp.headers().get("location").is_none());
401    }
402
403    #[tokio::test]
404    async fn https_redirect_accepts_an_ipv4_mapped_trusted_peer() {
405        let resp = enforcing_router()
406            .oneshot(from_peer(spoofed_request(), "[::ffff:10.0.0.2]"))
407            .await
408            .unwrap();
409
410        assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
411    }
412
413    #[tokio::test]
414    async fn https_redirect_redirects_plain_http() {
415        let resp = enforcing_router()
416            .oneshot(from_peer(https_request("GET", Some("http")), TRUSTED_PEER))
417            .await
418            .unwrap();
419
420        assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
421        assert_eq!(
422            resp.headers().get("location").unwrap(),
423            "https://example.com/api/v1/health-check?x=1"
424        );
425        assert!(resp.headers().get("strict-transport-security").is_some());
426    }
427
428    #[tokio::test]
429    async fn https_redirect_uses_first_value_and_ignores_case() {
430        let resp = enforcing_router()
431            .oneshot(from_peer(
432                https_request("GET", Some("HTTP, https")),
433                TRUSTED_PEER,
434            ))
435            .await
436            .unwrap();
437
438        assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
439    }
440
441    #[tokio::test]
442    async fn https_redirect_prefers_forwarded_host() {
443        let req = Request::builder()
444            .uri("/api/v1/health-check")
445            .header("host", "internal:3000")
446            .header("x-forwarded-host", "public.example.com")
447            .header("x-forwarded-proto", "http")
448            .body(Body::empty())
449            .unwrap();
450
451        let resp = enforcing_router()
452            .oneshot(from_peer(req, TRUSTED_PEER))
453            .await
454            .unwrap();
455
456        assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
457        assert_eq!(
458            resp.headers().get("location").unwrap(),
459            "https://public.example.com/api/v1/health-check"
460        );
461    }
462
463    #[tokio::test]
464    async fn https_redirect_passes_https_through() {
465        let resp = enforcing_router()
466            .oneshot(from_peer(https_request("GET", Some("https")), TRUSTED_PEER))
467            .await
468            .unwrap();
469
470        assert_eq!(resp.status(), StatusCode::OK);
471    }
472
473    #[tokio::test]
474    async fn https_redirect_passes_without_forwarded_proto() {
475        let resp = enforcing_router()
476            .oneshot(from_peer(https_request("GET", None), TRUSTED_PEER))
477            .await
478            .unwrap();
479
480        assert_eq!(resp.status(), StatusCode::OK);
481    }
482
483    #[tokio::test]
484    async fn https_redirect_without_host_passes_through() {
485        let req = Request::builder()
486            .uri("/api/v1/health-check")
487            .header("x-forwarded-proto", "http")
488            .body(Body::empty())
489            .unwrap();
490
491        let resp = enforcing_router()
492            .oneshot(from_peer(req, TRUSTED_PEER))
493            .await
494            .unwrap();
495
496        assert_eq!(resp.status(), StatusCode::OK);
497    }
498
499    #[tokio::test]
500    async fn https_redirect_disabled_by_default() {
501        let app = create_router(test_state(), RouterConfig::default());
502
503        let resp = app
504            .oneshot(from_peer(https_request("GET", Some("http")), TRUSTED_PEER))
505            .await
506            .unwrap();
507
508        assert_eq!(resp.status(), StatusCode::OK);
509    }
510
511    #[tokio::test]
512    async fn https_redirect_post_gets_308() {
513        let resp = enforcing_router()
514            .oneshot(from_peer(https_request("POST", Some("http")), TRUSTED_PEER))
515            .await
516            .unwrap();
517
518        assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
519    }
520
521    /// Router serving a dashboard from a temporary directory. The directory
522    /// is returned so it outlives the router.
523    fn dashboard_router(enforce_https: bool) -> (Router, TempDir) {
524        let dir = TempDir::new().unwrap();
525        write(dir.path().join("index.html"), "<!doctype html>").unwrap();
526        let config = RouterConfig {
527            enforce_https,
528            trusted_proxies: TRUSTED_NETWORK.parse().unwrap(),
529            dashboard_dir: Some(dir.path().to_path_buf()),
530            ..RouterConfig::default()
531        };
532        (create_router(test_state(), config), dir)
533    }
534
535    #[tokio::test]
536    async fn https_redirect_covers_dashboard() {
537        // Non-regression: the dashboard fallback used to be attached after the
538        // layers, so plain HTTP on the dashboard was served without redirect.
539        let (app, _dir) = dashboard_router(true);
540        let req = Request::builder()
541            .uri("/runs?page=2")
542            .header("host", "example.com")
543            .header("x-forwarded-proto", "http")
544            .body(Body::empty())
545            .unwrap();
546
547        let resp = app.oneshot(from_peer(req, TRUSTED_PEER)).await.unwrap();
548
549        assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT);
550        assert_eq!(
551            resp.headers().get("location").unwrap(),
552            "https://example.com/runs?page=2"
553        );
554    }
555
556    #[tokio::test]
557    async fn security_headers_cover_dashboard() {
558        let (app, _dir) = dashboard_router(false);
559        let req = Request::builder().uri("/").body(Body::empty()).unwrap();
560
561        let resp = app.oneshot(req).await.unwrap();
562
563        assert_eq!(resp.status(), StatusCode::OK);
564        assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
565        assert!(resp.headers().get("strict-transport-security").is_some());
566        assert!(resp.headers().get("content-security-policy").is_some());
567    }
568
569    #[tokio::test]
570    async fn security_headers_present() {
571        let state = test_state();
572        let app = create_router(state, RouterConfig::default());
573
574        let req = Request::builder()
575            .uri("/api/v1/health-check")
576            .body(Body::empty())
577            .unwrap();
578
579        let resp = app.oneshot(req).await.unwrap();
580
581        assert_eq!(
582            resp.headers().get("x-content-type-options").unwrap(),
583            "nosniff"
584        );
585        assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
586        assert_eq!(
587            resp.headers().get("x-xss-protection").unwrap(),
588            "1; mode=block"
589        );
590        assert_eq!(
591            resp.headers().get("strict-transport-security").unwrap(),
592            "max-age=63072000; includeSubDomains"
593        );
594        assert!(
595            resp.headers()
596                .get("content-security-policy")
597                .unwrap()
598                .to_str()
599                .unwrap()
600                .contains("default-src 'self'")
601        );
602    }
603
604    #[cfg(feature = "prometheus")]
605    #[test]
606    fn method_label_bounds_methods() {
607        use axum::http::Method;
608
609        use super::{OTHER_METHOD_LABEL, method_label};
610
611        assert_eq!(method_label(&Method::GET), "GET");
612        assert_eq!(method_label(&Method::POST), "POST");
613        assert_eq!(method_label(&Method::PUT), "PUT");
614        assert_eq!(method_label(&Method::PATCH), "PATCH");
615        assert_eq!(method_label(&Method::DELETE), "DELETE");
616        assert_eq!(method_label(&Method::HEAD), "HEAD");
617        assert_eq!(method_label(&Method::OPTIONS), "OPTIONS");
618        assert_eq!(method_label(&Method::TRACE), OTHER_METHOD_LABEL);
619        assert_eq!(method_label(&Method::CONNECT), OTHER_METHOD_LABEL);
620        let custom = Method::from_bytes(b"FOOBAR").expect("valid method token");
621        assert_eq!(method_label(&custom), OTHER_METHOD_LABEL);
622    }
623}