Fix dioxus-fullstack::Streaming<T,E> only deserializing one client value (#4975)

* dioxus-fullstack::Streaming: decode multiple client values from Bytes

* small nits

* remove the bidirectional stream since it doesn't work in the browser

* final cleanups

* nits

---------

Co-authored-by: Olivier Kannengieser <[email protected]>
Co-authored-by: Jonathan Kelley <[email protected]>

Author: 3 people

examples/07-fullstack/streaming.rs

@@ -21,7 +21,6 @@ use dioxus::{
     fullstack::{JsonEncoding, Streaming, TextStream},
     prelude::*,
 };
-use futures::{StreamExt as _, TryStreamExt};
 
 fn main() {
     dioxus::launch(app)
@@ -30,8 +29,6 @@ fn main() {
 fn app() -> Element {
     let mut text_responses = use_signal(String::new);
     let mut json_responses = use_signal(Vec::new);
-    let mut echo_responses = use_signal(Vec::new);
-    let mut transform_responses = use_signal(Vec::new);
 
     let mut start_text_stream = use_action(move || async move {
         text_responses.clear();
@@ -56,64 +53,6 @@ fn app() -> Element {
         dioxus::Ok(())
     });
 
-    let mut continue_echo_stream = use_signal_sync(|| false);
-    let mut start_echo_stream = use_action(move || async move {
-        continue_echo_stream.set(true);
-        echo_responses.clear();
-        let stream = echo_stream(Streaming::new(futures::stream::unfold(
-            0,
-            move |index| async move {
-                if !continue_echo_stream() {
-                    return None;
-                }
-                tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
-                let dog = Dog {
-                    name: format!("Dog {}", index),
-                    age: (index % 10) as u8,
-                };
-                Some((dog, index + 1))
-            },
-        )))
-        .await?;
-        stream
-            .into_inner()
-            .try_for_each(move |dog| async move {
-                echo_responses.push(dog);
-                Ok(())
-            })
-            .await?;
-        dioxus::Ok(())
-    });
-
-    let mut continue_transform_stream = use_signal_sync(|| false);
-    let mut start_transform_stream = use_action(move || async move {
-        continue_transform_stream.set(true);
-        transform_responses.clear();
-        let stream = transform_stream(Streaming::new(futures::stream::unfold(
-            0,
-            move |index| async move {
-                if !continue_transform_stream() {
-                    return None;
-                }
-                tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
-                let dog = Dog {
-                    name: format!("Dog {}", index),
-                    age: (index % 10) as u8,
-                };
-                Some((dog, index + 1))
-            },
-        )))
-        .await?;
-        stream
-            .into_inner()
-            .try_for_each(move |text| async move {
-                transform_responses.push(text);
-                Ok(())
-            })
-            .await?;
-        dioxus::Ok(())
-    });
-
     rsx! {
         div {
             button { onclick: move |_| start_text_stream.call(), "Start text stream" }
@@ -127,20 +66,6 @@ fn app() -> Element {
                 pre { "{dog:?}" }
             }
         }
-        div {
-            button { onclick: move |_| start_echo_stream.call(), "Start echo stream" }
-            button { onclick: move |_| continue_echo_stream.set(false), "Stop echo stream" }
-            for dog in echo_responses.read().iter() {
-                pre { "{dog:?}" }
-            }
-        }
-        div {
-            button { onclick: move |_| start_transform_stream.call(), "Start transform stream" }
-            button { onclick: move |_| continue_transform_stream.set(false), "Stop transform stream" }
-            for text in transform_responses.read().iter() {
-                pre { "{text}" }
-            }
-        }
     }
 }
 
@@ -220,20 +145,3 @@ async fn byte_stream() -> Result<Streaming<Bytes>> {
 
     Ok(Streaming::new(rx))
 }
-
-/// An example of echoing the stream back to the client.
-#[post("/api/echo_stream")]
-async fn echo_stream(stream: Streaming<Dog, JsonEncoding>) -> Result<Streaming<Dog, JsonEncoding>> {
-    Ok(stream)
-}
-
-/// An example of transforming the stream on the server.
-#[post("/api/transform_stream")]
-async fn transform_stream(stream: Streaming<Dog, JsonEncoding>) -> Result<TextStream> {
-    Ok(Streaming::new(stream.into_inner().filter_map(
-        |dog| async {
-            dog.ok()
-                .map(|dog| format!("name: {}, age: {}", dog.name, dog.age))
-        },
-    )))
-}

packages/fullstack/src/encoding.rs

@@ -4,7 +4,7 @@ use serde::{de::DeserializeOwned, Serialize};
 /// A trait for encoding and decoding data.
 ///
 /// This takes an owned self to make it easier for zero-copy encodings.
-pub trait Encoding {
+pub trait Encoding: 'static {
     fn content_type() -> &'static str;
     fn stream_content_type() -> &'static str;
     fn to_bytes(data: impl Serialize) -> Option<Bytes> {

packages/fullstack/src/payloads/stream.rs

@@ -6,7 +6,7 @@ use crate::{
 };
 use axum::extract::{FromRequest, Request};
 use axum_core::response::IntoResponse;
-use bytes::Bytes;
+use bytes::{Buf as _, Bytes};
 use dioxus_fullstack_core::{HttpError, RequestError};
 use futures::{Stream, StreamExt};
 #[cfg(feature = "server")]
@@ -276,18 +276,8 @@ impl<T: DeserializeOwned + Serialize + 'static + Send, E: Encoding> FromResponse
 {
     fn from_response(res: ClientResponse) -> impl Future<Output = Result<Self, ServerFnError>> {
         SendWrapper::new(async move {
-            let client_stream = Box::pin(SendWrapper::new(res.bytes_stream().map(
-                |byte| match byte {
-                    Ok(bytes) => match decode_stream_frame::<T, E>(bytes) {
-                        Some(res) => Ok(res),
-                        None => Err(StreamingError::Decoding),
-                    },
-                    Err(_) => Err(StreamingError::Failed),
-                },
-            )));
-
             Ok(Self {
-                stream: client_stream,
+                stream: byte_stream_to_client_stream::<E, _, _, _>(res.bytes_stream()),
                 encoding: PhantomData,
             })
         })
@@ -385,13 +375,7 @@ impl<T: DeserializeOwned + Serialize + 'static + Send, E: Encoding, S> FromReque
             let stream = body.into_data_stream();
 
             Ok(Self {
-                stream: Box::pin(stream.map(|byte| match byte {
-                    Ok(bytes) => match decode_stream_frame::<T, E>(bytes) {
-                        Some(res) => Ok(res),
-                        None => Err(StreamingError::Decoding),
-                    },
-                    Err(_) => Err(StreamingError::Failed),
-                })),
+                stream: byte_stream_to_client_stream::<E, _, _, _>(stream),
                 encoding: PhantomData,
             })
         }
@@ -504,22 +488,98 @@ pub fn encode_stream_frame<T: Serialize, E: Encoding>(data: T) -> Option<Bytes>
     Some(Bytes::from(bytes).slice(offset..))
 }
 
+fn byte_stream_to_client_stream<E, T, S, E1>(
+    stream: S,
+) -> Pin<Box<dyn Stream<Item = Result<T, StreamingError>> + Send>>
+where
+    S: Stream<Item = Result<Bytes, E1>> + 'static + Send,
+    E: Encoding,
+    T: DeserializeOwned + 'static,
+{
+    Box::pin(stream.flat_map(|bytes| {
+        enum DecodeIteratorState {
+            Empty,
+            Failed,
+            Checked(Bytes),
+            UnChecked(Bytes),
+        }
+
+        let mut state = match bytes {
+            Ok(bytes) => DecodeIteratorState::UnChecked(bytes),
+            Err(_) => DecodeIteratorState::Failed,
+        };
+
+        futures::stream::iter(std::iter::from_fn(move || {
+            match std::mem::replace(&mut state, DecodeIteratorState::Empty) {
+                DecodeIteratorState::Empty => None,
+                DecodeIteratorState::Failed => Some(Err(StreamingError::Failed)),
+                DecodeIteratorState::Checked(mut bytes) => {
+                    let r = decode_stream_frame_multi::<T, E>(&mut bytes);
+                    if r.is_some() {
+                        state = DecodeIteratorState::Checked(bytes)
+                    }
+                    r
+                }
+                DecodeIteratorState::UnChecked(mut bytes) => {
+                    let r = decode_stream_frame_multi::<T, E>(&mut bytes);
+                    if r.is_some() {
+                        state = DecodeIteratorState::Checked(bytes);
+                        r
+                    } else {
+                        Some(Err(StreamingError::Decoding))
+                    }
+                }
+            }
+        }))
+    }))
+}
+
 /// Decode a websocket-framed streaming payload produced by [`encode_stream_frame`].
 ///
 /// This function returns `None` if the frame is invalid or cannot be decoded.
 ///
 /// It cannot handle masked frames, as those are not produced by our encoding function.
-pub fn decode_stream_frame<T, E>(frame: Bytes) -> Option<T>
+pub fn decode_stream_frame<T, E>(mut frame: Bytes) -> Option<T>
 where
     E: Encoding,
     T: DeserializeOwned,
 {
+    decode_stream_frame_multi::<T, E>(&mut frame).and_then(|r| r.ok())
+}
+
+/// Decode one value and advance the bytes pointer
+///
+/// If the frame is empty return None.
+///
+/// Otherwise, if the initial opcode is not the one expected for binary stream
+/// or the frame is not large enough return error StreamingError::Decoding
+fn decode_stream_frame_multi<T, E>(frame: &mut Bytes) -> Option<Result<T, StreamingError>>
+where
+    E: Encoding,
+    T: DeserializeOwned,
+{
+    let (offset, payload_len) = match offset_payload_len(frame)? {
+        Ok(r) => r,
+        Err(e) => return Some(Err(e)),
+    };
+
+    let r = E::decode(frame.slice(offset..offset + payload_len));
+    frame.advance(offset + payload_len);
+    r.map(|r| Ok(r))
+}
+
+/// Compute (offset,len) for decoding data
+fn offset_payload_len(frame: &Bytes) -> Option<Result<(usize, usize), StreamingError>> {
     let data = frame.as_ref();
 
-    if data.len() < 2 {
+    if data.is_empty() {
         return None;
     }
 
+    if data.len() < 2 {
+        return Some(Err(StreamingError::Decoding));
+    }
+
     let first = data[0];
     let second = data[1];
 
@@ -528,44 +588,43 @@ where
     let opcode = first & 0x0F;
     let rsv = first & 0x70;
     if !fin || opcode != 0x02 || rsv != 0 {
-        return None;
+        return Some(Err(StreamingError::Decoding));
     }
 
     // Mask bit must be zero for our framing
     if second & 0x80 != 0 {
-        return None;
+        return Some(Err(StreamingError::Decoding));
     }
 
     let mut offset = 2usize;
     let mut payload_len = (second & 0x7F) as usize;
 
     if payload_len == 126 {
         if data.len() < offset + 2 {
-            return None;
+            return Some(Err(StreamingError::Decoding));
         }
 
         payload_len = u16::from_be_bytes([data[offset], data[offset + 1]]) as usize;
         offset += 2;
     } else if payload_len == 127 {
         if data.len() < offset + 8 {
-            return None;
+            return Some(Err(StreamingError::Decoding));
         }
 
         let mut len_bytes = [0u8; 8];
         len_bytes.copy_from_slice(&data[offset..offset + 8]);
         let len_u64 = u64::from_be_bytes(len_bytes);
 
         if len_u64 > usize::MAX as u64 {
-            return None;
+            return Some(Err(StreamingError::Decoding));
         }
 
         payload_len = len_u64 as usize;
         offset += 8;
     }
 
     if data.len() < offset + payload_len {
-        return None;
+        return Some(Err(StreamingError::Decoding));
     }
-
-    E::decode(frame.slice(offset..offset + payload_len))
+    Some(Ok((offset, payload_len)))
 }