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