gracefultls.rs

struct GracefulTlsStream {
    inner: Arc<Mutex<TlsStream<TcpStream, ClientSession>>>,
}

impl Write for GracefulTlsStream {
    fn write(&mut self, buf: &[u8]) -> Result<usize, std::io::Error> {
        let mut guard = self.inner.lock().unwrap();
        guard.write(buf)
    }

    fn flush(&mut self) -> Result<(), std::io::Error> {
        let mut guard = self.inner.lock().unwrap();
        guard.flush()
    }

    fn write_all(&mut self, buf: &[u8]) -> Result<(), std::io::Error> {
        let mut guard = self.inner.lock().unwrap();
        guard.write_all(buf)
    }

    fn write_fmt(&mut self, fmt: Arguments) -> Result<(), std::io::Error> {
        let mut guard = self.inner.lock().unwrap();
        guard.write_fmt(fmt)
    }
}

impl AsyncWrite for GracefulTlsStream {
    fn poll_write(&mut self, buf: &[u8]) -> Result<Async<usize>, std::io::Error> {
        let mut guard = self.inner.lock().unwrap();
        guard.poll_write(buf)
    }

    fn poll_flush(&mut self) -> Result<Async<()>, std::io::Error> {
        let mut guard = self.inner.lock().unwrap();
        guard.poll_flush()
    }

    fn shutdown(&mut self) -> Result<Async<()>, std::io::Error> {
        let mut guard = self.inner.lock().unwrap();
        guard.shutdown()
    }

    fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Result<Async<usize>, std::io::Error>
    where
        Self: Sized,
    {
        let mut guard = self.inner.lock().unwrap();
        guard.write_buf(buf)
    }
}

impl Read for GracefulTlsStream {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error> {
        let mut guard = self.inner.lock().unwrap();
        guard.read(buf)
    }
}

impl AsyncRead for GracefulTlsStream {}

struct GracefulTlsStreamShutdown {
    inner: Arc<Mutex<TlsStream<TcpStream, ClientSession>>>,
}

impl Future for GracefulTlsStreamShutdown {
    type Item = ();
    type Error = std::io::Error;

    fn poll(&mut self) -> Result<Async<Self::Item>, Self::Error> {
        let mut guard = self.inner.lock().unwrap();
        guard.shutdown()
    }
}

impl Drop for GracefulTlsStream {
    fn drop(&mut self) {
        tokio::spawn(
            GracefulTlsStreamShutdown {
                inner: self.inner.clone(),
            }
            .map_err(|e| eprintln!("could not shutdown TLS stream gracefully: {}!", e)),
        );
    }
}

/* Then in my implementation of tokio_connect::Connect I return my GracefulTlsStream like:
                let connect = tls_connector
                    .connect(domain, sock)
                    .map(|tls_stream| GracefulTlsStream {
                        inner: Arc::new(Mutex::new(tls_stream)),
                    })
*/