Return a custom TlsStream rather than a ChannelBinding up front

postgres-openssl/src/lib.rs

@@ -6,7 +6,7 @@
 //! use openssl::ssl::{SslConnector, SslMethod};
 //! use postgres_openssl::MakeTlsConnector;
 //!
-//! # fn main() -> Result<(), Box<std::error::Error>> {
+//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
 //! let mut builder = SslConnector::builder(SslMethod::tls())?;
 //! builder.set_ca_file("database_cert.pem")?;
 //! let connector = MakeTlsConnector::new(builder.build());
@@ -25,7 +25,7 @@
 //! use openssl::ssl::{SslConnector, SslMethod};
 //! use postgres_openssl::MakeTlsConnector;
 //!
-//! # fn main() -> Result<(), Box<std::error::Error>> {
+//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
 //! let mut builder = SslConnector::builder(SslMethod::tls())?;
 //! builder.set_ca_file("database_cert.pem")?;
 //! let connector = MakeTlsConnector::new(builder.build());
@@ -42,6 +42,8 @@
 #![doc(html_root_url = "https://docs.rs/postgres-openssl/0.3")]
 #![warn(rust_2018_idioms, clippy::all, missing_docs)]
 
+use futures::task::Context;
+use futures::Poll;
 #[cfg(feature = "runtime")]
 use openssl::error::ErrorStack;
 use openssl::hash::MessageDigest;
@@ -51,11 +53,13 @@ use openssl::ssl::SslConnector;
 use openssl::ssl::{ConnectConfiguration, SslRef};
 use std::fmt::Debug;
 use std::future::Future;
+use std::io;
 use std::pin::Pin;
 #[cfg(feature = "runtime")]
 use std::sync::Arc;
-use tokio_io::{AsyncRead, AsyncWrite};
+use tokio_io::{AsyncRead, AsyncWrite, Buf, BufMut};
 use tokio_openssl::{HandshakeError, SslStream};
+use tokio_postgres::tls;
 #[cfg(feature = "runtime")]
 use tokio_postgres::tls::MakeTlsConnect;
 use tokio_postgres::tls::{ChannelBinding, TlsConnect};
@@ -99,7 +103,7 @@ impl<S> MakeTlsConnect<S> for MakeTlsConnector
 where
     S: AsyncRead + AsyncWrite + Unpin + Debug + 'static + Sync + Send,
 {
-    type Stream = SslStream<S>;
+    type Stream = TlsStream<S>;
     type TlsConnect = TlsConnector;
     type Error = ErrorStack;
 
@@ -130,29 +134,96 @@ impl<S> TlsConnect<S> for TlsConnector
 where
     S: AsyncRead + AsyncWrite + Unpin + Debug + 'static + Sync + Send,
 {
-    type Stream = SslStream<S>;
+    type Stream = TlsStream<S>;
     type Error = HandshakeError<S>;
     #[allow(clippy::type_complexity)]
-    type Future = Pin<
-        Box<dyn Future<Output = Result<(SslStream<S>, ChannelBinding), HandshakeError<S>>> + Send>,
-    >;
+    type Future = Pin<Box<dyn Future<Output = Result<TlsStream<S>, HandshakeError<S>>> + Send>>;
 
     fn connect(self, stream: S) -> Self::Future {
         let future = async move {
             let stream = tokio_openssl::connect(self.ssl, &self.domain, stream).await?;
-
-            let channel_binding = match tls_server_end_point(stream.ssl()) {
-                Some(buf) => ChannelBinding::tls_server_end_point(buf),
-                None => ChannelBinding::none(),
-            };
-
-            Ok((stream, channel_binding))
+            Ok(TlsStream(stream))
         };
 
         Box::pin(future)
     }
 }
 
+/// The stream returned by `TlsConnector`.
+pub struct TlsStream<S>(SslStream<S>);
+
+impl<S> AsyncRead for TlsStream<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
+        self.0.prepare_uninitialized_buffer(buf)
+    }
+
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut self.0).poll_read(cx, buf)
+    }
+
+    fn poll_read_buf<B: BufMut>(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut B,
+    ) -> Poll<io::Result<usize>>
+    where
+        Self: Sized,
+    {
+        Pin::new(&mut self.0).poll_read_buf(cx, buf)
+    }
+}
+
+impl<S> AsyncWrite for TlsStream<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut self.0).poll_write(cx, buf)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_flush(cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_shutdown(cx)
+    }
+
+    fn poll_write_buf<B: Buf>(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut B,
+    ) -> Poll<io::Result<usize>>
+    where
+        Self: Sized,
+    {
+        Pin::new(&mut self.0).poll_write_buf(cx, buf)
+    }
+}
+
+impl<S> tls::TlsStream for TlsStream<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    fn channel_binding(&self) -> ChannelBinding {
+        match tls_server_end_point(self.0.ssl()) {
+            Some(buf) => ChannelBinding::tls_server_end_point(buf),
+            None => ChannelBinding::none(),
+        }
+    }
+}
+
 fn tls_server_end_point(ssl: &SslRef) -> Option<Vec<u8>> {
     let cert = ssl.peer_certificate()?;
     let algo_nid = cert.signature_algorithm().object().nid();