1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use futures_core::task::{Context, Poll};
use futures_io::{AsyncRead, AsyncWrite, AsyncSeek, AsyncBufRead};
use std::{fmt, io};
use std::pin::Pin;

/// A simple wrapper type which allows types which implement only
/// implement `std::io::Read` or `std::io::Write`
/// to be used in contexts which expect an `AsyncRead` or `AsyncWrite`.
///
/// If these types issue an error with the kind `io::ErrorKind::WouldBlock`,
/// it is expected that they will notify the current task on readiness.
/// Synchronous `std` types should not issue errors of this kind and
/// are safe to use in this context. However, using these types with
/// `AllowStdIo` will cause the event loop to block, so they should be used
/// with care.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct AllowStdIo<T>(T);

impl<T> Unpin for AllowStdIo<T> {}

macro_rules! try_with_interrupt {
    ($e:expr) => {
        loop {
            match $e {
                Ok(e) => {
                    break e;
                }
                Err(ref e) if e.kind() == ::std::io::ErrorKind::Interrupted => {
                    continue;
                }
                Err(e) => {
                    return Poll::Ready(Err(e));
                }
            }
        }
    }
}

impl<T> AllowStdIo<T> {
    /// Creates a new `AllowStdIo` from an existing IO object.
    pub fn new(io: T) -> Self {
        AllowStdIo(io)
    }

    /// Returns a reference to the contained IO object.
    pub fn get_ref(&self) -> &T {
        &self.0
    }

    /// Returns a mutable reference to the contained IO object.
    pub fn get_mut(&mut self) -> &mut T {
        &mut self.0
    }

    /// Consumes self and returns the contained IO object.
    pub fn into_inner(self) -> T {
        self.0
    }
}

impl<T> io::Write for AllowStdIo<T> where T: io::Write {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.0.write_all(buf)
    }
    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
        self.0.write_fmt(fmt)
    }
}

impl<T> AsyncWrite for AllowStdIo<T> where T: io::Write {
    fn poll_write(mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &[u8])
        -> Poll<io::Result<usize>>
    {
        Poll::Ready(Ok(try_with_interrupt!(self.0.write(buf))))
    }

    fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        try_with_interrupt!(self.0.flush());
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        self.poll_flush(cx)
    }
}

impl<T> io::Read for AllowStdIo<T> where T: io::Read {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
    // TODO: implement the `initializer` fn when it stabilizes.
    // See rust-lang/rust #42788
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        self.0.read_to_end(buf)
    }
    fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
        self.0.read_to_string(buf)
    }
    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        self.0.read_exact(buf)
    }
}

impl<T> AsyncRead for AllowStdIo<T> where T: io::Read {
    fn poll_read(mut self: Pin<&mut Self>, _: &mut Context<'_>, buf: &mut [u8])
        -> Poll<io::Result<usize>>
    {
        Poll::Ready(Ok(try_with_interrupt!(self.0.read(buf))))
    }
}

impl<T> io::Seek for AllowStdIo<T> where T: io::Seek {
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        self.0.seek(pos)
    }
}

impl<T> AsyncSeek for AllowStdIo<T> where T: io::Seek {
    fn poll_seek(mut self: Pin<&mut Self>, _: &mut Context<'_>, pos: io::SeekFrom)
        -> Poll<io::Result<u64>>
    {
        Poll::Ready(Ok(try_with_interrupt!(self.0.seek(pos))))
    }
}

impl<T> io::BufRead for AllowStdIo<T> where T: io::BufRead {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        self.0.fill_buf()
    }
    fn consume(&mut self, amt: usize) {
        self.0.consume(amt)
    }
}

impl<T> AsyncBufRead for AllowStdIo<T> where T: io::BufRead {
    fn poll_fill_buf<'a>(mut self: Pin<&'a mut Self>, _: &mut Context<'_>)
        -> Poll<io::Result<&'a [u8]>>
    {
        let this: *mut Self = &mut *self as *mut _;
        Poll::Ready(Ok(try_with_interrupt!(unsafe { &mut *this }.0.fill_buf())))
    }

    fn consume(mut self: Pin<&mut Self>, amt: usize) {
        self.0.consume(amt)
    }
}