wasm_streams/readable/
into_underlying_byte_source.rs

1use std::cell::RefCell;
2use std::pin::Pin;
3use std::rc::Rc;
4
5use futures_util::future::{abortable, AbortHandle, TryFutureExt};
6use futures_util::io::{AsyncRead, AsyncReadExt};
7use js_sys::{Error as JsError, Promise, Uint8Array};
8use wasm_bindgen::prelude::*;
9use wasm_bindgen_futures::future_to_promise;
10
11use crate::util::{checked_cast_to_u32, clamp_to_usize};
12
13use super::sys;
14
15#[wasm_bindgen]
16pub(crate) struct IntoUnderlyingByteSource {
17    inner: Rc<RefCell<Inner>>,
18    default_buffer_len: usize,
19    controller: Option<sys::ReadableByteStreamController>,
20    pull_handle: Option<AbortHandle>,
21}
22
23impl IntoUnderlyingByteSource {
24    pub fn new(async_read: Box<dyn AsyncRead>, default_buffer_len: usize) -> Self {
25        IntoUnderlyingByteSource {
26            inner: Rc::new(RefCell::new(Inner::new(async_read))),
27            default_buffer_len,
28            controller: None,
29            pull_handle: None,
30        }
31    }
32}
33
34#[allow(clippy::await_holding_refcell_ref)]
35#[wasm_bindgen]
36impl IntoUnderlyingByteSource {
37    #[wasm_bindgen(getter, js_name = type)]
38    pub fn type_(&self) -> sys::ReadableStreamType {
39        sys::ReadableStreamType::Bytes
40    }
41
42    #[wasm_bindgen(getter, js_name = autoAllocateChunkSize)]
43    pub fn auto_allocate_chunk_size(&self) -> usize {
44        self.default_buffer_len
45    }
46
47    pub fn start(&mut self, controller: sys::ReadableByteStreamController) {
48        self.controller = Some(controller);
49    }
50
51    pub fn pull(&mut self, controller: sys::ReadableByteStreamController) -> Promise {
52        let inner = self.inner.clone();
53        let fut = async move {
54            // This mutable borrow can never panic, since the ReadableStream always queues
55            // each operation on the underlying source.
56            let mut inner = inner.try_borrow_mut().unwrap_throw();
57            inner.pull(controller).await
58        };
59
60        // Allow aborting the future from cancel().
61        let (fut, handle) = abortable(fut);
62        // Ignore errors from aborting the future.
63        let fut = fut.unwrap_or_else(|_| Ok(JsValue::undefined()));
64
65        self.pull_handle = Some(handle);
66        future_to_promise(fut)
67    }
68
69    pub fn cancel(self) {
70        // The stream has been canceled, drop everything.
71        drop(self);
72    }
73}
74
75impl Drop for IntoUnderlyingByteSource {
76    fn drop(&mut self) {
77        // Abort the pending pull, if any.
78        if let Some(handle) = self.pull_handle.take() {
79            handle.abort();
80        }
81    }
82}
83
84struct Inner {
85    async_read: Option<Pin<Box<dyn AsyncRead>>>,
86    buffer: Vec<u8>,
87}
88
89impl Inner {
90    fn new(async_read: Box<dyn AsyncRead>) -> Self {
91        Inner {
92            async_read: Some(async_read.into()),
93            buffer: Vec::new(),
94        }
95    }
96
97    async fn pull(
98        &mut self,
99        controller: sys::ReadableByteStreamController,
100    ) -> Result<JsValue, JsValue> {
101        // The AsyncRead should still exist, since pull() will not be called again
102        // after the stream has closed or encountered an error.
103        let async_read = self.async_read.as_mut().unwrap_throw();
104        // We set autoAllocateChunkSize, so there should always be a BYOB request.
105        let request = controller.byob_request().unwrap_throw();
106        // Resize the buffer to fit the BYOB request.
107        let request_view = request.view().unwrap_throw().unchecked_into::<Uint8Array>();
108        let request_len = clamp_to_usize(request_view.byte_length());
109        if self.buffer.len() < request_len {
110            self.buffer.resize(request_len, 0);
111        }
112        match async_read.read(&mut self.buffer[0..request_len]).await {
113            Ok(0) => {
114                // The stream has closed, drop it.
115                self.discard();
116                controller.close()?;
117                request.respond_with_u32(0)?;
118            }
119            Ok(bytes_read) => {
120                // Copy read bytes from buffer to BYOB request view
121                debug_assert!(bytes_read <= request_len);
122                let bytes_read_u32 = checked_cast_to_u32(bytes_read);
123                let dest = Uint8Array::new_with_byte_offset_and_length(
124                    &request_view.buffer(),
125                    request_view.byte_offset(),
126                    bytes_read_u32,
127                );
128                dest.copy_from(&self.buffer[0..bytes_read]);
129                // Respond to BYOB request
130                request.respond_with_u32(bytes_read_u32)?;
131            }
132            Err(err) => {
133                // The stream encountered an error, drop it.
134                self.discard();
135                return Err(JsError::new(&err.to_string()).into());
136            }
137        };
138        Ok(JsValue::undefined())
139    }
140
141    #[inline]
142    fn discard(&mut self) {
143        self.async_read = None;
144        self.buffer = Vec::new();
145    }
146}