WASM & Flake

wasm/src/client.rs

@@ -193,7 +193,6 @@ impl WasmClient {
         self.set_state(ConnectionState::Connecting);
         let transport =
             WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?;
-        let inner = transport.inner().clone();
 
         let version_str = format!("{}", PROTOCOL_VERSION);
         let ident = CommunicationValue::new(CommunicationType::Identification)
@@ -207,6 +206,7 @@ impl WasmClient {
             .map_err(|e| js_error(&format!("encode failed: {}", e)))?;
         transport.send_frame(&ident_bytes).await?;
 
+        let loop_transport = transport.clone();
         self.transport = Some(transport);
         self.set_state(ConnectionState::Connected);
 
@@ -214,9 +214,7 @@ impl WasmClient {
         let on_msg = self.on_message.clone();
         let on_err = self.on_error.clone();
         wasm_bindgen_futures::spawn_local(async move {
-            WasmTransport::from_inner(inner)
-                .receive_loop(on_msg, on_err)
-                .await;
+            loop_transport.receive_loop(on_msg, on_err).await;
             state.set(ConnectionState::Disconnected);
         });
         Ok(())
@@ -250,7 +248,6 @@ impl WasmClient {
 
         let transport =
             WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?;
-        let inner = transport.inner().clone();
 
         // 1. Send the unsigned Identification hello.
         let hello = CommunicationValue::new(CommunicationType::Identification)
@@ -360,6 +357,7 @@ impl WasmClient {
             }
         };
 
+        let loop_transport = transport.clone();
         self.transport = Some(transport);
         self.set_state(ConnectionState::Connected);
 
@@ -367,9 +365,7 @@ impl WasmClient {
         let on_msg = self.on_message.clone();
         let on_err = self.on_error.clone();
         wasm_bindgen_futures::spawn_local(async move {
-            WasmTransport::from_inner(inner)
-                .receive_loop(on_msg, on_err)
-                .await;
+            loop_transport.receive_loop(on_msg, on_err).await;
             state.set(ConnectionState::Disconnected);
         });
 
@@ -403,7 +399,6 @@ impl WasmClient {
 
         let transport =
             WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?;
-        let inner = transport.inner().clone();
 
         // 1. Send the unsigned Register hello (version + public-key bundle).
         let hello = CommunicationValue::new(CommunicationType::Register)
@@ -510,6 +505,7 @@ impl WasmClient {
             return Err(e);
         }
 
+        let loop_transport = transport.clone();
         self.transport = Some(transport);
         self.set_state(ConnectionState::Connected);
 
@@ -517,9 +513,7 @@ impl WasmClient {
         let on_msg = self.on_message.clone();
         let on_err = self.on_error.clone();
         wasm_bindgen_futures::spawn_local(async move {
-            WasmTransport::from_inner(inner)
-                .receive_loop(on_msg, on_err)
-                .await;
+            loop_transport.receive_loop(on_msg, on_err).await;
             state.set(ConnectionState::Disconnected);
         });
 

wasm/src/transport.rs

@@ -1,3 +1,6 @@
+use std::cell::RefCell;
+use std::rc::Rc;
+
 use wasm_bindgen::JsCast;
 use wasm_bindgen::prelude::*;
 use wasm_bindgen_futures::JsFuture;
@@ -27,9 +30,37 @@ fn resolve_stream_readable(recv_stream: &JsValue) -> Result<JsValue, JsValue> {
     }
 }
 
+/// Outcome of reading the next framed message from the incoming stream(s).
+enum FrameOutcome {
+    /// A complete application frame.
+    Frame(Vec<u8>),
+    /// The peer sent an explicit close frame (length == `u32::MAX`).
+    Closed,
+    /// The incoming-streams readable ended (transport gone), no more frames.
+    Ended,
+}
+
+/*
+ * WebTransport client transport.
+ *
+ * The native host sends with a *persistent* uni-directional stream: the auth
+ * `Challenge` and the final `IdentificationResponse`/`RegisterResponse` arrive
+ * as two length-prefixed frames on the *same* QUIC stream, and later
+ * application messages may arrive on subsequent streams. The reader state
+ * (`streams_reader`, `stream_reader`, `buffer`) is therefore shared via `Rc`
+ * between the handshake (`read_one_frame`) and the background `receive_loop`,
+ * so frames are never lost across the boundary and multiple frames can be read
+ * from one stream.
+ */
 #[derive(Clone)]
 pub struct WasmTransport {
     inner: WebTransport,
+    /// Reader over `incoming_unidirectional_streams()` (a singleton stream of streams).
+    streams_reader: Rc<RefCell<Option<JsValue>>>,
+    /// Reader over the host's current uni-directional stream, if one is open.
+    stream_reader: Rc<RefCell<Option<JsValue>>>,
+    /// Bytes already read from the current stream but not yet consumed as a frame.
+    buffer: Rc<RefCell<Vec<u8>>>,
 }
 
 impl WasmTransport {
@@ -58,17 +89,18 @@ impl WasmTransport {
         JsFuture::from(transport.ready())
             .await
             .map_err(|e| js_error(&format!("WebTransport ready failed: {:?}", e)))?;
-        Ok(Self { inner: transport })
+        Ok(Self {
+            inner: transport,
+            streams_reader: Rc::new(RefCell::new(None)),
+            stream_reader: Rc::new(RefCell::new(None)),
+            buffer: Rc::new(RefCell::new(Vec::new())),
+        })
     }
 
     pub fn inner(&self) -> &WebTransport {
         &self.inner
     }
 
-    pub fn from_inner(inner: WebTransport) -> Self {
-        Self { inner }
-    }
-
     pub async fn send_frame(&self, frame: &[u8]) -> Result<(), JsValue> {
         let stream_promise = self.inner.create_unidirectional_stream();
         let stream = JsFuture::from(stream_promise).await?;
@@ -110,247 +142,178 @@ impl WasmTransport {
         Ok(())
     }
 
-    /// Read exactly one frame from incoming uni streams, then release the reader
-    /// so `receive_loop` can pick up from where we left off.
-    pub async fn read_one_frame(&self) -> Result<Vec<u8>, JsValue> {
+    /// Get (creating once) the reader over `incoming_unidirectional_streams()`.
+    fn ensure_streams_reader(&self) -> Result<JsValue, JsValue> {
+        if let Some(reader) = self.streams_reader.borrow().clone() {
+            return Ok(reader);
+        }
         let incoming = self.inner.incoming_unidirectional_streams();
-
-        let reader_fn = js_sys::Reflect::get(&incoming, &JsValue::from_str("getReader"))
+        let reader = js_sys::Reflect::get(&incoming, &JsValue::from_str("getReader"))
             .map_err(|_| js_error("missing getReader"))?
             .dyn_into::<js_sys::Function>()
-            .map_err(|_| js_error("getReader not a function"))?;
-        let reader_val = reader_fn
+            .map_err(|_| js_error("getReader not a function"))?
             .call0(&incoming)
             .map_err(|_| js_error("getReader call failed"))?;
+        *self.streams_reader.borrow_mut() = Some(reader.clone());
+        Ok(reader)
+    }
 
-        let read_fn = js_sys::Reflect::get(&reader_val, &JsValue::from_str("read"))
+    /// Accept the next incoming uni-directional stream and make it current.
+    /// Returns `false` if the incoming-streams readable has ended.
+    async fn open_next_stream(&self) -> Result<bool, JsValue> {
+        let streams_reader = self.ensure_streams_reader()?;
+
+        let read_fn = js_sys::Reflect::get(&streams_reader, &JsValue::from_str("read"))
             .map_err(|_| js_error("missing read"))?
             .dyn_into::<js_sys::Function>()
             .map_err(|_| js_error("read not a function"))?;
-        let result_promise = read_fn
-            .call0(&reader_val)
-            .map_err(|_| js_error("read call failed"))?;
-        let result = JsFuture::from(result_promise.unchecked_into::<js_sys::Promise>())
-            .await
-            .map_err(|e| js_error(&format!("read failed: {:?}", e)))?;
-
-        // Release the reader lock so receive_loop can create its own reader
-        if let Some(release_fn) =
-            js_sys::Reflect::get(&reader_val, &JsValue::from_str("releaseLock"))
-                .ok()
-                .and_then(|f| f.dyn_into::<js_sys::Function>().ok())
-        {
-            let _ = release_fn.call0(&reader_val);
-        }
+        let result = JsFuture::from(
+            read_fn
+                .call0(&streams_reader)
+                .map_err(|_| js_error("read call failed"))?
+                .unchecked_into::<js_sys::Promise>(),
+        )
+        .await
+        .map_err(|e| js_error(&format!("accept stream failed: {:?}", e)))?;
 
         let done = js_sys::Reflect::get(&result, &JsValue::from_str("done"))
             .ok()
             .and_then(|v| v.as_bool())
             .unwrap_or(false);
         if done {
-            return Err(js_error("stream ended before frame"));
+            return Ok(false);
         }
 
         let recv_stream = js_sys::Reflect::get(&result, &JsValue::from_str("value"))
             .map_err(|_| js_error("missing value"))?;
-
-        let readable_or_stream = resolve_stream_readable(&recv_stream)?;
-        let stream_reader_fn =
-            js_sys::Reflect::get(&readable_or_stream, &JsValue::from_str("getReader"))
-                .map_err(|_| js_error("missing stream getReader"))?
-                .dyn_into::<js_sys::Function>()
-                .map_err(|_| js_error("stream getReader not a function"))?;
-        let stream_reader = stream_reader_fn
-            .call0(&readable_or_stream)
+        let readable = resolve_stream_readable(&recv_stream)?;
+        let reader = js_sys::Reflect::get(&readable, &JsValue::from_str("getReader"))
+            .map_err(|_| js_error("missing stream getReader"))?
+            .dyn_into::<js_sys::Function>()
+            .map_err(|_| js_error("stream getReader not a function"))?
+            .call0(&readable)
             .map_err(|_| js_error("stream getReader call failed"))?;
 
-        let mut buffer: Vec<u8> = Vec::new();
+        *self.stream_reader.borrow_mut() = Some(reader);
+        Ok(true)
+    }
+
+    /// Read one chunk from the current stream. `Ok(None)` means the stream ended.
+    async fn read_chunk(&self) -> Result<Option<Vec<u8>>, JsValue> {
+        let reader = match self.stream_reader.borrow().clone() {
+            Some(r) => r,
+            None => return Ok(None),
+        };
+
+        let read_fn = js_sys::Reflect::get(&reader, &JsValue::from_str("read"))
+            .map_err(|_| js_error("missing read"))?
+            .dyn_into::<js_sys::Function>()
+            .map_err(|_| js_error("read not a function"))?;
+        let result = JsFuture::from(
+            read_fn
+                .call0(&reader)
+                .map_err(|_| js_error("read call failed"))?
+                .unchecked_into::<js_sys::Promise>(),
+        )
+        .await
+        .map_err(|e| js_error(&format!("read failed: {:?}", e)))?;
+
+        let done = js_sys::Reflect::get(&result, &JsValue::from_str("done"))
+            .ok()
+            .and_then(|v| v.as_bool())
+            .unwrap_or(true);
+        if done {
+            return Ok(None);
+        }
+
+        let value = js_sys::Reflect::get(&result, &JsValue::from_str("value"))
+            .map_err(|_| js_error("missing value"))?;
+        Ok(Some(js_sys::Uint8Array::new(&value).to_vec()))
+    }
+
+    /// Try to pull one complete frame out of the buffer without reading more.
+    fn parse_buffer(&self) -> Result<Option<FrameOutcome>, JsValue> {
+        let buf = self.buffer.borrow();
+        if buf.len() < 4 {
+            return Ok(None);
+        }
+        let frame_len = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
+        if frame_len == CLOSE_FRAME_LEN {
+            return Ok(Some(FrameOutcome::Closed));
+        }
+        let frame_len = frame_len as usize;
+        let Some(frame_end) = 4usize.checked_add(frame_len) else {
+            return Err(js_error("invalid frame length"));
+        };
+        if frame_end > buf.len() {
+            return Ok(None);
+        }
+        let frame = buf[4..frame_end].to_vec();
+        drop(buf);
+        self.buffer.borrow_mut().drain(..frame_end);
+        Ok(Some(FrameOutcome::Frame(frame)))
+    }
+
+    /*
+     * Read the next framed message from the host. Frames are length-prefixed
+     * (u32 big-endian) and may be packed several-per-stream (the host reuses a
+     * persistent uni stream) or one-per-stream; both are handled by buffering
+     * across reads and advancing to the next stream when the current one ends.
+     */
+    async fn next_frame(&self) -> Result<FrameOutcome, JsValue> {
         loop {
-            let stream_read_fn =
-                match js_sys::Reflect::get(&stream_reader, &JsValue::from_str("read"))
-                    .ok()
-                    .and_then(|f| f.dyn_into::<js_sys::Function>().ok())
-                {
-                    Some(f) => f,
-                    None => break,
-                };
-            let chunk_promise = match stream_read_fn.call0(&stream_reader) {
-                Ok(p) => p,
-                Err(_) => break,
-            };
-            let chunk_result =
-                match JsFuture::from(chunk_promise.unchecked_into::<js_sys::Promise>()).await {
-                    Ok(v) => v,
-                    Err(_) => break,
-                };
-
-            let chunk_done = js_sys::Reflect::get(&chunk_result, &JsValue::from_str("done"))
-                .ok()
-                .and_then(|v| v.as_bool())
-                .unwrap_or(true);
-            if chunk_done {
-                return Err(js_error("stream closed before complete frame"));
+            if let Some(outcome) = self.parse_buffer()? {
+                return Ok(outcome);
             }
 
-            if let Ok(chunk_val) = js_sys::Reflect::get(&chunk_result, &JsValue::from_str("value"))
-            {
-                let arr = js_sys::Uint8Array::new(&chunk_val).to_vec();
-                if !arr.is_empty() {
-                    buffer.extend_from_slice(&arr);
-                }
+            let have_stream = self.stream_reader.borrow().is_some();
+            if !have_stream && !self.open_next_stream().await? {
+                return Ok(FrameOutcome::Ended);
             }
 
-            if buffer.len() >= 4 {
-                let frame_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
-                if frame_len == CLOSE_FRAME_LEN {
-                    return Err(js_error("connection closed before frame"));
+            match self.read_chunk().await? {
+                Some(chunk) => {
+                    if !chunk.is_empty() {
+                        self.buffer.borrow_mut().extend_from_slice(&chunk);
+                    }
                 }
-
-                let frame_len = frame_len as usize;
-                let Some(frame_end) = 4usize.checked_add(frame_len) else {
-                    return Err(js_error("invalid frame length"));
-                };
-                if frame_end <= buffer.len() {
-                    return Ok(buffer[4..frame_end].to_vec());
+                None => {
+                    // Current stream finished; the next frame (if any) is on a
+                    // subsequent stream. Any trailing partial bytes are dropped
+                    // since the host never splits a frame across streams.
+                    *self.stream_reader.borrow_mut() = None;
+                    self.buffer.borrow_mut().clear();
                 }
             }
         }
-
-        Err(js_error("stream ended before frame complete"))
     }
 
+    /// Read exactly one application frame (used during the auth handshake).
+    pub async fn read_one_frame(&self) -> Result<Vec<u8>, JsValue> {
+        match self.next_frame().await? {
+            FrameOutcome::Frame(frame) => Ok(frame),
+            FrameOutcome::Closed => Err(js_error("connection closed before frame")),
+            FrameOutcome::Ended => Err(js_error("stream ended before frame")),
+        }
+    }
+
+    /// Background loop: deliver every incoming frame to `on_message` until the
+    /// connection closes. Shares reader state with `read_one_frame`, so frames
+    /// buffered during the handshake are not lost.
     pub async fn receive_loop(&self, on_message: js_sys::Function, on_error: js_sys::Function) {
-        let incoming = self.inner.incoming_unidirectional_streams();
-
-        let reader_fn = match js_sys::Reflect::get(&incoming, &JsValue::from_str("getReader")) {
-            Ok(f) => match f.dyn_into::<js_sys::Function>() {
-                Ok(f) => f,
-                Err(_) => return,
-            },
-            Err(_) => return,
-        };
-        let reader_val = match reader_fn.call0(&incoming) {
-            Ok(v) => v,
-            Err(_) => return,
-        };
-
         loop {
-            let read_fn = match js_sys::Reflect::get(&reader_val, &JsValue::from_str("read")) {
-                Ok(f) => match f.dyn_into::<js_sys::Function>() {
-                    Ok(f) => f,
-                    Err(_) => break,
-                },
-                Err(_) => break,
-            };
-            let result = match read_fn.call0(&reader_val) {
-                Ok(p) => match JsFuture::from(p.unchecked_into::<js_sys::Promise>()).await {
-                    Ok(v) => v,
-                    Err(_) => break,
-                },
-                Err(_) => break,
-            };
-
-            let done = js_sys::Reflect::get(&result, &JsValue::from_str("done"))
-                .ok()
-                .and_then(|v| v.as_bool())
-                .unwrap_or(false);
-            if done {
-                break;
-            }
-
-            let recv_stream = match js_sys::Reflect::get(&result, &JsValue::from_str("value")) {
-                Ok(v) => v,
-                Err(_) => continue,
-            };
-
-            let this = self.clone();
-            let on_msg = on_message.clone();
-            let on_err = on_error.clone();
-            wasm_bindgen_futures::spawn_local(async move {
-                let _ = this.handle_stream(recv_stream, on_msg, on_err).await;
-            });
-        }
-    }
-
-    async fn handle_stream(
-        &self,
-        recv_stream: JsValue,
-        on_message: js_sys::Function,
-        _on_error: js_sys::Function,
-    ) -> Result<(), JsValue> {
-        let readable_or_stream = resolve_stream_readable(&recv_stream)?;
-        let stream_reader_fn =
-            js_sys::Reflect::get(&readable_or_stream, &JsValue::from_str("getReader"))?
-                .dyn_into::<js_sys::Function>()?;
-        let stream_reader = stream_reader_fn.call0(&readable_or_stream)?;
-
-        let mut buffer: Vec<u8> = Vec::new();
-
-        loop {
-            let stream_read_fn = js_sys::Reflect::get(&stream_reader, &JsValue::from_str("read"))?
-                .dyn_into::<js_sys::Function>()?;
-            let chunk_promise = stream_read_fn.call0(&stream_reader)?;
-            let chunk_result =
-                JsFuture::from(chunk_promise.unchecked_into::<js_sys::Promise>()).await?;
-
-            let chunk_done = js_sys::Reflect::get(&chunk_result, &JsValue::from_str("done"))
-                .ok()
-                .and_then(|v| v.as_bool())
-                .unwrap_or(true);
-            if chunk_done {
-                break;
-            }
-
-            if let Ok(chunk_val) = js_sys::Reflect::get(&chunk_result, &JsValue::from_str("value"))
-            {
-                let arr = js_sys::Uint8Array::new(&chunk_val).to_vec();
-                if !arr.is_empty() {
-                    buffer.extend_from_slice(&arr);
+            match self.next_frame().await {
+                Ok(FrameOutcome::Frame(frame)) => {
+                    let arr = js_sys::Uint8Array::from(&frame[..]);
+                    let _ = on_message.call1(&JsValue::NULL, &arr);
                 }
-            }
-
-            // Extract all complete frames from the buffer
-            while buffer.len() >= 4 {
-                let frame_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
-                if frame_len == CLOSE_FRAME_LEN {
-                    return Ok(());
-                }
-
-                let frame_len = frame_len as usize;
-                let Some(frame_end) = 4usize.checked_add(frame_len) else {
-                    return Err(js_error("invalid frame length"));
-                };
-                if frame_end > buffer.len() {
+                Ok(FrameOutcome::Closed) | Ok(FrameOutcome::Ended) => break,
+                Err(e) => {
+                    let _ = on_error.call1(&JsValue::NULL, &e);
                     break;
                 }
-                let frame = buffer[4..frame_end].to_vec();
-                let arr = js_sys::Uint8Array::from(&frame[..]);
-                let _ = on_message.call1(&JsValue::NULL, &arr);
-                buffer.drain(..frame_end);
             }
         }
-
-        // Process any remaining complete frames after stream closes
-        while buffer.len() >= 4 {
-            let frame_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
-            if frame_len == CLOSE_FRAME_LEN {
-                return Ok(());
-            }
-
-            let frame_len = frame_len as usize;
-            let Some(frame_end) = 4usize.checked_add(frame_len) else {
-                return Err(js_error("invalid frame length"));
-            };
-            if frame_end > buffer.len() {
-                break;
-            }
-            let frame = buffer[4..frame_end].to_vec();
-            let arr = js_sys::Uint8Array::from(&frame[..]);
-            let _ = on_message.call1(&JsValue::NULL, &arr);
-            buffer.drain(..frame_end);
-        }
-
-        Ok(())
     }
 
     pub fn close(&self) {