Allow access to underlying connection after handshake

Author: mmastrac

src/connection_stream.rs

@@ -104,6 +104,10 @@ impl ConnectionStream {
     (Arc::try_unwrap(self.tcp).unwrap(), self.tls)
   }
 
+  pub fn connection(&self) -> &Connection {
+    &self.tls
+  }
+
   pub(crate) fn tcp_stream(&self) -> &Arc<TcpStream> {
     &self.tcp
   }

src/stream.rs

@@ -433,6 +433,14 @@ impl TlsStream {
     (read, write)
   }
 
+  /// If the stream is open, returns the underlying rustls connection.
+  pub fn connection(&self) -> Option<&rustls::Connection> {
+    match &self.state {
+      TlsStreamState::Open(stm) => Some(stm.connection()),
+      _ => None,
+    }
+  }
+
   pub async fn into_inner(mut self) -> io::Result<(TcpStream, Connection)> {
     poll_fn(|cx| self.poll_pending_handshake(cx)).await?;
     match std::mem::replace(&mut self.state, TlsStreamState::Closed) {
@@ -489,12 +497,14 @@ impl TlsStream {
     &mut self,
     cx: &mut Context,
   ) -> Poll<io::Result<TlsHandshake>> {
+    // Transition to the open state if necessary
+    ready!(self.poll_pending_handshake(cx)?);
+
     // TODO(mmastrac): Handshake shouldn't need to be cloned
     match &*self.handshake.handshake.lock().unwrap() {
       None => {
         // Register both wakers just in case we get split
         self.handshake.rx_waker.register(cx.waker());
-        self.handshake.tx_waker.register(cx.waker());
         Poll::Pending
       }
       Some(handshake) => Poll::Ready(clone_result(handshake)),
@@ -1619,6 +1629,25 @@ pub(super) mod tests {
     Ok(())
   }
 
+  /// Test that the handshake fails, and we get the correct errors on both ends.
+  #[tokio::test]
+  #[ntest::timeout(60000)]
+  async fn test_client_server_raw_connection() -> TestResult {
+    let (mut server, mut client) =
+      tls_pair_alpn(&["a"], None, &["a"], None).await;
+
+    assert!(server.connection().is_none());
+    assert!(client.connection().is_none());
+
+    server.handshake().await?;
+    client.handshake().await?;
+
+    assert!(server.connection().is_some());
+    assert!(client.connection().is_some());
+
+    Ok(())
+  }
+
   #[tokio::test]
   async fn test_peer_and_local_addresses() {
     let (server, client) = tls_pair_slow_handshake(true, true, true).await;