//! Mocking helpers for testing with futures::io types. //! //! Note that some of this code might be of general use, but for now //! we're only trying it for testing. use futures::channel::mpsc; use futures::io::{AsyncRead, AsyncWrite}; use futures::sink::{Sink, SinkExt}; use futures::stream::Stream; use std::io::{Error as IoError, ErrorKind, Result as IoResult}; use std::pin::Pin; use std::task::{Context, Poll}; /// Channel capacity for our internal MPSC channels. /// /// We keep this intentionally low to make sure that some blocking /// will occur occur. const CAPACITY: usize = 4; /// Maximum size for a queued buffer on a local chunk. /// /// This size is deliberately weird, to try to find errors. const CHUNKSZ: usize = 213; /// Construct a new pair of linked LocalStream objects. /// /// Any bytes written to one will be readable on the other, and vice /// versa. These streams will behave more or less like a socketpair, /// except without actually going through the operating system. /// /// Note that this implementation is intended for testing only, and /// isn't optimized. pub fn stream_pair() -> (LocalStream, LocalStream) { let (w1, r2) = mpsc::channel(CAPACITY); let (w2, r1) = mpsc::channel(CAPACITY); let s1 = LocalStream { w: w1, r: r1, pending_bytes: Vec::new(), tls_cert: None, }; let s2 = LocalStream { w: w2, r: r2, pending_bytes: Vec::new(), tls_cert: None, }; (s1, s2) } /// One half of a pair of linked streams returned by [`stream_pair`]. // // Implementation notes: linked streams are made out a pair of mpsc // channels. There's one channel for sending bytes in each direction. // Bytes are sent as IoResult>: sending an error causes an error // to occur on the other side. pub struct LocalStream { /// The writing side of the channel that we use to implement this /// stream. /// /// The reading side is held by the other linked stream. w: mpsc::Sender>>, /// The reading side of the channel that we use to implement this /// stream. /// /// The writing side is held by the other linked stream. r: mpsc::Receiver>>, /// Bytes that we have read from `r` but not yet delivered. pending_bytes: Vec, /// Data about the other side of this stream's fake TLS certificate, if any. /// If this is present, I/O operations will fail with an error. /// /// How this is intended to work: things that return `LocalStream`s that could potentially /// be connected to a fake TLS listener should set this field. Then, a fake TLS wrapper /// type would clear this field (after checking its contents are as expected). /// /// FIXME(eta): this is a bit of a layering violation, but it's hard to do otherwise pub(crate) tls_cert: Option>, } /// Helper: pull bytes off the front of `pending_bytes` and put them /// onto `buf. Return the number of bytes moved. fn drain_helper(buf: &mut [u8], pending_bytes: &mut Vec) -> usize { let n_to_drain = std::cmp::min(buf.len(), pending_bytes.len()); buf[..n_to_drain].copy_from_slice(&pending_bytes[..n_to_drain]); pending_bytes.drain(..n_to_drain); n_to_drain } impl AsyncRead for LocalStream { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8], ) -> Poll> { if buf.is_empty() { return Poll::Ready(Ok(0)); } if self.tls_cert.is_some() { return Poll::Ready(Err(std::io::Error::new( std::io::ErrorKind::Other, "attempted to treat a TLS stream as non-TLS!", ))); } if !self.pending_bytes.is_empty() { return Poll::Ready(Ok(drain_helper(buf, &mut self.pending_bytes))); } match futures::ready!(Pin::new(&mut self.r).poll_next(cx)) { Some(Err(e)) => Poll::Ready(Err(e)), Some(Ok(bytes)) => { self.pending_bytes = bytes; let n = drain_helper(buf, &mut self.pending_bytes); Poll::Ready(Ok(n)) } None => Poll::Ready(Ok(0)), // This is an EOF } } } impl AsyncWrite for LocalStream { fn poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { if self.tls_cert.is_some() { return Poll::Ready(Err(std::io::Error::new( std::io::ErrorKind::Other, "attempted to treat a TLS stream as non-TLS!", ))); } match futures::ready!(Pin::new(&mut self.w).poll_ready(cx)) { Ok(()) => (), Err(e) => return Poll::Ready(Err(IoError::new(ErrorKind::BrokenPipe, e))), } let buf = if buf.len() > CHUNKSZ { &buf[..CHUNKSZ] } else { buf }; let len = buf.len(); match Pin::new(&mut self.w).start_send(Ok(buf.to_vec())) { Ok(()) => Poll::Ready(Ok(len)), Err(e) => Poll::Ready(Err(IoError::new(ErrorKind::BrokenPipe, e))), } } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.w) .poll_flush(cx) .map_err(|e| IoError::new(ErrorKind::BrokenPipe, e)) } fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.w) .poll_close(cx) .map_err(|e| IoError::new(ErrorKind::Other, e)) } } /// An error generated by [`LocalStream::send_err`]. #[derive(Debug, Clone, Eq, PartialEq)] #[non_exhaustive] pub struct SyntheticError; impl std::error::Error for SyntheticError {} impl std::fmt::Display for SyntheticError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "Synthetic error") } } impl LocalStream { /// Send an error to the other linked local stream. /// /// When the other stream reads this message, it will generate a /// [`std::io::Error`] with the provided `ErrorKind`. pub async fn send_err(&mut self, kind: ErrorKind) { let _ignore = self.w.send(Err(IoError::new(kind, SyntheticError))).await; } } #[cfg(test)] mod test { #![allow(clippy::unwrap_used)] use super::*; use futures::io::{AsyncReadExt, AsyncWriteExt}; use futures_await_test::async_test; use rand::thread_rng; use rand::Rng; #[async_test] async fn basic_rw() { let (mut s1, mut s2) = stream_pair(); let mut text1 = vec![0_u8; 9999]; thread_rng().fill(&mut text1[..]); let (v1, v2): (IoResult<()>, IoResult<()>) = futures::join!( async { for _ in 0_u8..10 { s1.write_all(&text1[..]).await?; } s1.close().await?; Ok(()) }, async { let mut text2: Vec = Vec::new(); let mut buf = [0_u8; 33]; loop { let n = s2.read(&mut buf[..]).await?; if n == 0 { break; } text2.extend(&buf[..n]); } for ch in text2[..].chunks(text1.len()) { assert_eq!(ch, &text1[..]); } Ok(()) } ); v1.unwrap(); v2.unwrap(); } #[async_test] async fn send_error() { let (mut s1, mut s2) = stream_pair(); let (v1, v2): (IoResult<()>, IoResult<()>) = futures::join!( async { s1.write_all(b"hello world").await?; s1.send_err(ErrorKind::PermissionDenied).await; Ok(()) }, async { let mut buf = [0_u8; 33]; loop { let n = s2.read(&mut buf[..]).await?; if n == 0 { break; } } Ok(()) } ); v1.unwrap(); let e = v2.err().unwrap(); assert_eq!(e.kind(), ErrorKind::PermissionDenied); let synth = e.into_inner().unwrap(); assert_eq!(synth.to_string(), "Synthetic error"); } #[async_test] async fn drop_reader() { let (mut s1, s2) = stream_pair(); let (v1, v2): (IoResult<()>, IoResult<()>) = futures::join!( async { for _ in 0_u16..1000 { s1.write_all(&[9_u8; 9999]).await?; } Ok(()) }, async { drop(s2); Ok(()) } ); v2.unwrap(); let e = v1.err().unwrap(); assert_eq!(e.kind(), ErrorKind::BrokenPipe); } }