From d70aa8feec7f2e4ad791dcde1f1ba019dceab15c Mon Sep 17 00:00:00 2001
From: Dominik Nakamura <dnaka91@gmail.com>
Date: Sat, 6 Feb 2021 12:36:23 +0900
Subject: [PATCH] Move TLS related errors to a separate enum

---
 src/client.rs | 16 +++++++++++-----
 src/error.rs  | 41 ++++++++++++++++++++++++++++-------------
 2 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 096663a..d2cb533 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -26,17 +26,20 @@ mod encryption {
     /// TCP stream switcher (plain/TLS).
     pub type AutoStream = StreamSwitcher<TcpStream, TlsStream<TcpStream>>;
 
-    use crate::{error::Result, stream::Mode};
+    use crate::{
+        error::{Result, TlsError},
+        stream::Mode,
+    };
 
     pub fn wrap_stream(stream: TcpStream, domain: &str, mode: Mode) -> Result<AutoStream> {
         match mode {
             Mode::Plain => Ok(StreamSwitcher::Plain(stream)),
             Mode::Tls => {
-                let connector = TlsConnector::builder().build()?;
+                let connector = TlsConnector::builder().build().map_err(TlsError::Native)?;
                 connector
                     .connect(domain, stream)
                     .map_err(|e| match e {
-                        TlsHandshakeError::Failure(f) => f.into(),
+                        TlsHandshakeError::Failure(f) => TlsError::Native(f).into(),
                         TlsHandshakeError::WouldBlock(_) => {
                             panic!("Bug: TLS handshake not blocked")
                         }
@@ -58,7 +61,10 @@ mod encryption {
     /// TCP stream switcher (plain/TLS).
     pub type AutoStream = StreamSwitcher<TcpStream, StreamOwned<ClientSession, TcpStream>>;
 
-    use crate::{error::Result, stream::Mode};
+    use crate::{
+        error::{Result, TlsError},
+        stream::Mode,
+    };
 
     pub fn wrap_stream(stream: TcpStream, domain: &str, mode: Mode) -> Result<AutoStream> {
         match mode {
@@ -70,7 +76,7 @@ mod encryption {
 
                     Arc::new(config)
                 };
-                let domain = DNSNameRef::try_from_ascii_str(domain)?;
+                let domain = DNSNameRef::try_from_ascii_str(domain).map_err(TlsError::Dns)?;
                 let client = ClientSession::new(&config, domain);
                 let stream = StreamOwned::new(client, stream);
 
diff --git a/src/error.rs b/src/error.rs
index 313c929..2760e8a 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -11,7 +11,6 @@ pub type Result<T> = result::Result<T, Error>;
 
 /// Possible WebSocket errors.
 #[derive(Error, Debug)]
-#[non_exhaustive]
 pub enum Error {
     /// WebSocket connection closed normally. This informs you of the close.
     /// It's not an error as such and nothing wrong happened.
@@ -39,18 +38,12 @@ pub enum Error {
     /// underlying connection and you should probably consider them fatal.
     #[error("IO error: {0}")]
     Io(#[from] io::Error),
-    /// TLS error
-    #[cfg(feature = "native-tls")]
-    #[error("TLS (native-tls) error: {0}")]
-    TlsNative(#[from] native_tls_crate::Error),
-    /// TLS error
-    #[cfg(feature = "rustls-tls")]
-    #[error("TLS (rustls) error: {0}")]
-    TlsRustls(#[from] rustls::TLSError),
-    /// DNS name resolution error.
-    #[cfg(feature = "rustls-tls")]
-    #[error("Invalid DNS name: {0}")]
-    Dns(#[from] webpki::InvalidDNSNameError),
+    /// TLS error.
+    ///
+    /// Note that this error variant is enabled unconditionally even if no TLS feature is enabled,
+    /// to provide a feature-agnostic API surface.
+    #[error("TLS error: {0}")]
+    TlsNative(#[from] TlsError),
     /// - When reading: buffer capacity exhausted.
     /// - When writing: your message is bigger than the configured max message size
     ///   (64MB by default).
@@ -251,3 +244,25 @@ pub enum UrlError {
     #[error("No path/query in URL")]
     NoPathOrQuery,
 }
+
+/// TLS errors.
+///
+/// Note that even if you enable only the rustls-based TLS support, the error at runtime could still
+/// be `Native`, as another crate in the dependency graph may enable native TLS support.
+#[allow(missing_copy_implementations)]
+#[derive(Error, Debug)]
+#[non_exhaustive]
+pub enum TlsError {
+    /// Native TLS error.
+    #[cfg(feature = "native-tls")]
+    #[error("native-tls error: {0}")]
+    Native(#[from] native_tls_crate::Error),
+    /// Rustls error.
+    #[cfg(feature = "rustls-tls")]
+    #[error("rustls error: {0}")]
+    Rustls(#[from] rustls::TLSError),
+    /// DNS name resolution error.
+    #[cfg(feature = "rustls-tls")]
+    #[error("Invalid DNS name: {0}")]
+    Dns(#[from] webpki::InvalidDNSNameError),
+}