Expand Shadowsocks runtime compatibility

third_party/tokio-rustls-fork-shadow-tls/src/client.rs

@@ -0,0 +1,86 @@
+use std::sync::Arc;
+
+use tokio::io::{AsyncRead, AsyncWrite};
+use rustls_fork_shadow_tls::{ClientConfig, ClientConnection, ClientSessionIdGenerators};
+
+use crate::{
+    split::{ReadHalf, WriteHalf},
+    stream::Stream,
+    TlsError,
+};
+
+/// A wrapper around an underlying raw stream which implements the TLS protocol.
+pub type TlsStream<IO> = Stream<IO, ClientConnection>;
+/// TlsStream for read only.
+pub type TlsStreamReadHalf<IO> = ReadHalf<IO, ClientConnection>;
+/// TlsStream for write only.
+pub type TlsStreamWriteHalf<IO> = WriteHalf<IO, ClientConnection>;
+
+/// A wrapper around a `rustls::ClientConfig`, providing an async `connect` method.
+#[derive(Clone)]
+pub struct TlsConnector {
+    inner: Arc<ClientConfig>,
+}
+
+impl From<Arc<ClientConfig>> for TlsConnector {
+    fn from(inner: Arc<ClientConfig>) -> TlsConnector {
+        TlsConnector { inner }
+    }
+}
+
+impl From<ClientConfig> for TlsConnector {
+    fn from(inner: ClientConfig) -> TlsConnector {
+        TlsConnector {
+            inner: Arc::new(inner),
+        }
+    }
+}
+
+impl TlsConnector {
+    pub async fn connect<IO>(
+        &self,
+        domain: rustls_fork_shadow_tls::ServerName,
+        stream: IO,
+    ) -> Result<TlsStream<IO>, TlsError>
+    where
+        IO: AsyncRead + AsyncWrite + Unpin,
+    {
+        let session = ClientConnection::new(self.inner.clone(), domain)?;
+        let mut stream = Stream::new(stream, session);
+        stream.handshake().await?;
+        Ok(stream)
+    }
+
+    pub async fn connect_with_session_id_generator<IO, F>(
+        &self,
+        domain: rustls_fork_shadow_tls::ServerName,
+        stream: IO,
+        generator: F,
+    ) -> Result<TlsStream<IO>, TlsError>
+    where
+        IO: AsyncRead + AsyncWrite + Unpin,
+        F: Fn(&[u8]) -> [u8; 32] + Send + Sync + 'static,
+    {
+        let session =
+            ClientConnection::new_with_session_id_generator(self.inner.clone(), domain, generator)?;
+        let mut stream = Stream::new(stream, session);
+        stream.handshake().await?;
+        Ok(stream)
+    }
+
+    pub async fn connect_with_session_id_generators<IO>(
+        &self,
+        domain: rustls_fork_shadow_tls::ServerName,
+        stream: IO,
+        generators: ClientSessionIdGenerators,
+    ) -> Result<TlsStream<IO>, TlsError>
+    where
+        IO: AsyncRead + AsyncWrite + Unpin,
+    {
+        let session =
+            ClientConnection::new_with_session_id_generators(self.inner.clone(), domain, generators)?;
+        let mut stream = Stream::new(stream, session);
+        stream.handshake().await?;
+        Ok(stream)
+    }
+}

third_party/tokio-rustls-fork-shadow-tls/src/error.rs

@@ -0,0 +1,20 @@
+use std::io;
+
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum TlsError {
+    #[error("io error")]
+    Io(#[from] std::io::Error),
+    #[error("rustls error")]
+    Rustls(#[from] rustls_fork_shadow_tls::Error),
+}
+
+impl From<TlsError> for io::Error {
+    fn from(e: TlsError) -> Self {
+        match e {
+            TlsError::Io(e) => e,
+            TlsError::Rustls(e) => io::Error::new(io::ErrorKind::Other, e),
+        }
+    }
+}

third_party/tokio-rustls-fork-shadow-tls/src/lib.rs

@@ -0,0 +1,21 @@
+#![allow(stable_features)]
+
+mod client;
+mod error;
+#[cfg(not(feature = "unsafe_io"))]
+mod safe_io;
+mod server;
+mod split;
+mod stream;
+#[cfg(feature = "unsafe_io")]
+mod unsafe_io;
+
+pub use client::{
+    TlsConnector, TlsStream as ClientTlsStream, TlsStreamReadHalf as ClientTlsStreamReadHalf,
+    TlsStreamWriteHalf as ClientTlsStreamWriteHalf,
+};
+pub use error::TlsError;
+pub use server::{
+    TlsAcceptor, TlsStream as ServerTlsStream, TlsStreamReadHalf as ServerTlsStreamReadHalf,
+    TlsStreamWriteHalf as ServerTlsStreamWriteHalf,
+};

third_party/tokio-rustls-fork-shadow-tls/src/safe_io.rs

@@ -0,0 +1,227 @@
+use std::{
+    fmt::Debug, hint::unreachable_unchecked, io
+};
+
+use tokio::{
+    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}
+};
+
+const BUFFER_SIZE: usize = 16 * 1024;
+
+struct Buffer {
+    read: usize,
+    write: usize,
+    buf: Box<[u8]>,
+}
+
+impl Buffer {
+    fn new() -> Self {
+        Self {
+            read: 0,
+            write: 0,
+            buf: vec![0; BUFFER_SIZE].into_boxed_slice(),
+        }
+    }
+
+    fn len(&self) -> usize {
+        self.write - self.read
+    }
+
+    fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    fn available(&self) -> usize {
+        self.buf.len() - self.write
+    }
+
+    fn is_full(&self) -> bool {
+        self.available() == 0
+    }
+
+    fn advance(&mut self, n: usize) {
+        assert!(self.write - self.read >= n);
+        self.read += n;
+        if self.read == self.write {
+            self.read = 0;
+            self.write = 0;
+        }
+    }
+}
+
+pub(crate) struct SafeRead {
+    // the option is only meant for temporary take, it always should be some
+    buffer: Option<Buffer>,
+    status: ReadStatus,
+}
+
+impl Debug for SafeRead {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SafeRead")
+            .field("status", &self.status)
+            .finish()
+    }
+}
+
+#[derive(Debug)]
+enum ReadStatus {
+    Eof,
+    Err(io::Error),
+    Ok,
+}
+
+impl Default for SafeRead {
+    fn default() -> Self {
+        Self {
+            buffer: Some(Buffer::new()),
+            status: ReadStatus::Ok,
+        }
+    }
+}
+
+impl SafeRead {
+    pub(crate) async fn do_io<IO: AsyncRead + Unpin>(&mut self, mut io: IO) -> io::Result<usize> {
+        // if there are some data inside the buffer, just return.
+        let buffer = self.buffer.as_ref().expect("buffer ref expected");
+        if !buffer.is_empty() {
+            return Ok(buffer.len());
+        }
+
+        // read from raw io
+        let buffer = self.buffer.as_mut().expect("buffer ownership expected");
+        let buf = &mut buffer.buf.as_mut()[buffer.write..];
+        let result = io.read(buf).await;
+        match result {
+            Ok(0) => {
+                self.status = ReadStatus::Eof;
+                result
+            }
+            Ok(n) => {
+                buffer.write += n;
+                self.status = ReadStatus::Ok;
+                result
+            }
+            Err(e) => {
+                let rerr = e.kind().into();
+                self.status = ReadStatus::Err(e);
+                Err(rerr)
+            }
+        }
+    }
+}
+
+impl io::Read for SafeRead {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        // if buffer is empty, return WoundBlock.
+        let buffer = self.buffer.as_mut().expect("buffer mut expected");
+        if buffer.is_empty() {
+            if !matches!(self.status, ReadStatus::Ok) {
+                match std::mem::replace(&mut self.status, ReadStatus::Ok) {
+                    ReadStatus::Eof => return Ok(0),
+                    ReadStatus::Err(e) => return Err(e),
+                    ReadStatus::Ok => unsafe { unreachable_unchecked() },
+                }
+            }
+            return Err(io::ErrorKind::WouldBlock.into());
+        }
+
+        // now buffer is not empty. copy it.
+        let to_copy = buffer.len().min(buf.len());
+        unsafe { std::ptr::copy_nonoverlapping(buffer.buf.as_ptr().add(buffer.read), buf.as_mut_ptr(), to_copy) };
+        buffer.advance(to_copy);
+
+        Ok(to_copy)
+    }
+}
+
+pub(crate) struct SafeWrite {
+    // the option is only meant for temporary take, it always should be some
+    buffer: Option<Buffer>,
+    status: WriteStatus,
+}
+
+impl Debug for SafeWrite {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SafeWrite")
+            .field("status", &self.status)
+            .finish()
+    }
+}
+
+#[derive(Debug)]
+enum WriteStatus {
+    Err(io::Error),
+    Ok,
+}
+
+impl Default for SafeWrite {
+    fn default() -> Self {
+        Self {
+            buffer: Some(Buffer::new()),
+            status: WriteStatus::Ok,
+        }
+    }
+}
+
+impl SafeWrite {
+    pub(crate) async fn do_io<IO: AsyncWrite + Unpin>(&mut self, mut io: IO) -> io::Result<usize> {
+        // if the buffer is empty, just return.
+        let buffer = self.buffer.as_ref().expect("buffer ref expected");
+        if buffer.is_empty() {
+            return Ok(0);
+        }
+
+        // buffer is not empty now. write it.
+        let buffer = self.buffer.as_mut().expect("buffer ownership expected");
+        let buf = &buffer.buf.as_ref()[buffer.read..buffer.write];
+        let result = io.write_all(buf).await;
+        match result {
+            Ok(_) => {
+                let n = buffer.write - buffer.read;
+                buffer.advance(n);
+                Ok(n)
+            }
+            Err(e) => {
+                let rerr = e.kind().into();
+                self.status = WriteStatus::Err(e);
+                Err(rerr)
+            }
+        }
+    }
+}
+
+impl io::Write for SafeWrite {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        // if there is too much data inside the buffer, return WoundBlock
+        let buffer = self.buffer.as_mut().expect("buffer mut expected");
+        if !matches!(self.status, WriteStatus::Ok) {
+            match std::mem::replace(&mut self.status, WriteStatus::Ok) {
+                WriteStatus::Err(e) => return Err(e),
+                WriteStatus::Ok => unsafe { unreachable_unchecked() },
+            }
+        }
+        if buffer.is_full() {
+            return Err(io::ErrorKind::WouldBlock.into());
+        }
+
+        // there is space inside the buffer, copy to it.
+        let to_copy = buf.len().min(buffer.available());
+        unsafe { std::ptr::copy_nonoverlapping(buf.as_ptr(), buffer.buf.as_mut_ptr().add(buffer.write), to_copy); }
+        buffer.write += to_copy;
+        Ok(to_copy)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        let buffer = self.buffer.as_mut().expect("buffer mut expected");
+        if !matches!(self.status, WriteStatus::Ok) {
+            match std::mem::replace(&mut self.status, WriteStatus::Ok) {
+                WriteStatus::Err(e) => return Err(e),
+                WriteStatus::Ok => unsafe { unreachable_unchecked() },
+            }
+        }
+        if !buffer.is_empty() {
+            return Err(io::ErrorKind::WouldBlock.into());
+        }
+        Ok(())
+    }
+}

third_party/tokio-rustls-fork-shadow-tls/src/server.rs

@@ -0,0 +1,49 @@
+use std::sync::Arc;
+
+use tokio::io::{AsyncRead, AsyncWrite};
+use rustls_fork_shadow_tls::{ServerConfig, ServerConnection};
+
+use crate::{
+    split::{ReadHalf, WriteHalf},
+    stream::Stream,
+    TlsError,
+};
+
+/// A wrapper around an underlying raw stream which implements the TLS protocol.
+pub type TlsStream<IO> = Stream<IO, ServerConnection>;
+/// TlsStream for read only.
+pub type TlsStreamReadHalf<IO> = ReadHalf<IO, ServerConnection>;
+/// TlsStream for write only.
+pub type TlsStreamWriteHalf<IO> = WriteHalf<IO, ServerConnection>;
+
+/// A wrapper around a `rustls::ServerConfig`, providing an async `accept` method.
+#[derive(Clone)]
+pub struct TlsAcceptor {
+    inner: Arc<ServerConfig>,
+}
+
+impl From<Arc<ServerConfig>> for TlsAcceptor {
+    fn from(inner: Arc<ServerConfig>) -> TlsAcceptor {
+        TlsAcceptor { inner }
+    }
+}
+
+impl From<ServerConfig> for TlsAcceptor {
+    fn from(inner: ServerConfig) -> TlsAcceptor {
+        TlsAcceptor {
+            inner: Arc::new(inner),
+        }
+    }
+}
+
+impl TlsAcceptor {
+    pub async fn accept<IO>(&self, stream: IO) -> Result<TlsStream<IO>, TlsError>
+    where
+        IO: AsyncRead + AsyncWrite + Unpin,
+    {
+        let session = ServerConnection::new(self.inner.clone())?;
+        let mut stream = Stream::new(stream, session);
+        stream.handshake().await?;
+        Ok(stream)
+    }
+}

third_party/tokio-rustls-fork-shadow-tls/src/split.rs

@@ -0,0 +1,139 @@
+//! Split implement for TlsStream.
+//! Note: Here we depends on the behavior of monoio TcpStream.
+//! Though it is not a good assumption, it can really make it
+//! more efficient with less code. The read and write will not
+//! interfere each other.
+use std::{
+    cell::UnsafeCell,
+    future::Future,
+    io::IoSlice,
+    ops::{Deref, DerefMut},
+    pin::Pin,
+    rc::Rc,
+    task::{Context, Poll},
+};
+
+use tokio::{
+    pin,
+    io::{AsyncRead, AsyncWrite, ReadBuf}
+};
+
+use rustls_fork_shadow_tls::{ConnectionCommon, SideData};
+
+use crate::stream::Stream;
+
+#[derive(Debug)]
+pub struct ReadHalf<IO, C> {
+    pub(crate) inner: Rc<UnsafeCell<Stream<IO, C>>>,
+}
+
+#[derive(Debug)]
+pub struct WriteHalf<IO, C> {
+    pub(crate) inner: Rc<UnsafeCell<Stream<IO, C>>>,
+}
+
+impl<IO: AsyncRead + AsyncWrite + Unpin, C, SD: SideData + 'static> AsyncRead
+    for ReadHalf<IO, C>
+where
+    C: DerefMut + Deref<Target = ConnectionCommon<SD>>,
+{
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>
+    ) -> Poll<std::io::Result<()>> {
+        let inner = unsafe { &mut *self.inner.get() };
+        let ex = inner.read_inner(buf, true);
+        pin!(ex);
+        ex.poll(cx)
+    }
+}
+
+impl<IO, C> ReadHalf<IO, C> {
+    pub fn reunite(self, other: WriteHalf<IO, C>) -> Result<Stream<IO, C>, ReuniteError<IO, C>> {
+        reunite(self, other)
+    }
+}
+
+impl<IO: AsyncRead + AsyncWrite + Unpin, C: Unpin, SD: SideData + 'static> AsyncWrite
+    for WriteHalf<IO, C>
+where
+    C: DerefMut + Deref<Target = ConnectionCommon<SD>>,
+{
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8]
+    ) -> Poll<std::io::Result<usize>> {
+        let inner = unsafe { &mut *self.inner.get() };
+        Pin::new(inner).poll_write(cx, buf)
+    }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[IoSlice<'_>]
+    ) -> Poll<std::io::Result<usize>> {
+        let inner = unsafe { &mut *self.inner.get() };
+        Pin::new(inner).poll_write_vectored(cx, bufs)
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>
+    ) -> Poll<std::io::Result<()>> {
+        let inner = unsafe { &mut *self.inner.get() };
+        Pin::new(inner).poll_flush(cx)
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>
+    ) -> Poll<std::io::Result<()>> {
+        let inner = unsafe { &mut *self.inner.get() };
+        Pin::new(inner).poll_shutdown(cx)
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        let inner = unsafe { &mut *self.inner.get() };
+        Pin::new(inner).is_write_vectored()
+    }
+}
+
+impl<IO, C> WriteHalf<IO, C> {
+    pub fn reunite(self, other: ReadHalf<IO, C>) -> Result<Stream<IO, C>, ReuniteError<IO, C>> {
+        reunite(other, self)
+    }
+}
+
+pub(crate) fn reunite<IO, C>(
+    read: ReadHalf<IO, C>,
+    write: WriteHalf<IO, C>,
+) -> Result<Stream<IO, C>, ReuniteError<IO, C>> {
+    if Rc::ptr_eq(&read.inner, &write.inner) {
+        drop(write);
+        // This unwrap cannot fail as the api does not allow creating more than two Rcs,
+        // and we just dropped the other half.
+        Ok(Rc::try_unwrap(read.inner)
+            .expect("TlsStream: try_unwrap failed in reunite")
+            .into_inner())
+    } else {
+        Err(ReuniteError(read, write))
+    }
+}
+
+/// Error indicating that two halves were not from the same socket, and thus could
+/// not be reunited.
+#[derive(Debug)]
+pub struct ReuniteError<IO, C>(pub ReadHalf<IO, C>, pub WriteHalf<IO, C>);
+
+impl<IO: std::fmt::Debug, C: std::fmt::Debug> std::fmt::Display for ReuniteError<IO, C> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "tried to reunite halves that are not from the same socket"
+        )
+    }
+}
+
+impl<IO: std::fmt::Debug, C: std::fmt::Debug> std::error::Error for ReuniteError<IO, C> {}

third_party/tokio-rustls-fork-shadow-tls/src/stream.rs

@@ -0,0 +1,297 @@
+use std::{
+    cell::UnsafeCell,
+    future::Future,
+    io::{IoSlice, Read, self, Write},
+    ops::{Deref, DerefMut},
+    pin::Pin,
+    rc::Rc,
+    task::{Context, Poll},
+};
+
+use tokio::{
+    pin,
+    io::{AsyncRead, AsyncWrite, ReadBuf}
+};
+
+use rustls_fork_shadow_tls::{ConnectionCommon, SideData};
+
+use crate::split::{ReadHalf, WriteHalf};
+
+#[derive(Debug)]
+pub struct Stream<IO, C> {
+    pub(crate) io: IO,
+    pub(crate) session: C,
+    #[cfg(not(feature = "unsafe_io"))]
+    r_buffer: crate::safe_io::SafeRead,
+    #[cfg(not(feature = "unsafe_io"))]
+    w_buffer: crate::safe_io::SafeWrite,
+    #[cfg(feature = "unsafe_io")]
+    r_buffer: crate::unsafe_io::UnsafeRead,
+    #[cfg(feature = "unsafe_io")]
+    w_buffer: crate::unsafe_io::UnsafeWrite,
+}
+
+impl<IO, C> Stream<IO, C> {
+    pub fn new(io: IO, session: C) -> Self {
+        Self {
+            io,
+            session,
+            r_buffer: Default::default(),
+            w_buffer: Default::default(),
+        }
+    }
+
+    pub fn split(self) -> (ReadHalf<IO, C>, WriteHalf<IO, C>) {
+        let shared = Rc::new(UnsafeCell::new(self));
+        (
+            ReadHalf {
+                inner: shared.clone(),
+            },
+            WriteHalf { inner: shared },
+        )
+    }
+
+    pub fn into_parts(self) -> (IO, C) {
+        (self.io, self.session)
+    }
+}
+
+impl<IO: AsyncRead + AsyncWrite + Unpin, C, SD: SideData> Stream<IO, C>
+where
+    C: DerefMut + Deref<Target = ConnectionCommon<SD>>,
+{
+    pub(crate) async fn read_io(&mut self, splitted: bool) -> io::Result<usize> {
+        let n = loop {
+            match self.session.read_tls(&mut self.r_buffer) {
+                Ok(n) => {
+                    break n;
+                }
+                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => (),
+                Err(err) => return Err(err),
+            }
+            #[allow(unused_unsafe)]
+            unsafe {
+                self.r_buffer.do_io(&mut self.io).await?
+            };
+        };
+
+        let state = match self.session.process_new_packets() {
+            Ok(state) => state,
+            Err(err) => {
+                // When to write_io? If we do this in read call, the UnsafeWrite may crash
+                // when we impl split in an UnsafeCell way.
+                // Here we choose not to do write when read.
+                // User should manually shutdown it on error.
+                if !splitted {
+                    let _ = self.write_io().await;
+                }
+                return Err(io::Error::new(io::ErrorKind::InvalidData, err));
+            }
+        };
+
+        if state.peer_has_closed() && self.session.is_handshaking() {
+            return Err(io::Error::new(
+                io::ErrorKind::UnexpectedEof,
+                "tls handshake alert",
+            ));
+        }
+
+        Ok(n)
+    }
+
+    pub(crate) async fn write_io(&mut self) -> io::Result<usize> {
+        let n = loop {
+            match self.session.write_tls(&mut self.w_buffer) {
+                Ok(n) => {
+                    break n;
+                }
+                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => (),
+                Err(err) => return Err(err),
+            }
+            #[allow(unused_unsafe)]
+            unsafe {
+                self.w_buffer.do_io(&mut self.io).await?
+            };
+        };
+        // Flush buffered data, only needed for safe_io.
+        #[cfg(not(feature = "unsafe_io"))]
+        self.w_buffer.do_io(&mut self.io).await?;
+
+        Ok(n)
+    }
+
+    pub(crate) async fn handshake(&mut self) -> io::Result<(usize, usize)> {
+        let mut wrlen = 0;
+        let mut rdlen = 0;
+        let mut eof = false;
+
+        loop {
+            while self.session.wants_write() && self.session.is_handshaking() {
+                wrlen += self.write_io().await?;
+            }
+            while !eof && self.session.wants_read() && self.session.is_handshaking() {
+                let n = self.read_io(false).await?;
+                rdlen += n;
+                if n == 0 {
+                    eof = true;
+                }
+            }
+
+            match (eof, self.session.is_handshaking()) {
+                (true, true) => {
+                    let err = io::Error::new(io::ErrorKind::UnexpectedEof, "tls handshake eof");
+                    return Err(err);
+                }
+                (false, true) => (),
+                (_, false) => {
+                    break;
+                }
+            }
+        }
+
+        // flush buffer
+        while self.session.wants_write() {
+            wrlen += self.write_io().await?;
+        }
+
+        Ok((rdlen, wrlen))
+    }
+
+    pub(crate) async fn read_inner(
+        &mut self,
+        buf: &mut ReadBuf<'_>,
+        splitted: bool,
+    ) -> std::io::Result<()> {
+        if buf.remaining() == 0 {
+            return Ok(());
+        }
+        let slice = buf.initialize_unfilled();
+        loop {
+            // read from rustls to buffer
+            match self.session.reader().read(slice) {
+                Ok(n) => {
+                    buf.advance(n);
+                    return Ok(());
+                }
+                // we need more data, read something.
+                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => (),
+                Err(e) => {
+                    return Err(e);
+                }
+            }
+
+            // now we need data, read something into rustls
+            match self.read_io(splitted).await {
+                Ok(0) => {
+                    return
+                        Err(io::Error::new(
+                            io::ErrorKind::UnexpectedEof,
+                            "tls raw stream eof",
+                        ),
+                    );
+                }
+                Ok(_) => (),
+                Err(e) => {
+                    return Err(e);
+                }
+            }
+        }
+    }
+}
+
+impl<IO: AsyncRead + AsyncWrite + Unpin, C, SD: SideData + 'static> AsyncRead for Stream<IO, C>
+where
+    C: DerefMut + Deref<Target = ConnectionCommon<SD>> + Unpin,
+{
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>
+    ) -> Poll<std::io::Result<()>> {
+        let ex = self.read_inner(buf, false);
+        pin!(ex);
+        ex.poll(cx)
+    }
+}
+
+impl<IO: AsyncRead + AsyncWrite + Unpin, C, SD: SideData + 'static> AsyncWrite for Stream<IO, C>
+where
+    C: DerefMut + Deref<Target = ConnectionCommon<SD>> + Unpin,
+{
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8]
+    ) -> Poll<std::io::Result<usize>> {
+        // write buf to rustls
+        let n = match self.session.writer().write(buf) {
+            Ok(n) => n,
+            Err(e) => return Poll::Ready(Err(e)),
+        };
+
+        // write from rustls to connection
+        while self.session.wants_write() {
+            let ex = self.write_io();
+            pin!(ex);
+            match ex.poll(cx) {
+                Poll::Ready(Ok(0)) => {
+                    break;
+                }
+                Poll::Ready(Ok(_)) => (),
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
+            }
+        }
+        Poll::Ready(Ok(n))
+    }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[IoSlice<'_>]
+    ) -> Poll<std::io::Result<usize>> {
+        let buf = bufs
+            .iter()
+            .find(|b| !b.is_empty())
+            .map_or(&[][..], |b| &**b);
+        self.poll_write(cx, buf)
+    }
+
+    fn poll_flush(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>
+    ) -> Poll<std::io::Result<()>> {
+        self.session.writer().flush()?;
+        while self.session.wants_write() {
+            let ex = self.write_io();
+            pin!(ex);
+            match ex.poll(cx) {
+                Poll::Ready(Ok(_)) => (),
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
+            }
+        }
+        Pin::new(&mut self.io).poll_flush(cx)
+    }
+
+    fn poll_shutdown(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>
+    ) -> Poll<std::io::Result<()>> {
+        self.session.send_close_notify();
+        while self.session.wants_write() {
+            let ex = self.write_io();
+            pin!(ex);
+            match ex.poll(cx) {
+                Poll::Ready(Ok(_)) => (),
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
+            }
+        }
+        Pin::new(&mut self.io).poll_shutdown(cx)
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        Pin::new(&self.io).is_write_vectored()
+    }
+}

third_party/tokio-rustls-fork-shadow-tls/src/unsafe_io.rs

@@ -0,0 +1,124 @@
+use std::{
+    io,
+    slice::{from_raw_parts, from_raw_parts_mut}
+};
+
+use tokio::{
+    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}
+};
+
+/// Used by both UnsafeRead and UnsafeWrite.
+#[derive(Debug)]
+enum Status {
+    /// We haven't do real io, and maybe the dest is recorded.
+    WaitFill(Option<(*const u8, usize)>),
+    /// We have already do real io. The length maybe zero or non-zero.
+    Filled(usize),
+}
+
+impl Default for Status {
+    fn default() -> Self {
+        Status::WaitFill(None)
+    }
+}
+
+/// UnsafeRead is a wrapper of some meta data.
+/// It implements std::io::Read trait. But it do real io in an async way.
+/// On the first read, it may returns WouldBlock error, which means the
+/// `fullfill` should be called to do real io.
+/// The data is read directly into the dest that last std read passes.
+/// Note that this action is an unsafe hack to avoid data copy.
+/// You can only use this wrapper when you make sure the read dest is always
+/// a valid buffer.
+#[derive(Default, Debug)]
+pub(crate) struct UnsafeRead {
+    status: Status,
+}
+
+impl UnsafeRead {
+    /// `do_io` must be called after calling to io::Read::read.
+    pub(crate) async unsafe fn do_io<IO: AsyncRead + Unpin>(
+        &mut self,
+        mut io: IO,
+    ) -> io::Result<usize> {
+        match self.status {
+            Status::WaitFill(Some((ptr, len))) => {
+                let buf = unsafe { from_raw_parts_mut(ptr as *mut u8, len) };
+                let n = io.read(buf).await?;
+                self.status = Status::Filled(n);
+                Ok(n)
+            }
+            Status::Filled(len) => Ok(len),
+            Status::WaitFill(None) => Err(io::ErrorKind::WouldBlock.into()),
+        }
+    }
+}
+
+impl io::Read for UnsafeRead {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        match self.status {
+            Status::WaitFill(_) => {
+                let ptr = buf.as_ptr();
+                let len = buf.len();
+                self.status = Status::WaitFill(Some((ptr, len)));
+                Err(io::ErrorKind::WouldBlock.into())
+            }
+            Status::Filled(len) => {
+                if len != 0 {
+                    // reset only when not eof
+                    self.status = Status::WaitFill(None);
+                }
+                Ok(len)
+            }
+        }
+    }
+}
+
+/// UnsafeWrite behaves like `UnsafeRead`.
+#[derive(Default, Debug)]
+pub(crate) struct UnsafeWrite {
+    status: Status,
+}
+
+impl UnsafeWrite {
+    /// `do_io` must be called after calling to io::Write::write.
+    pub(crate) async unsafe fn do_io<IO: AsyncWrite + Unpin>(
+        &mut self,
+        mut io: IO,
+    ) -> io::Result<usize> {
+        match self.status {
+            Status::WaitFill(Some((ptr, len))) => {
+                let buf = unsafe { from_raw_parts(ptr, len) };
+                let n = io.write(buf).await?;
+                self.status = Status::Filled(n);
+                Ok(n)
+            }
+            Status::Filled(len) => Ok(len),
+            Status::WaitFill(None) => Err(io::ErrorKind::WouldBlock.into()),
+        }
+    }
+}
+
+impl io::Write for UnsafeWrite {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        match self.status {
+            Status::WaitFill(_) => {
+                let ptr = buf.as_ptr();
+                let len = buf.len();
+                self.status = Status::WaitFill(Some((ptr, len)));
+                Err(io::ErrorKind::WouldBlock.into())
+            }
+            Status::Filled(len) => {
+                if len != 0 {
+                    // reset only when not eof
+                    self.status = Status::WaitFill(None);
+                }
+                Ok(len)
+            }
+        }
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}