rw_stream_sink/
lib.rs

1// Copyright 2017-2020 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21//! This crate provides the [`RwStreamSink`] type. It wraps around a [`Stream`]
22//! and [`Sink`] that produces and accepts byte arrays, and implements
23//! [`AsyncRead`] and [`AsyncWrite`].
24//!
25//! Each call to [`AsyncWrite::poll_write`] will send one packet to the sink.
26//! Calls to [`AsyncRead::poll_read`] will read from the stream's incoming packets.
27
28#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
29
30use std::{
31    io::{self, Read},
32    mem,
33    pin::Pin,
34    task::{Context, Poll},
35};
36
37use futures::{prelude::*, ready};
38
39static_assertions::const_assert!(mem::size_of::<usize>() <= mem::size_of::<u64>());
40
41/// Wraps a [`Stream`] and [`Sink`] whose items are buffers.
42/// Implements [`AsyncRead`] and [`AsyncWrite`].
43#[pin_project::pin_project]
44pub struct RwStreamSink<S: TryStream> {
45    #[pin]
46    inner: S,
47    current_item: Option<std::io::Cursor<<S as TryStream>::Ok>>,
48}
49
50impl<S: TryStream> RwStreamSink<S> {
51    /// Wraps around `inner`.
52    pub fn new(inner: S) -> Self {
53        RwStreamSink {
54            inner,
55            current_item: None,
56        }
57    }
58}
59
60impl<S> AsyncRead for RwStreamSink<S>
61where
62    S: TryStream<Error = io::Error>,
63    <S as TryStream>::Ok: AsRef<[u8]>,
64{
65    fn poll_read(
66        self: Pin<&mut Self>,
67        cx: &mut Context,
68        buf: &mut [u8],
69    ) -> Poll<io::Result<usize>> {
70        let mut this = self.project();
71
72        // Grab the item to copy from.
73        let item_to_copy = loop {
74            if let Some(ref mut i) = this.current_item {
75                if i.position() < i.get_ref().as_ref().len() as u64 {
76                    break i;
77                }
78            }
79            *this.current_item = Some(match ready!(this.inner.as_mut().try_poll_next(cx)) {
80                Some(Ok(i)) => std::io::Cursor::new(i),
81                Some(Err(e)) => return Poll::Ready(Err(e)),
82                None => return Poll::Ready(Ok(0)), // EOF
83            });
84        };
85
86        // Copy it!
87        Poll::Ready(Ok(item_to_copy.read(buf)?))
88    }
89}
90
91impl<S> AsyncWrite for RwStreamSink<S>
92where
93    S: TryStream + Sink<<S as TryStream>::Ok, Error = io::Error>,
94    <S as TryStream>::Ok: for<'r> From<&'r [u8]>,
95{
96    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
97        let mut this = self.project();
98        ready!(this.inner.as_mut().poll_ready(cx)?);
99        let n = buf.len();
100        if let Err(e) = this.inner.start_send(buf.into()) {
101            return Poll::Ready(Err(e));
102        }
103        Poll::Ready(Ok(n))
104    }
105
106    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
107        let this = self.project();
108        this.inner.poll_flush(cx)
109    }
110
111    fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
112        let this = self.project();
113        this.inner.poll_close(cx)
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use std::{
120        pin::Pin,
121        task::{Context, Poll},
122    };
123
124    use async_std::task;
125    use futures::{channel::mpsc, prelude::*};
126
127    use super::RwStreamSink;
128
129    // This struct merges a stream and a sink and is quite useful for tests.
130    struct Wrapper<St, Si>(St, Si);
131
132    impl<St, Si> Stream for Wrapper<St, Si>
133    where
134        St: Stream + Unpin,
135        Si: Unpin,
136    {
137        type Item = St::Item;
138
139        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
140            self.0.poll_next_unpin(cx)
141        }
142    }
143
144    impl<St, Si, T> Sink<T> for Wrapper<St, Si>
145    where
146        St: Unpin,
147        Si: Sink<T> + Unpin,
148    {
149        type Error = Si::Error;
150
151        fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
152            Pin::new(&mut self.1).poll_ready(cx)
153        }
154
155        fn start_send(mut self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
156            Pin::new(&mut self.1).start_send(item)
157        }
158
159        fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
160            Pin::new(&mut self.1).poll_flush(cx)
161        }
162
163        fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
164            Pin::new(&mut self.1).poll_close(cx)
165        }
166    }
167
168    #[test]
169    fn basic_reading() {
170        let (tx1, _) = mpsc::channel::<Vec<u8>>(10);
171        let (mut tx2, rx2) = mpsc::channel(10);
172
173        let mut wrapper = RwStreamSink::new(Wrapper(rx2.map(Ok), tx1));
174
175        task::block_on(async move {
176            tx2.send(Vec::from("hel")).await.unwrap();
177            tx2.send(Vec::from("lo wor")).await.unwrap();
178            tx2.send(Vec::from("ld")).await.unwrap();
179            tx2.close().await.unwrap();
180
181            let mut data = Vec::new();
182            wrapper.read_to_end(&mut data).await.unwrap();
183            assert_eq!(data, b"hello world");
184        })
185    }
186
187    #[test]
188    fn skip_empty_stream_items() {
189        let data: Vec<&[u8]> = vec![b"", b"foo", b"", b"bar", b"", b"baz", b""];
190        let mut rws = RwStreamSink::new(stream::iter(data).map(Ok));
191        let mut buf = [0; 9];
192        task::block_on(async move {
193            assert_eq!(3, rws.read(&mut buf).await.unwrap());
194            assert_eq!(3, rws.read(&mut buf[3..]).await.unwrap());
195            assert_eq!(3, rws.read(&mut buf[6..]).await.unwrap());
196            assert_eq!(0, rws.read(&mut buf).await.unwrap());
197            assert_eq!(b"foobarbaz", &buf[..])
198        })
199    }
200
201    #[test]
202    fn partial_read() {
203        let data: Vec<&[u8]> = vec![b"hell", b"o world"];
204        let mut rws = RwStreamSink::new(stream::iter(data).map(Ok));
205        let mut buf = [0; 3];
206        task::block_on(async move {
207            assert_eq!(3, rws.read(&mut buf).await.unwrap());
208            assert_eq!(b"hel", &buf[..3]);
209            assert_eq!(0, rws.read(&mut buf[..0]).await.unwrap());
210            assert_eq!(1, rws.read(&mut buf).await.unwrap());
211            assert_eq!(b"l", &buf[..1]);
212            assert_eq!(3, rws.read(&mut buf).await.unwrap());
213            assert_eq!(b"o w", &buf[..3]);
214            assert_eq!(0, rws.read(&mut buf[..0]).await.unwrap());
215            assert_eq!(3, rws.read(&mut buf).await.unwrap());
216            assert_eq!(b"orl", &buf[..3]);
217            assert_eq!(1, rws.read(&mut buf).await.unwrap());
218            assert_eq!(b"d", &buf[..1]);
219            assert_eq!(0, rws.read(&mut buf).await.unwrap());
220        })
221    }
222}
rw_stream_sink/lib.rs

rw_stream_sink/
lib.rs
