From a3c9ebc7fa4f05e855f47d183fa667db86d2c91e Mon Sep 17 00:00:00 2001
From: fakeshadow <24548779@qq.com>
Date: Wed, 24 Mar 2021 09:32:04 -0700
Subject: [PATCH] fix rustls panic when generating dns name from ip (#296)

* fix rustls panic when generating dns name from ip

* Update rustls.rs

* update changelog

Co-authored-by: Rob Ede <robjtede@icloud.com>
---
 actix-tls/CHANGES.md                |  4 ++
 actix-tls/src/connect/ssl/rustls.rs | 67 +++++++++++++++++------------
 2 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/actix-tls/CHANGES.md b/actix-tls/CHANGES.md
index 824663b0b9..5bf21d4ed6 100644
--- a/actix-tls/CHANGES.md
+++ b/actix-tls/CHANGES.md
@@ -1,6 +1,10 @@
 # Changes
 
 ## Unreleased - 2021-xx-xx
+* Changed `connect::ssl::rustls::RustlsConnectorService` to return error when `DNSNameRef` 
+  generation failed instead of panic. [#296]
+
+[#296]: https://github.com/actix/actix-net/pull/296
 
 
 ## 3.0.0-beta.4 - 2021-02-24
diff --git a/actix-tls/src/connect/ssl/rustls.rs b/actix-tls/src/connect/ssl/rustls.rs
index 46b4b11dd1..b03d4b7fb1 100755
--- a/actix-tls/src/connect/ssl/rustls.rs
+++ b/actix-tls/src/connect/ssl/rustls.rs
@@ -1,6 +1,6 @@
 use std::{
-    fmt,
     future::Future,
+    io,
     pin::Pin,
     sync::Arc,
     task::{Context, Poll},
@@ -10,7 +10,7 @@ pub use tokio_rustls::rustls::Session;
 pub use tokio_rustls::{client::TlsStream, rustls::ClientConfig};
 pub use webpki_roots::TLS_SERVER_ROOTS;
 
-use actix_codec::{AsyncRead, AsyncWrite};
+use actix_rt::net::ActixStream;
 use actix_service::{Service, ServiceFactory};
 use futures_core::{future::LocalBoxFuture, ready};
 use log::trace;
@@ -44,12 +44,13 @@ impl Clone for RustlsConnector {
     }
 }
 
-impl<T: Address, U> ServiceFactory<Connection<T, U>> for RustlsConnector
+impl<T, U> ServiceFactory<Connection<T, U>> for RustlsConnector
 where
-    U: AsyncRead + AsyncWrite + Unpin + fmt::Debug,
+    T: Address,
+    U: ActixStream,
 {
     type Response = Connection<T, TlsStream<U>>;
-    type Error = std::io::Error;
+    type Error = io::Error;
     type Config = ();
     type Service = RustlsConnectorService;
     type InitError = ();
@@ -76,43 +77,55 @@ impl Clone for RustlsConnectorService {
 impl<T, U> Service<Connection<T, U>> for RustlsConnectorService
 where
     T: Address,
-    U: AsyncRead + AsyncWrite + Unpin + fmt::Debug,
+    U: ActixStream,
 {
     type Response = Connection<T, TlsStream<U>>;
-    type Error = std::io::Error;
-    type Future = ConnectAsyncExt<T, U>;
+    type Error = io::Error;
+    type Future = RustlsConnectorServiceFuture<T, U>;
 
     actix_service::always_ready!();
 
-    fn call(&self, stream: Connection<T, U>) -> Self::Future {
-        trace!("SSL Handshake start for: {:?}", stream.host());
-        let (io, stream) = stream.replace_io(());
-        let host = DNSNameRef::try_from_ascii_str(stream.host())
-            .expect("rustls currently only handles hostname-based connections. See https://github.com/briansmith/webpki/issues/54");
-        ConnectAsyncExt {
-            fut: TlsConnector::from(self.connector.clone()).connect(host, io),
-            stream: Some(stream),
+    fn call(&self, connection: Connection<T, U>) -> Self::Future {
+        trace!("SSL Handshake start for: {:?}", connection.host());
+        let (stream, connection) = connection.replace_io(());
+
+        match DNSNameRef::try_from_ascii_str(connection.host()) {
+            Ok(host) => RustlsConnectorServiceFuture::Future {
+                connect: TlsConnector::from(self.connector.clone()).connect(host, stream),
+                connection: Some(connection),
+            },
+            Err(_) => RustlsConnectorServiceFuture::InvalidDns,
         }
     }
 }
 
-pub struct ConnectAsyncExt<T, U> {
-    fut: Connect<U>,
-    stream: Option<Connection<T, ()>>,
+pub enum RustlsConnectorServiceFuture<T, U> {
+    /// See issue https://github.com/briansmith/webpki/issues/54
+    InvalidDns,
+    Future {
+        connect: Connect<U>,
+        connection: Option<Connection<T, ()>>,
+    },
 }
 
-impl<T, U> Future for ConnectAsyncExt<T, U>
+impl<T, U> Future for RustlsConnectorServiceFuture<T, U>
 where
     T: Address,
-    U: AsyncRead + AsyncWrite + Unpin + fmt::Debug,
+    U: ActixStream,
 {
-    type Output = Result<Connection<T, TlsStream<U>>, std::io::Error>;
+    type Output = Result<Connection<T, TlsStream<U>>, io::Error>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let this = self.get_mut();
-        let stream = ready!(Pin::new(&mut this.fut).poll(cx))?;
-        let s = this.stream.take().unwrap();
-        trace!("SSL Handshake success: {:?}", s.host());
-        Poll::Ready(Ok(s.replace_io(stream).1))
+        match self.get_mut() {
+            Self::InvalidDns => Poll::Ready(Err(
+                io::Error::new(io::ErrorKind::Other, "rustls currently only handles hostname-based connections. See https://github.com/briansmith/webpki/issues/54")
+            )),
+            Self::Future { connect, connection } => {
+                let stream = ready!(Pin::new(connect).poll(cx))?;
+                let connection = connection.take().unwrap();
+                trace!("SSL Handshake success: {:?}", connection.host());
+                Poll::Ready(Ok(connection.replace_io(stream).1))
+            }
+        }
     }
 }