fix(operations): bound handler lifetimes

req: operations/003

examples/kanban/src/main.rs

@@ -1,7 +1,8 @@
-use axum::extract::{Query, State};
+use axum::extract::{Query, Request, State};
 use axum::http::{HeaderMap, StatusCode};
+use axum::middleware::{self, Next};
 use axum::response::sse::{Event, KeepAlive, Sse};
-use axum::response::IntoResponse;
+use axum::response::{IntoResponse, Response};
 use axum::routing::{get, post};
 use axum::{Json, Router};
 use futures_util::{stream, StreamExt};
@@ -18,6 +19,7 @@ use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use std::convert::Infallible;
 use std::fs::{self, OpenOptions};
+use std::future::Future;
 use std::io::Write;
 use std::net::SocketAddr;
 use std::path::PathBuf;
@@ -26,6 +28,9 @@ use std::time::Duration;
 
 const COLUMNS: [(&str, &str); 3] = [("backlog", "Backlog"), ("doing", "Doing"), ("done", "Done")];
 const ACKNOWLEDGEMENT_STREAM_BUFFER_LIMIT: usize = 64;
+const ORDINARY_HANDLER_TIMEOUT: Duration = Duration::from_secs(10);
+const STARTUP_REPLAY_TIMEOUT: Duration = Duration::from_secs(5);
+const MAX_SYNC_STORE_BYTES: usize = 1024 * 1024;
 const ACKNOWLEDGEMENT_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15);
 const ACKNOWLEDGEMENT_RECONNECT_BACKOFF: [Duration; 3] = [
     Duration::from_millis(100),
@@ -228,15 +233,26 @@ struct PersistedAcknowledgement {
 }
 
 impl SyncStore {
-    fn load(&self) -> Result<SyncState, String> {
+    async fn load(&self) -> Result<SyncState, String> {
         if !self.0.exists() {
             return Ok(SyncState {
                 next_sequence: 1,
                 ..SyncState::default()
             });
         }
-        let bytes =
-            fs::read(&self.0).map_err(|error| format!("read {}: {error}", self.0.display()))?;
+        let bytes = run_with_timeout(
+            STARTUP_REPLAY_TIMEOUT,
+            "startup sync-store read/replay",
+            tokio::fs::read(&self.0),
+        )
+        .await?
+        .map_err(|error| format!("read {}: {error}", self.0.display()))?;
+        if bytes.len() > MAX_SYNC_STORE_BYTES {
+            return Err(format!(
+                "sync store {} exceeds {MAX_SYNC_STORE_BYTES} bytes",
+                self.0.display()
+            ));
+        }
         let persisted: PersistedSync = serde_json::from_slice(&bytes)
             .map_err(|error| format!("decode {}: {error}", self.0.display()))?;
         if !matches!(persisted.schema_version, 1 | 2) || persisted.next_sequence == 0 {
@@ -408,20 +424,46 @@ struct Presence {
     count: u64,
 }
 
+async fn run_with_timeout<F, T>(
+    duration: Duration,
+    operation: &'static str,
+    future: F,
+) -> Result<T, String>
+where
+    F: Future<Output = T>,
+{
+    tokio::time::timeout(duration, future)
+        .await
+        .map_err(|_| format!("{operation} timed out after {} ms", duration.as_millis()))
+}
+
+async fn bounded_handler(duration: Duration, request: Request, next: Next) -> Response {
+    match run_with_timeout(duration, "ordinary request", next.run(request)).await {
+        Ok(response) => response,
+        Err(message) => (StatusCode::GATEWAY_TIMEOUT, message).into_response(),
+    }
+}
+
+async fn ordinary_handler_timeout(request: Request, next: Next) -> Response {
+    bounded_handler(ORDINARY_HANDLER_TIMEOUT, request, next).await
+}
+
 #[tokio::main]
 async fn main() {
     let sync_store = std::env::var_os("HEMX_KANBAN_SYNC_STORE")
         .map(PathBuf::from)
         .map(SyncStore);
-    let mut sync = sync_store
-        .as_ref()
-        .map(SyncStore::load)
-        .transpose()
-        .unwrap_or_else(|error| panic!("cannot start with sync store: {error}"))
-        .unwrap_or_else(|| SyncState {
+    let mut sync = if let Some(store) = sync_store.as_ref() {
+        store
+            .load()
+            .await
+            .unwrap_or_else(|error| panic!("cannot start with sync store: {error}"))
+    } else {
+        SyncState {
             next_sequence: 1,
             ..SyncState::default()
-        });
+        }
+    };
     sync.retained_after = std::env::var("HEMX_KANBAN_RETAINED_AFTER")
         .ok()
         .and_then(|value| value.parse::<u64>().ok())
@@ -453,16 +495,18 @@ async fn main() {
             .unwrap_or(ACKNOWLEDGEMENT_HEARTBEAT_INTERVAL),
     });
 
-    let app = Router::new()
+    let ordinary_routes = Router::new()
         .route("/", get(home).post(interact))
         .route("/events", get(events))
         .route("/sync-demo", get(sync_demo))
         .route("/sync.js", get(sync_js))
         .route("/sync/context", get(sync_context))
         .route("/sync/commands", post(sync_command))
-        .route("/sync/acknowledgements", get(sync_acknowledgements))
         .route("/sync/snapshot", get(sync_snapshot))
         .route(runtime_js_path(), get(runtime))
+        .layer(middleware::from_fn(ordinary_handler_timeout));
+    let app = ordinary_routes
+        .merge(Router::new().route("/sync/acknowledgements", get(sync_acknowledgements)))
         .with_state(state);
 
     let addr = std::env::var("HEMX_KANBAN_ADDR")
@@ -1067,6 +1111,77 @@ mod tests {
         select_options_selector, small_text_selector, strong_text_selector,
     };
     use scraper::{Html, Selector};
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use tower::ServiceExt;
+
+    struct CancelProof(Arc<AtomicBool>);
+
+    impl Drop for CancelProof {
+        fn drop(&mut self) {
+            self.0.store(true, Ordering::SeqCst);
+        }
+    }
+
+    // req: operations/003
+    #[tokio::test]
+    async fn ordinary_handlers_timeout_and_cancel_inflight_work() {
+        let cancelled = Arc::new(AtomicBool::new(false));
+        let proof = Arc::clone(&cancelled);
+        let app = Router::new()
+            .route(
+                "/slow",
+                get(move || {
+                    let proof = Arc::clone(&proof);
+                    async move {
+                        let _cancel_proof = CancelProof(proof);
+                        std::future::pending::<Response>().await
+                    }
+                }),
+            )
+            .layer(middleware::from_fn(|request, next| async move {
+                bounded_handler(Duration::from_millis(20), request, next).await
+            }));
+        let response = app
+            .oneshot(
+                Request::get("/slow")
+                    .body(axum::body::Body::empty())
+                    .unwrap(),
+            )
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::GATEWAY_TIMEOUT);
+        assert!(cancelled.load(Ordering::SeqCst));
+    }
+
+    // req: operations/003
+    #[tokio::test]
+    async fn startup_replay_timeout_cancels_inflight_work_with_a_named_error() {
+        let cancelled = Arc::new(AtomicBool::new(false));
+        let proof = Arc::clone(&cancelled);
+        let error = run_with_timeout(Duration::from_millis(20), "startup replay", async move {
+            let _cancel_proof = CancelProof(proof);
+            std::future::pending::<()>().await;
+        })
+        .await
+        .expect_err("startup replay must time out");
+        assert_eq!(error, "startup replay timed out after 20 ms");
+        assert!(cancelled.load(Ordering::SeqCst));
+    }
+
+    #[tokio::test]
+    async fn startup_replay_rejects_oversized_store_before_decoding() {
+        let path = std::env::temp_dir().join(format!(
+            "hemx-kanban-oversized-store-{}.json",
+            std::process::id()
+        ));
+        fs::write(&path, vec![b' '; MAX_SYNC_STORE_BYTES + 1]).unwrap();
+        let error = match SyncStore(path.clone()).load().await {
+            Ok(_) => panic!("oversized store must be rejected"),
+            Err(error) => error,
+        };
+        let _ = fs::remove_file(path);
+        assert!(error.contains("exceeds 1048576 bytes"), "{error}");
+    }
 
     fn selector(value: &str) -> Selector {
         Selector::parse(value).expect("test selector parses")