|
| 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 | +use futures::{prelude::*, ready}; |
| 29 | +use std::{ |
| 30 | + io::{self, Read}, |
| 31 | + mem, |
| 32 | + pin::Pin, |
| 33 | + task::{Context, Poll}, |
| 34 | +}; |
| 35 | + |
| 36 | +static_assertions::const_assert!(mem::size_of::<usize>() <= mem::size_of::<u64>()); |
| 37 | + |
| 38 | +/// Wraps a [`Stream`] and [`Sink`] whose items are buffers. |
| 39 | +/// Implements [`AsyncRead`] and [`AsyncWrite`]. |
| 40 | +#[pin_project::pin_project] |
| 41 | +pub struct RwStreamSink<S: TryStream> { |
| 42 | + #[pin] |
| 43 | + inner: S, |
| 44 | + current_item: Option<std::io::Cursor<<S as TryStream>::Ok>>, |
| 45 | +} |
| 46 | + |
| 47 | +impl<S: TryStream> RwStreamSink<S> { |
| 48 | + /// Wraps around `inner`. |
| 49 | + pub fn new(inner: S) -> Self { |
| 50 | + RwStreamSink { |
| 51 | + inner, |
| 52 | + current_item: None, |
| 53 | + } |
| 54 | + } |
| 55 | +} |
| 56 | + |
| 57 | +impl<S> AsyncRead for RwStreamSink<S> |
| 58 | +where |
| 59 | + S: TryStream<Error = io::Error>, |
| 60 | + <S as TryStream>::Ok: AsRef<[u8]>, |
| 61 | +{ |
| 62 | + fn poll_read( |
| 63 | + self: Pin<&mut Self>, |
| 64 | + cx: &mut Context, |
| 65 | + buf: &mut [u8], |
| 66 | + ) -> Poll<io::Result<usize>> { |
| 67 | + let mut this = self.project(); |
| 68 | + |
| 69 | + // Grab the item to copy from. |
| 70 | + let item_to_copy = loop { |
| 71 | + if let Some(ref mut i) = this.current_item { |
| 72 | + if i.position() < i.get_ref().as_ref().len() as u64 { |
| 73 | + break i; |
| 74 | + } |
| 75 | + } |
| 76 | + *this.current_item = Some(match ready!(this.inner.as_mut().try_poll_next(cx)) { |
| 77 | + Some(Ok(i)) => std::io::Cursor::new(i), |
| 78 | + Some(Err(e)) => return Poll::Ready(Err(e)), |
| 79 | + None => return Poll::Ready(Ok(0)), // EOF |
| 80 | + }); |
| 81 | + }; |
| 82 | + |
| 83 | + // Copy it! |
| 84 | + Poll::Ready(Ok(item_to_copy.read(buf)?)) |
| 85 | + } |
| 86 | +} |
| 87 | + |
| 88 | +impl<S> AsyncWrite for RwStreamSink<S> |
| 89 | +where |
| 90 | + S: TryStream + Sink<<S as TryStream>::Ok, Error = io::Error>, |
| 91 | + <S as TryStream>::Ok: for<'r> From<&'r [u8]>, |
| 92 | +{ |
| 93 | + fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> { |
| 94 | + let mut this = self.project(); |
| 95 | + ready!(this.inner.as_mut().poll_ready(cx)?); |
| 96 | + let n = buf.len(); |
| 97 | + if let Err(e) = this.inner.start_send(buf.into()) { |
| 98 | + return Poll::Ready(Err(e)); |
| 99 | + } |
| 100 | + Poll::Ready(Ok(n)) |
| 101 | + } |
| 102 | + |
| 103 | + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> { |
| 104 | + let this = self.project(); |
| 105 | + this.inner.poll_flush(cx) |
| 106 | + } |
| 107 | + |
| 108 | + fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> { |
| 109 | + let this = self.project(); |
| 110 | + this.inner.poll_close(cx) |
| 111 | + } |
| 112 | +} |
| 113 | + |
| 114 | +#[cfg(test)] |
| 115 | +mod tests { |
| 116 | + use super::RwStreamSink; |
| 117 | + use async_std::task; |
| 118 | + use futures::{channel::mpsc, prelude::*, stream}; |
| 119 | + use std::{ |
| 120 | + pin::Pin, |
| 121 | + task::{Context, Poll}, |
| 122 | + }; |
| 123 | + |
| 124 | + // This struct merges a stream and a sink and is quite useful for tests. |
| 125 | + struct Wrapper<St, Si>(St, Si); |
| 126 | + |
| 127 | + impl<St, Si> Stream for Wrapper<St, Si> |
| 128 | + where |
| 129 | + St: Stream + Unpin, |
| 130 | + Si: Unpin, |
| 131 | + { |
| 132 | + type Item = St::Item; |
| 133 | + |
| 134 | + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> { |
| 135 | + self.0.poll_next_unpin(cx) |
| 136 | + } |
| 137 | + } |
| 138 | + |
| 139 | + impl<St, Si, T> Sink<T> for Wrapper<St, Si> |
| 140 | + where |
| 141 | + St: Unpin, |
| 142 | + Si: Sink<T> + Unpin, |
| 143 | + { |
| 144 | + type Error = Si::Error; |
| 145 | + |
| 146 | + fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> { |
| 147 | + Pin::new(&mut self.1).poll_ready(cx) |
| 148 | + } |
| 149 | + |
| 150 | + fn start_send(mut self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> { |
| 151 | + Pin::new(&mut self.1).start_send(item) |
| 152 | + } |
| 153 | + |
| 154 | + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> { |
| 155 | + Pin::new(&mut self.1).poll_flush(cx) |
| 156 | + } |
| 157 | + |
| 158 | + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> { |
| 159 | + Pin::new(&mut self.1).poll_close(cx) |
| 160 | + } |
| 161 | + } |
| 162 | + |
| 163 | + #[test] |
| 164 | + fn basic_reading() { |
| 165 | + let (tx1, _) = mpsc::channel::<Vec<u8>>(10); |
| 166 | + let (mut tx2, rx2) = mpsc::channel(10); |
| 167 | + |
| 168 | + let mut wrapper = RwStreamSink::new(Wrapper(rx2.map(Ok), tx1)); |
| 169 | + |
| 170 | + task::block_on(async move { |
| 171 | + tx2.send(Vec::from("hel")).await.unwrap(); |
| 172 | + tx2.send(Vec::from("lo wor")).await.unwrap(); |
| 173 | + tx2.send(Vec::from("ld")).await.unwrap(); |
| 174 | + tx2.close().await.unwrap(); |
| 175 | + |
| 176 | + let mut data = Vec::new(); |
| 177 | + wrapper.read_to_end(&mut data).await.unwrap(); |
| 178 | + assert_eq!(data, b"hello world"); |
| 179 | + }) |
| 180 | + } |
| 181 | + |
| 182 | + #[test] |
| 183 | + fn skip_empty_stream_items() { |
| 184 | + let data: Vec<&[u8]> = vec![b"", b"foo", b"", b"bar", b"", b"baz", b""]; |
| 185 | + let mut rws = RwStreamSink::new(stream::iter(data).map(Ok)); |
| 186 | + let mut buf = [0; 9]; |
| 187 | + task::block_on(async move { |
| 188 | + assert_eq!(3, rws.read(&mut buf).await.unwrap()); |
| 189 | + assert_eq!(3, rws.read(&mut buf[3..]).await.unwrap()); |
| 190 | + assert_eq!(3, rws.read(&mut buf[6..]).await.unwrap()); |
| 191 | + assert_eq!(0, rws.read(&mut buf).await.unwrap()); |
| 192 | + assert_eq!(b"foobarbaz", &buf[..]) |
| 193 | + }) |
| 194 | + } |
| 195 | + |
| 196 | + #[test] |
| 197 | + fn partial_read() { |
| 198 | + let data: Vec<&[u8]> = vec![b"hell", b"o world"]; |
| 199 | + let mut rws = RwStreamSink::new(stream::iter(data).map(Ok)); |
| 200 | + let mut buf = [0; 3]; |
| 201 | + task::block_on(async move { |
| 202 | + assert_eq!(3, rws.read(&mut buf).await.unwrap()); |
| 203 | + assert_eq!(b"hel", &buf[..3]); |
| 204 | + assert_eq!(0, rws.read(&mut buf[..0]).await.unwrap()); |
| 205 | + assert_eq!(1, rws.read(&mut buf).await.unwrap()); |
| 206 | + assert_eq!(b"l", &buf[..1]); |
| 207 | + assert_eq!(3, rws.read(&mut buf).await.unwrap()); |
| 208 | + assert_eq!(b"o w", &buf[..3]); |
| 209 | + assert_eq!(0, rws.read(&mut buf[..0]).await.unwrap()); |
| 210 | + assert_eq!(3, rws.read(&mut buf).await.unwrap()); |
| 211 | + assert_eq!(b"orl", &buf[..3]); |
| 212 | + assert_eq!(1, rws.read(&mut buf).await.unwrap()); |
| 213 | + assert_eq!(b"d", &buf[..1]); |
| 214 | + assert_eq!(0, rws.read(&mut buf).await.unwrap()); |
| 215 | + }) |
| 216 | + } |
| 217 | +} |
0 commit comments