wrpc_transport/
pack.rs

1//! Packing and unpacking of wRPC values to/from singular, flat byte buffers
2
3use std::collections::VecDeque;
4
5use bytes::{BufMut as _, Bytes, BytesMut};
6use futures::{Stream, StreamExt as _};
7use tokio::io::{AsyncRead, AsyncReadExt as _};
8use tokio::sync::mpsc;
9use tokio_util::codec::{Decoder, Encoder};
10
11use crate::{Decode, Deferred as _, DeferredFn, Encode, Index, Sink, encode_ready_bytes_header};
12
13/// [Sink] used by [pack], which collects the packed asynchronous values
14#[derive(Clone)]
15pub struct Packer(mpsc::UnboundedSender<(usize, VecDeque<Bytes>)>);
16
17impl Packer {
18    fn push(&self, offset: usize, bufs: VecDeque<Bytes>) -> std::io::Result<()> {
19        self.0
20            .send((offset, bufs))
21            .map_err(|_| std::io::Error::from(std::io::ErrorKind::BrokenPipe))
22    }
23
24    fn push_bytes(&self, offset: usize, mut chunks: VecDeque<Bytes>) -> std::io::Result<()> {
25        let mut header = BytesMut::default();
26        encode_ready_bytes_header(chunks.iter().map(Bytes::len).sum(), &mut header)?;
27        chunks.push_front(header.freeze());
28        self.push(offset, chunks)
29    }
30}
31
32impl Index for Packer {
33    fn index(&self, _: &[usize]) -> std::io::Result<Self> {
34        Ok(self.clone())
35    }
36}
37
38impl Sink for Packer {
39    async fn write_future<T>(
40        self,
41        offset: usize,
42        item: impl Future<Output = T> + Send + 'static,
43    ) -> std::io::Result<()>
44    where
45        T: Encode<Self>,
46        std::io::Error: From<<T::Encoder as Encoder<T>>::Error>,
47    {
48        let mut buf = BytesMut::default();
49        buf.put_u8(0x01);
50        let deferred = T::encode(item.await, &mut T::Encoder::default(), &mut buf)?;
51        fill(deferred, &mut buf).await?;
52        self.push(offset, VecDeque::from([buf.freeze()]))
53    }
54
55    async fn write_stream<T>(
56        self,
57        offset: usize,
58        items: impl Stream<Item = Vec<T>> + Send + Unpin + 'static,
59    ) -> std::io::Result<()>
60    where
61        T: Encode<Self> + Send + 'static,
62        std::io::Error: From<<T::Encoder as Encoder<T>>::Error>,
63    {
64        self.write_future(offset, items.concat()).await
65    }
66
67    async fn write_stream_bytes(
68        self,
69        offset: usize,
70        items: impl Stream<Item = Bytes> + Send + Unpin + 'static,
71    ) -> std::io::Result<()> {
72        self.push_bytes(offset, items.collect().await)
73    }
74
75    async fn write_stream_read(
76        self,
77        offset: usize,
78        mut ready: BytesMut,
79        mut items: impl AsyncRead + Send + Unpin + 'static,
80    ) -> std::io::Result<()> {
81        while items.read_buf(&mut ready).await? > 0 {}
82        self.push_bytes(offset, VecDeque::from([ready.freeze()]))
83    }
84}
85
86async fn fill(deferred: Option<DeferredFn<Packer>>, dst: &mut BytesMut) -> std::io::Result<()> {
87    let Some(f) = deferred else {
88        return Ok(());
89    };
90    let (tx, mut rx) = mpsc::unbounded_channel();
91    f(Packer(tx), Vec::default()).await?;
92    let mut holes = Vec::default();
93    while let Some(hole) = rx.recv().await {
94        holes.push(hole);
95    }
96    holes.sort_unstable_by_key(|(offset, _)| *offset);
97    let n = holes
98        .iter()
99        .flat_map(|(_, bufs)| bufs)
100        .map(Bytes::len)
101        .sum::<usize>();
102    let mut out = BytesMut::with_capacity(dst.len() + n);
103    let mut prev = 0;
104    for (offset, bufs) in holes {
105        if offset < prev || dst.get(offset) != Some(&0x00) {
106            return Err(std::io::Error::new(
107                std::io::ErrorKind::InvalidData,
108                "asynchronous value was not encoded at the expected offset",
109            ));
110        }
111        out.extend_from_slice(&dst[prev..offset]);
112        for buf in bufs {
113            out.extend_from_slice(&buf);
114        }
115        prev = offset + 1;
116    }
117    out.extend_from_slice(&dst[prev..]);
118    *dst = out;
119    Ok(())
120}
121
122/// Pack a [`Encode`] into a singular byte buffer `dst`.
123///
124/// All asynchronous values contained in `v` are awaited and encoded as ready values:
125///
126/// - `future<T>` is encoded as `ready` variant of `future<T>`
127/// - `stream<T>` is encoded as `ready` variant of `stream<T>` containing all elements of the stream
128///
129/// Note, that this function will never complete if any of the asynchronous values contained in
130/// `v` never resolve, use e.g. [`tokio::time::timeout`] to bound the time spent packing.
131///
132/// This function does not verify that all asynchronous values have been written. If a deferred
133/// function returned by an encoder does not write every asynchronous value it is responsible for,
134/// the corresponding pending placeholders are left in `dst` and [`unpack`] will fail.
135///
136/// If this function returns an error, contents of `dst` are undefined.
137pub async fn pack<T: Encode<Packer>>(
138    v: T,
139    dst: &mut BytesMut,
140) -> Result<(), <T::Encoder as Encoder<T>>::Error> {
141    let deferred = v.encode(&mut T::Encoder::default(), dst)?;
142    fill(deferred, dst).await?;
143    Ok(())
144}
145
146/// Synchronously pack a [`Encode`] into a singular byte buffer `dst`.
147///
148/// All asynchronous values contained in `v` must be ready, this function returns an error of
149/// kind [`std::io::ErrorKind::WouldBlock`] if any of them is pending.
150///
151/// Futures are polled exactly once and streams are polled until they are pending or finished,
152/// so values backed by Tokio I/O types (e.g. `tokio::fs::File`) are typically pending and require
153/// this function to be called within the context of a Tokio runtime.
154///
155/// If this function returns an error, contents of `dst` are undefined.
156pub fn pack_ready<T: Encode<Packer>>(
157    v: T,
158    dst: &mut BytesMut,
159) -> Result<(), <T::Encoder as Encoder<T>>::Error> {
160    let deferred = v.encode(&mut T::Encoder::default(), dst)?;
161    if deferred.is_some() {
162        return Err(std::io::Error::new(
163            std::io::ErrorKind::WouldBlock,
164            "asynchronous value is pending",
165        )
166        .into());
167    }
168    Ok(())
169}
170
171/// Unpack a [`Decode`] from a byte buffer `buf` produced by [`pack`] or [`pack_ready`].
172///
173/// This function returns an error if `buf` contains pending asynchronous values or if `buf` is
174/// not fully consumed.
175///
176/// If this function returns an error, contents of `buf` are undefined.
177pub fn unpack<T: Decode>(buf: &mut BytesMut) -> Result<T, <T::Decoder as Decoder>::Error> {
178    let mut dec = T::Decoder::default();
179    let Some(v) = dec.decode(buf)? else {
180        return Err(
181            std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "buffer is incomplete").into(),
182        );
183    };
184    if dec.take_deferred().is_some() {
185        return Err(std::io::Error::new(
186            std::io::ErrorKind::InvalidData,
187            "buffer contains pending asynchronous values, which cannot be unpacked",
188        )
189        .into());
190    }
191    if !buf.is_empty() {
192        return Err(std::io::Error::new(
193            std::io::ErrorKind::InvalidData,
194            "buffer contains trailing bytes",
195        )
196        .into());
197    }
198    Ok(v)
199}
200
201#[cfg(test)]
202mod tests {
203    use core::future::{pending, ready};
204    use core::pin::Pin;
205    use core::time::Duration;
206
207    use futures::FutureExt as _;
208    use futures::stream::{self, BoxStream};
209    use tokio::sync::oneshot;
210    use tokio::task::yield_now;
211    use tokio_util::io::StreamReader;
212
213    use super::*;
214
215    #[test_log::test(tokio::test)]
216    async fn sync() -> anyhow::Result<()> {
217        let mut buf = BytesMut::default();
218        pack((0x42_u8, "foo".to_string()), &mut buf).await?;
219        assert_eq!(buf.as_ref(), b"\x42\x03foo");
220        let (a, b) = unpack::<(u8, String)>(&mut buf)?;
221        assert_eq!(a, 0x42);
222        assert_eq!(b, "foo");
223        Ok(())
224    }
225
226    #[test_log::test(tokio::test)]
227    async fn future() -> anyhow::Result<()> {
228        let (tx, rx) = oneshot::channel::<u32>();
229        let fut: Pin<Box<dyn Future<Output = Pin<Box<dyn Future<Output = u32> + Send>>> + Send>> =
230            Box::pin(ready(
231                Box::pin(async { rx.await.unwrap() }) as Pin<Box<dyn Future<Output = u32> + Send>>
232            ));
233        let mut buf = BytesMut::default();
234        let (res, ()) = tokio::join!(pack((fut, 0x01_u8), &mut buf), async {
235            tx.send(0x42).unwrap();
236        });
237        res?;
238        assert_eq!(buf.as_ref(), b"\x01\x01\x42\x01");
239
240        let (fut, v) = unpack::<(
241            Pin<Box<dyn Future<Output = Pin<Box<dyn Future<Output = u32> + Send>>> + Send>>,
242            u8,
243        )>(&mut buf)?;
244        assert_eq!(v, 0x01);
245        assert_eq!(fut.await.await, 0x42);
246        Ok(())
247    }
248
249    #[test_log::test(tokio::test)]
250    async fn stream() -> anyhow::Result<()> {
251        let s: Pin<Box<dyn Stream<Item = Vec<Pin<Box<dyn Future<Output = u8> + Send>>>> + Send>> =
252            Box::pin(stream::iter([
253                vec![
254                    Box::pin(async { 1_u8 }) as Pin<Box<dyn Future<Output = u8> + Send>>,
255                    Box::pin(async { 2 }),
256                ],
257                vec![],
258                vec![Box::pin(async { 3 })],
259            ]));
260        let mut buf = BytesMut::default();
261        pack(s, &mut buf).await?;
262        assert_eq!(buf.as_ref(), b"\x01\x03\x01\x01\x01\x02\x01\x03");
263
264        let s = unpack::<
265            Pin<Box<dyn Stream<Item = Vec<Pin<Box<dyn Future<Output = u8> + Send>>>> + Send>>,
266        >(&mut buf)?;
267        let mut vs = vec![];
268        for fut in s.concat().await {
269            vs.push(fut.await);
270        }
271        assert_eq!(vs, [1, 2, 3]);
272        Ok(())
273    }
274
275    #[test_log::test(tokio::test)]
276    async fn stream_bytes() -> anyhow::Result<()> {
277        let s: Pin<Box<dyn Stream<Item = Bytes> + Send>> = Box::pin(stream::iter([
278            Bytes::from_static(b"foo"),
279            Bytes::from_static(b"bar"),
280        ]));
281        let r: Pin<Box<dyn AsyncRead + Send>> = Box::pin(std::io::Cursor::new(b"baz"));
282        let mut buf = BytesMut::default();
283        pack((s, r), &mut buf).await?;
284        assert_eq!(buf.as_ref(), b"\x01\x06foobar\x01\x03baz");
285
286        let (s, mut r) = unpack::<(
287            Pin<Box<dyn Stream<Item = Bytes> + Send>>,
288            Pin<Box<dyn AsyncRead + Send>>,
289        )>(&mut buf)?;
290        let s: Vec<Bytes> = s.collect().await;
291        assert_eq!(s.concat(), b"foobar");
292        let mut v = vec![];
293        r.read_to_end(&mut v).await?;
294        assert_eq!(v, b"baz");
295        Ok(())
296    }
297
298    #[test_log::test(tokio::test)]
299    async fn stream_empty() -> anyhow::Result<()> {
300        let s: Pin<Box<dyn Stream<Item = Vec<u8>> + Send>> = Box::pin(stream::empty());
301        let b: Pin<Box<dyn Stream<Item = Bytes> + Send>> = Box::pin(stream::empty());
302        let mut buf = BytesMut::default();
303        pack((s, b, tokio::io::empty()), &mut buf).await?;
304        assert_eq!(buf.as_ref(), b"\x01\x00\x01\x00\x01\x00");
305
306        let (s, b, mut r) = unpack::<(
307            Pin<Box<dyn Stream<Item = Vec<u8>> + Send>>,
308            Pin<Box<dyn Stream<Item = Bytes> + Send>>,
309            Pin<Box<dyn AsyncRead + Send>>,
310        )>(&mut buf)?;
311        assert!(s.concat().await.is_empty());
312        assert_eq!(b.count().await, 0);
313        let mut v = vec![];
314        r.read_to_end(&mut v).await?;
315        assert!(v.is_empty());
316        Ok(())
317    }
318
319    fn ready_then_pending<T: Send + 'static>(ready: T, pending: T) -> BoxStream<'static, T> {
320        stream::iter([ready])
321            .chain(stream::once(async {
322                yield_now().await;
323                pending
324            }))
325            .boxed()
326    }
327
328    #[test_log::test(tokio::test)]
329    async fn stream_partial() -> anyhow::Result<()> {
330        let s = ready_then_pending(vec![1_u8, 2], vec![3]);
331        let b = ready_then_pending(Bytes::from("foo"), Bytes::from("bar"));
332        let r: Pin<Box<dyn AsyncRead + Send>> = Box::pin(StreamReader::new(
333            ready_then_pending(Bytes::from("foo"), Bytes::from("bar")).map(Ok::<_, std::io::Error>),
334        ));
335        let mut buf = BytesMut::default();
336        pack((s, b, r), &mut buf).await?;
337        assert_eq!(
338            buf.as_ref(),
339            b"\x01\x03\x01\x02\x03\x01\x06foobar\x01\x06foobar"
340        );
341        Ok(())
342    }
343
344    #[test_log::test(tokio::test)]
345    async fn pending_future() -> anyhow::Result<()> {
346        let fut: Pin<Box<dyn Future<Output = u8> + Send>> = Box::pin(pending());
347        let mut buf = BytesMut::default();
348        let res = tokio::time::timeout(Duration::from_millis(10), pack(fut, &mut buf)).await;
349        assert!(res.is_err());
350        Ok(())
351    }
352
353    #[test]
354    fn ready_values() -> anyhow::Result<()> {
355        let fut = ready(0x42_u8).boxed();
356        let s = stream::iter([Bytes::from("foo")]).boxed();
357        let mut buf = BytesMut::default();
358        pack_ready((0x01_u8, fut, s), &mut buf)?;
359        assert_eq!(buf.as_ref(), b"\x01\x01\x42\x01\x03foo");
360        Ok(())
361    }
362
363    #[test]
364    fn ready_nested() -> anyhow::Result<()> {
365        let s = stream::iter((0..=255_u8).map(|i| vec![ready(i).boxed()])).boxed();
366        let mut buf = BytesMut::default();
367        pack_ready(s, &mut buf)?;
368        let expected: Vec<u8> = [0x01, 0x80, 0x02]
369            .into_iter()
370            .chain((0..=255).flat_map(|i| [0x01, i]))
371            .collect();
372        assert_eq!(buf.as_ref(), expected);
373        Ok(())
374    }
375
376    #[test]
377    fn ready_pending() {
378        let fut = ready(pending::<u8>().boxed()).boxed();
379        let mut buf = BytesMut::default();
380        let err = pack_ready(fut, &mut buf).unwrap_err();
381        assert_eq!(err.kind(), std::io::ErrorKind::WouldBlock);
382    }
383
384    #[test]
385    fn unpack_trailing() {
386        let mut buf = BytesMut::from(&b"\x42\x00"[..]);
387        let err = unpack::<u8>(&mut buf).err().unwrap();
388        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
389    }
390
391    #[test_log::test(tokio::test)]
392    async fn unpack_pending() {
393        let mut buf = BytesMut::from(&b"\x00"[..]);
394        let err = unpack::<Pin<Box<dyn Future<Output = u8> + Send>>>(&mut buf)
395            .err()
396            .unwrap();
397        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
398
399        let mut buf = BytesMut::from(&b"\x00"[..]);
400        let err = unpack::<Pin<Box<dyn Stream<Item = Vec<u8>> + Send>>>(&mut buf)
401            .err()
402            .unwrap();
403        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
404    }
405}
wrpc_transport/pack.rs

wrpc_transport/
pack.rs
