Address the review findings from #1: four production bugs and five hardening fixes

crates/tinybrowser/src/capture/mod.rs

@@ -27,7 +27,7 @@ use store::OutputStore;
 pub(crate) async fn screenshot(
     session: &Session,
     request: &ScreenshotRequest,
-    store: &tokio::sync::Mutex<OutputStore>,
+    store: &std::sync::Arc<tokio::sync::Mutex<OutputStore>>,
 ) -> Result<OutputRef> {
     if let Some(quality) = request.quality
         && !(1..=100).contains(&quality)
@@ -92,6 +92,8 @@ pub(crate) async fn screenshot(
         .and_then(Value::as_str)
         .ok_or_else(|| Error::page("browser captured no image data".to_string()))?;
 
+    within_cap(encoded.len())?;
+
     let bytes = BASE64
         .decode(encoded)
         .map_err(|error| Error::page(format!("screenshot was not valid base64: {error}")))?;
@@ -129,6 +131,36 @@ fn pixels(dimension: f64) -> u32 {
     }
 }
 
+/// Refuses an image too large to hold, from the length of its encoding.
+///
+/// Checked before decoding, not after. [`store::OutputStore::insert`] rejects an
+/// image larger than it will hold, but by then the decode has already allocated
+/// it — and the CDP frame carrying it is unbounded by design, because a
+/// full-page capture legitimately exceeds any frame cap worth setting. So the
+/// first place the size is known is the length of the encoded text, and this is
+/// the first place it can be refused without paying for it.
+///
+/// Base64 carries three bytes in four, so the encoded length gives the decoded
+/// size to within the padding.
+///
+/// # Errors
+///
+/// [`Error::LimitExceeded`] when the image would exceed the module's cap.
+fn within_cap(encoded_len: usize) -> Result<()> {
+    let decoded_len = encoded_len / 4 * 3;
+
+    if decoded_len > store::MAX_OUTPUT_BYTES {
+        return Err(Error::LimitExceeded {
+            message: format!(
+                "screenshot of about {decoded_len} bytes exceeds the {} byte cap",
+                store::MAX_OUTPUT_BYTES
+            ),
+        });
+    }
+
+    Ok(())
+}
+
 /// The clip rectangle covering one element.
 async fn element_clip(session: &Session, node: i64) -> Result<Value> {
     let model = session

crates/tinybrowser/src/capture/store.rs

@@ -30,10 +30,16 @@ const MAX_OUTPUTS: usize = 16;
 ///
 /// A full-page capture of a long article at 2x lands in the low megabytes; this
 /// is several times that and still far below what would matter to a host.
-const MAX_OUTPUT_BYTES: usize = 64 * 1024 * 1024;
+pub(crate) const MAX_OUTPUT_BYTES: usize = 64 * 1024 * 1024;
 
 /// How long an uncollected output survives.
-const TTL: Duration = Duration::from_secs(300);
+pub(crate) const TTL: Duration = Duration::from_secs(300);
+
+/// How often the sweeper looks for outputs to drop.
+///
+/// A fraction of [`TTL`], so an abandoned output is released within a minute or
+/// so of expiring rather than at some unbounded later moment.
+pub(crate) const SWEEP_INTERVAL: Duration = Duration::from_secs(60);
 
 /// The most a single [`read`](OutputStore::read) will return.
 ///
@@ -173,9 +179,28 @@ impl OutputStore {
     }
 
     /// Drops everything past its time to live.
-    fn expire(&mut self) {
+    ///
+    /// Called from the operations *and* from a sweeper, because an expiry that
+    /// only runs when something else happens is not an expiry: a host that takes
+    /// sixteen large screenshots and then goes quiet would hold every byte of
+    /// them, in somebody else's process, until it happened to call again.
+    pub(crate) fn expire(&mut self) {
         let now = Instant::now();
         self.held
             .retain(|_, held| now.duration_since(held.stored) < TTL);
     }
 }
+
+#[cfg(test)]
+impl OutputStore {
+    /// Ages every held output by `elapsed`, as if that much time had passed.
+    ///
+    /// Expiry is the one behaviour here that is a function of the clock, and a
+    /// test that waited five real minutes to check it would never be run. Moving
+    /// the timestamps back instead keeps the assertion exact and instant.
+    pub(crate) fn age(&mut self, elapsed: Duration) {
+        for held in self.held.values_mut() {
+            held.stored -= elapsed;
+        }
+    }
+}

crates/tinybrowser/src/capture/test.rs

@@ -11,7 +11,8 @@ use base64::Engine as _;
 use base64::engine::general_purpose::STANDARD as BASE64;
 use tinybrowser_bus::OutputId;
 
-use super::store::OutputStore;
+use super::store::{MAX_OUTPUT_BYTES, OutputStore, TTL};
+use super::within_cap;
 use crate::error::Error;
 
 fn store_with(bytes: Vec<u8>) -> (OutputStore, OutputId) {
@@ -182,3 +183,63 @@ fn an_output_larger_than_the_cap_is_refused_rather_than_held() {
     assert!(matches!(error, Error::LimitExceeded { .. }), "{error}");
     assert_eq!(store.len(), 0);
 }
+
+#[test]
+fn an_output_expires_once_its_time_to_live_has_passed() {
+    let (mut store, id) = store_with(b"hello".to_vec());
+    store.age(TTL + std::time::Duration::from_secs(1));
+
+    let error = store.read(&id, 0, 1024).expect_err("refused");
+    assert!(matches!(error, Error::NoSuchOutput { .. }), "{error}");
+}
+
+#[test]
+fn an_output_within_its_time_to_live_survives() {
+    let (mut store, id) = store_with(b"hello".to_vec());
+    store.age(TTL.saturating_sub(std::time::Duration::from_secs(1)));
+
+    assert!(store.read(&id, 0, 1024).is_ok());
+}
+
+#[test]
+fn sweeping_releases_what_nothing_else_would_have() {
+    // The case the sweeper exists for: outputs are held, nobody calls again, and
+    // without an independent sweep the bytes stay resident in the host process
+    // long past the point they were promised to be gone.
+    let mut store = OutputStore::default();
+    for _ in 0..4 {
+        store
+            .insert(vec![0; 1024], "image/png", 1, 1)
+            .expect("within the cap");
+    }
+    assert_eq!(store.len(), 4);
+
+    store.age(TTL + std::time::Duration::from_secs(1));
+    store.expire();
+
+    assert_eq!(store.len(), 0);
+}
+
+#[test]
+fn an_ordinary_screenshot_passes_the_pre_decode_check() {
+    // A full-page capture at 2x is a few megabytes; nothing near the cap.
+    assert!(within_cap(4 * 1024 * 1024).is_ok());
+    assert!(within_cap(0).is_ok());
+}
+
+#[test]
+fn an_oversized_screenshot_is_refused_before_it_is_decoded() {
+    // The point is the ordering: `OutputStore::insert` would refuse this too,
+    // but only after the decode had already allocated it in the host's process.
+    let encoded_len = MAX_OUTPUT_BYTES / 3 * 4 + 8;
+    let error = within_cap(encoded_len).expect_err("refused");
+
+    assert!(matches!(error, Error::LimitExceeded { .. }), "{error}");
+}
+
+#[test]
+fn the_pre_decode_check_agrees_with_the_store_it_is_guarding() {
+    // An encoding that decodes to exactly the cap must pass, or the check would
+    // refuse images the store would happily have held.
+    assert!(within_cap(MAX_OUTPUT_BYTES / 3 * 4).is_ok());
+}

crates/tinybrowser/src/cdp/launch.rs

@@ -137,6 +137,14 @@ pub(crate) struct LaunchedBrowser {
     pub(crate) websocket_url: String,
     child: Child,
     profile: Option<PathBuf>,
+    /// Keeps reading the browser's stderr for as long as it runs.
+    ///
+    /// Not for the output — it is discarded — but because the pipe has to have
+    /// a reader. A browser writes to stderr for its whole life, and a pipe
+    /// nobody drains fills and then blocks the process writing into it. Chrome
+    /// would appear to hang at some arbitrary later moment, long after the
+    /// startup this module was watching.
+    drain: tokio::task::JoinHandle<()>,
 }
 
 impl LaunchedBrowser {
@@ -146,6 +154,7 @@ impl LaunchedBrowser {
     /// browser that has already exited or a directory already gone are both the
     /// outcome being asked for.
     pub(crate) async fn shutdown(mut self) {
+        self.drain.abort();
         let _ = self.child.kill().await;
         if let Some(profile) = self.profile.take() {
             let _ = tokio::fs::remove_dir_all(profile).await;
@@ -286,26 +295,31 @@ pub(crate) async fn launch_within(
         ));
     };
 
-    let websocket_url = match tokio::time::timeout(startup, read_websocket_url(stderr)).await {
-        Ok(Ok(url)) => url,
-        Ok(Err(error)) => {
-            let _ = child.kill().await;
-            return Err(error);
-        }
-        Err(_) => {
-            let _ = child.kill().await;
-            return Err(Error::browser_unavailable(format!(
-                "{} did not report a devtools url within {}s",
-                executable.display(),
-                startup.as_secs()
-            )));
-        }
-    };
+    let (websocket_url, remaining) =
+        match tokio::time::timeout(startup, read_websocket_url(stderr)).await {
+            Ok(Ok(found)) => found,
+            Ok(Err(error)) => {
+                let _ = child.kill().await;
+                return Err(error);
+            }
+            Err(_) => {
+                let _ = child.kill().await;
+                return Err(Error::browser_unavailable(format!(
+                    "{} did not report a devtools url within {}s",
+                    executable.display(),
+                    startup.as_secs()
+                )));
+            }
+        };
 
     Ok(LaunchedBrowser {
         websocket_url,
         child,
         profile: owned.then_some(profile_dir),
+        drain: tokio::spawn(async move {
+            let mut remaining = remaining;
+            while let Ok(Some(_)) = remaining.next_line().await {}
+        }),
     })
 }
 
@@ -314,7 +328,9 @@ pub(crate) async fn launch_within(
 /// When the browser dies instead, the banner is the only account of why, so the
 /// first few lines of it are kept and handed to [`diagnose`] rather than
 /// discarded in favour of "it did not start".
-async fn read_websocket_url(stderr: tokio::process::ChildStderr) -> Result<String> {
+type StderrLines = tokio::io::Lines<BufReader<tokio::process::ChildStderr>>;
+
+async fn read_websocket_url(stderr: tokio::process::ChildStderr) -> Result<(String, StderrLines)> {
     const MARKER: &str = "DevTools listening on ";
     /// Enough to hold the fatal line and its context, and few enough that a
     /// browser logging steadily cannot grow this without bound.
@@ -325,7 +341,9 @@ async fn read_websocket_url(stderr: tokio::process::ChildStderr) -> Result<Strin
 
     while let Ok(Some(line)) = lines.next_line().await {
         if let Some(url) = line.split_once(MARKER) {
-            return Ok(url.1.trim().to_string());
+            // The reader goes back to the caller rather than being dropped here:
+            // see the note on `LaunchedBrowser::drain`.
+            return Ok((url.1.trim().to_string(), lines));
         }
         if banner.len() < KEPT_LINES {
             banner.push(line);