From d0af82c54cfd0e54602d66677d471b0f465cbaef Mon Sep 17 00:00:00 2001
From: Jun Kurihara <junkurihara@users.noreply.github.com>
Date: Fri, 8 Jul 2022 12:35:03 +0900
Subject: [PATCH] refactor

---
 src/proxy/backend.rs   | 11 +++++++++--
 src/proxy/proxy_tls.rs | 10 +++++-----
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/proxy/backend.rs b/src/proxy/backend.rs
index 62189e7..18794e0 100644
--- a/src/proxy/backend.rs
+++ b/src/proxy/backend.rs
@@ -99,8 +99,15 @@ impl Backend {
   }
   pub async fn update_server_config(&self) -> io::Result<()> {
     debug!("Update TLS server config");
-    let certs_path = self.tls_cert_path.as_ref().unwrap();
-    let certs_keys_path = self.tls_cert_key_path.as_ref().unwrap();
+    let (certs_path, certs_keys_path) =
+      if let (Some(c), Some(k)) = (self.tls_cert_path.as_ref(), self.tls_cert_key_path.as_ref()) {
+        (c, k)
+      } else {
+        return Err(io::Error::new(
+          io::ErrorKind::Other,
+          "Invalid certs and keys paths",
+        ));
+      };
     let certs: Vec<_> = {
       let certs_path_str = certs_path.display().to_string();
       let mut reader = BufReader::new(File::open(certs_path).map_err(|e| {
diff --git a/src/proxy/proxy_tls.rs b/src/proxy/proxy_tls.rs
index e2eb1fb..c9d8914 100644
--- a/src/proxy/proxy_tls.rs
+++ b/src/proxy/proxy_tls.rs
@@ -34,18 +34,18 @@ where
     loop {
       select! {
         tcp_cnx = tcp_listener.accept().fuse() => {
-          if tcp_cnx.is_err() {
+          // First check SNI
+          let rustls_acceptor = rustls::server::Acceptor::new();
+          if tcp_cnx.is_err() || rustls_acceptor.is_err() {
             continue;
           }
           let (raw_stream, _client_addr) = tcp_cnx.unwrap();
-
-          // First check SNI
-          let rustls_acceptor = rustls::server::Acceptor::new().unwrap();
-          let acceptor = tokio_rustls::LazyConfigAcceptor::new(rustls_acceptor, raw_stream).await;
+          let acceptor = tokio_rustls::LazyConfigAcceptor::new(rustls_acceptor.unwrap(), raw_stream).await;
           if acceptor.is_err() {
             continue;
           }
           let start = acceptor.unwrap();
+
           let client_hello = start.client_hello();
           debug!("SNI in ClientHello: {:?}", client_hello.server_name());
           // Find server config for given SNI