fix(script): bound and cancel socket writes (#1248)

Cargo.lock

@@ -1591,6 +1591,7 @@ dependencies = [
  "effectcraft-render",
  "effectcraft-time",
  "serde_json",
+ "socket2",
  "web-time",
 ]
 
@@ -4815,6 +4816,16 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "socket2"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4"
+dependencies = [
+ "libc",
+ "windows-sys 0.61.2",
+]
+
 [[package]]
 name = "spin"
 version = "0.9.9"

crates/script/Cargo.toml

@@ -27,3 +27,6 @@ boa_engine = { workspace = true, features = ["js"] }
 [dev-dependencies]
 effectcraft-expr = { workspace = true }
 effectcraft-psd = { workspace = true }
+
+[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies]
+socket2 = "0.6"

crates/script/src/socket.rs

@@ -51,6 +51,7 @@ mod imp {
     use std::collections::BTreeMap;
     use std::io::{ErrorKind, Read, Write};
     use std::net::{Shutdown, TcpListener, TcpStream, ToSocketAddrs};
+    use std::sync::atomic::{AtomicBool, Ordering};
     use std::sync::{Arc, Mutex};
     use std::time::{Duration, Instant};
 
@@ -65,9 +66,10 @@ mod imp {
     /// socket blocked on a slow peer doesn't hold up the others (#1248).
     struct Entry {
         sock: Mutex<Sock>,
-        /// A second handle on a stream, so `close` can shut it down while another call is
-        /// blocked reading or writing it.
+        /// A second handle on a stream, so `close` can wake a blocked read.
         wake: Option<TcpStream>,
+        /// Writes poll this flag because shutdown alone may not interrupt a Windows send.
+        closed: AtomicBool,
     }
 
     struct Table {
@@ -85,7 +87,7 @@ mod imp {
             Sock::Stream { s, .. } => s.try_clone().ok(),
             Sock::Listener(_) => None,
         };
-        let entry = Arc::new(Entry { sock: Mutex::new(s), wake });
+        let entry = Arc::new(Entry { sock: Mutex::new(s), wake, closed: AtomicBool::new(false) });
         let mut t = table().lock().unwrap_or_else(|e| e.into_inner());
         let id = t.next;
         t.next += 1;
@@ -101,11 +103,23 @@ mod imp {
     /// Forget socket `id`; a call blocked on it fails at once.
     fn close(id: u32) {
         let entry = table().lock().unwrap_or_else(|e| e.into_inner()).socks.remove(&id);
-        if let Some(w) = entry.as_ref().and_then(|e| e.wake.as_ref()) {
-            let _ = w.shutdown(Shutdown::Both);
+        if let Some(entry) = entry {
+            entry.closed.store(true, Ordering::Release);
+            if let Some(w) = &entry.wake {
+                let _ = w.shutdown(Shutdown::Both);
+            }
         }
     }
 
+    /// A test-only handle for establishing backpressure before exercising the public operations.
+    #[cfg(test)]
+    pub(super) fn test_stream(id: u32) -> TcpStream {
+        let entry = get(id).unwrap();
+        let sock = entry.sock.lock().unwrap();
+        let Sock::Stream { s, .. } = &*sock else { panic!("expected a stream") };
+        s.try_clone().unwrap()
+    }
+
     fn timeout(p: &J) -> Duration {
         Duration::from_secs_f64(p.get("timeout").and_then(J::as_f64).unwrap_or(10.0).clamp(0.0, 3600.0))
     }
@@ -150,6 +164,36 @@ mod imp {
         Ok(())
     }
 
+    /// Write with one deadline and an explicit close signal. On Windows, shutdown alone does
+    /// not reliably interrupt a blocking send, so never leave a send blocked inside the OS.
+    /// The caller holds this socket's lock; restore blocking mode before another operation.
+    fn write(s: &mut TcpStream, data: &[u8], limit: Duration, closed: &AtomicBool) -> std::io::Result<()> {
+        s.set_nonblocking(true)?;
+        let end = Instant::now() + limit;
+        let result = (|| {
+            let mut pending = data;
+            while !pending.is_empty() {
+                if closed.load(Ordering::Acquire) {
+                    return Err(std::io::Error::new(ErrorKind::ConnectionAborted, "the socket was closed"));
+                }
+                let left = end.saturating_duration_since(Instant::now());
+                if left.is_zero() {
+                    return Err(std::io::Error::from(ErrorKind::TimedOut));
+                }
+                match s.write(pending) {
+                    Ok(0) => return Err(std::io::Error::from(ErrorKind::WriteZero)),
+                    Ok(n) => pending = pending.get(n..).ok_or_else(|| std::io::Error::from(ErrorKind::InvalidData))?,
+                    Err(e) if e.kind() == ErrorKind::WouldBlock => std::thread::sleep(left.min(Duration::from_millis(5))),
+                    Err(e) if e.kind() == ErrorKind::Interrupted => continue,
+                    Err(e) => return Err(e),
+                }
+            }
+            Ok(())
+        })();
+        let restored = s.set_nonblocking(false);
+        result.and(restored)
+    }
+
     pub(super) fn run(op: &str, id: u32, p: &J) -> J {
         let utf8 = p.get("encoding").and_then(J::as_str).is_some_and(|e| e.eq_ignore_ascii_case("utf-8") || e.eq_ignore_ascii_case("utf8"));
         match op {
@@ -265,11 +309,8 @@ mod imp {
                         // A peer that stops reading fills the buffers: give up after `timeout`
                         // instead of blocking the script for good (#1248).
                         let limit = timeout(p).max(Duration::from_millis(1));
-                        if let Err(e) = s.set_write_timeout(Some(limit)) {
-                            return err(format!("cannot set the write timeout: {e}"));
-                        }
                         let data = p.get("data").and_then(J::as_str).unwrap_or("");
-                        match s.write_all(&encode(data, utf8)).and_then(|_| s.flush()) {
+                        match write(s, &encode(data, utf8), limit, &entry.closed) {
                             Ok(()) => json!({"ok": true}),
                             Err(e) if matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => err(format!(
                                 "write timed out after {} s: the other side isn't reading (some of the data may have been sent)",
@@ -289,6 +330,7 @@ mod imp {
 #[cfg(test)]
 #[cfg(not(target_arch = "wasm32"))]
 mod tests {
+    use std::io::{ErrorKind, Read, Write};
     use std::net::{TcpListener, TcpStream};
     use std::sync::mpsc;
     use std::time::{Duration, Instant};
@@ -301,6 +343,7 @@ mod tests {
     /// `host:port` and the accepted stream's keeper (drop it to let the connection go).
     fn stalled_peer() -> (String, mpsc::Sender<()>) {
         let l = TcpListener::bind("127.0.0.1:0").unwrap();
+        socket2::SockRef::from(&l).set_recv_buffer_size(4096).unwrap();
         let host = l.local_addr().unwrap().to_string();
         let (stop, wait) = mpsc::channel::<()>();
         std::thread::spawn(move || {
@@ -315,9 +358,26 @@ mod tests {
         r["id"].as_u64().unwrap_or_else(|| panic!("open {host}: {r}")) as u32
     }
 
-    /// More than the loopback buffers hold, so a peer that doesn't read stops the write.
-    fn big() -> String {
-        "x".repeat(64 * 1024 * 1024)
+    /// Fill bounded socket buffers nonblocking, rather than assuming a 64 MiB write stalls.
+    /// In particular, Windows may accept that entire write with its default send buffering.
+    fn backpressure(id: u32) {
+        let mut stream = super::imp::test_stream(id);
+        socket2::SockRef::from(&stream).set_send_buffer_size(4096).unwrap();
+        stream.set_nonblocking(true).unwrap();
+        let start = Instant::now();
+        let chunk = [b'x'; 4096];
+        let mut sent = 0;
+        loop {
+            assert!(sent < 16 * 1024 * 1024 && start.elapsed() < Duration::from_secs(5), "could not establish backpressure");
+            match stream.write(&chunk) {
+                Ok(0) => panic!("peer closed while establishing backpressure"),
+                Ok(n) => sent += n,
+                Err(e) if e.kind() == ErrorKind::WouldBlock => break,
+                Err(e) if e.kind() == ErrorKind::Interrupted => continue,
+                Err(e) => panic!("establishing backpressure: {e}"),
+            }
+        }
+        stream.set_nonblocking(false).unwrap();
     }
 
     /// Run `f` on a thread; its result, or `None` if it takes longer than `limit`.
@@ -334,11 +394,12 @@ mod tests {
         // #1248: it used to block for good.
         let (host, _keep) = stalled_peer();
         let id = open(&host);
-        let data = big();
+        backpressure(id);
+        let data = "x".repeat(64 * 1024);
         let t0 = Instant::now();
         let r = within(Duration::from_secs(20), move || run("write", id, &json!({"data": data, "timeout": 1}))).expect("the write gave up");
         assert!(r["error"].as_str().is_some_and(|e| e.contains("timed out")), "{r}");
-        assert!(t0.elapsed() < Duration::from_secs(15), "{:?}", t0.elapsed());
+        assert!(t0.elapsed() >= Duration::from_millis(500) && t0.elapsed() < Duration::from_secs(15), "{:?}", t0.elapsed());
         run("close", id, &J::Null);
     }
 
@@ -348,12 +409,14 @@ mod tests {
         let (host_b, _keep_b) = stalled_peer();
         let (a, b) = (open(&host_a), open(&host_b));
         // A write on A that would wait a minute.
-        let data = big();
+        backpressure(a);
+        let data = "x".repeat(64 * 1024);
         let (tx, blocked) = mpsc::channel();
         std::thread::spawn(move || {
             let _ = tx.send(run("write", a, &json!({"data": data, "timeout": 60})));
         });
         std::thread::sleep(Duration::from_millis(500));
+        assert!(matches!(blocked.try_recv(), Err(mpsc::TryRecvError::Empty)), "A must still be writing before close");
         // Socket B answers meanwhile (the table isn't locked by A's write).
         let r = within(Duration::from_secs(3), move || run("connected", b, &json!({}))).expect("socket B answered while A was blocked");
         assert_eq!(r["value"], true, "{r}");
@@ -363,4 +426,46 @@ mod tests {
         assert!(r.get("error").is_some(), "{r}");
         run("close", b, &J::Null);
     }
+
+    /// Both outgoing and accepted sockets deliver complete data, then return to blocking reads.
+    #[test]
+    fn writes_deliver_all_data_and_restore_blocking_reads() {
+        for accepted in [false, true] {
+            let (id, mut peer) = if accepted {
+                let listener = run("listen", 0, &json!({"port": 0}));
+                let listener_id = listener["id"].as_u64().unwrap() as u32;
+                let peer = TcpStream::connect(("127.0.0.1", listener["port"].as_u64().unwrap() as u16)).unwrap();
+                let start = Instant::now();
+                let id = loop {
+                    if let Some(id) = run("poll", listener_id, &J::Null)["id"].as_u64() {
+                        break id as u32;
+                    }
+                    assert!(start.elapsed() < Duration::from_secs(5));
+                    std::thread::sleep(Duration::from_millis(1));
+                };
+                run("close", listener_id, &J::Null);
+                (id, peer)
+            } else {
+                let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+                let id = open(&listener.local_addr().unwrap().to_string());
+                (id, listener.accept().unwrap().0)
+            };
+            socket2::SockRef::from(&super::imp::test_stream(id)).set_send_buffer_size(4096).unwrap();
+            peer.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
+            let reader = std::thread::spawn(move || {
+                let mut data = vec![0; 256 * 1024];
+                peer.read_exact(&mut data).unwrap();
+                assert!(data.iter().all(|b| *b == b'x'));
+                // A read on a socket accidentally left nonblocking would return before this.
+                std::thread::sleep(Duration::from_millis(100));
+                peer.write_all(b"ack").unwrap();
+            });
+            let r = within(Duration::from_secs(10), move || run("write", id, &json!({"data": "x".repeat(256 * 1024), "timeout": 5}))).unwrap();
+            assert_eq!(r["ok"], true, "{r}");
+            let r = run("read", id, &json!({"count": 3, "timeout": 2}));
+            assert_eq!(r["data"], "ack", "{r}");
+            run("close", id, &J::Null);
+            reader.join().unwrap();
+        }
+    }
 }