From 73b1646c7189b16561e4e8e4276a49f90ff6bf0e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Sat, 4 Jul 2020 13:10:57 +0200
Subject: [PATCH] aquatic_http: refactor TlsHandshakeMachine, adding error type

---
 aquatic_http/src/lib/network/connection.rs | 37 +++++++++++++---------
 aquatic_http/src/lib/network/mod.rs        | 33 +++++++++++++------
 2 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/aquatic_http/src/lib/network/connection.rs b/aquatic_http/src/lib/network/connection.rs
index 4aa6dae..afaad67 100644
--- a/aquatic_http/src/lib/network/connection.rs
+++ b/aquatic_http/src/lib/network/connection.rs
@@ -5,7 +5,6 @@ use std::sync::Arc;
 
 use either::Either;
 use hashbrown::HashMap;
-use log::info;
 use mio::Token;
 use mio::net::TcpStream;
 use native_tls::{TlsAcceptor, MidHandshakeTlsStream};
@@ -122,7 +121,13 @@ impl EstablishedConnection {
 }
 
 
-enum HandshakeMachineInner {
+pub enum TlsHandshakeMachineError {
+    WouldBlock(TlsHandshakeMachine),
+    Failure(native_tls::Error)
+}
+
+
+enum TlsHandshakeMachineInner {
     TcpStream(TcpStream),
     TlsMidHandshake(MidHandshakeTlsStream<TcpStream>),
 }
@@ -130,7 +135,7 @@ enum HandshakeMachineInner {
 
 pub struct TlsHandshakeMachine {
     tls_acceptor: Arc<TlsAcceptor>,
-    inner: HandshakeMachineInner,
+    inner: TlsHandshakeMachineInner,
 }
 
 
@@ -142,19 +147,19 @@ impl <'a>TlsHandshakeMachine {
     ) -> Self {
         Self {
             tls_acceptor,
-            inner: HandshakeMachineInner::TcpStream(tcp_stream)
+            inner: TlsHandshakeMachineInner::TcpStream(tcp_stream)
         }
     }
 
+    /// Attempt to establish a TLS connection. On a WouldBlock error, return
+    /// the machine wrapped in an error for later attempts.
     #[inline]
-    pub fn advance(
-        self,
-    ) -> (Option<Either<EstablishedConnection, Self>>, bool) { // bool = would block
+    pub fn establish_tls(self) -> Result<EstablishedConnection, TlsHandshakeMachineError> {
         let handshake_result = match self.inner {
-            HandshakeMachineInner::TcpStream(stream) => {
+            TlsHandshakeMachineInner::TcpStream(stream) => {
                 self.tls_acceptor.accept(stream)
             },
-            HandshakeMachineInner::TlsMidHandshake(handshake) => {
+            TlsHandshakeMachineInner::TlsMidHandshake(handshake) => {
                 handshake.handshake()
             },
         };
@@ -165,20 +170,22 @@ impl <'a>TlsHandshakeMachine {
                     Stream::TlsStream(stream)
                 );
 
-                (Some(Either::Left(established)), false)
+                Ok(established)
             },
             Err(native_tls::HandshakeError::WouldBlock(handshake)) => {
+                let inner = TlsHandshakeMachineInner::TlsMidHandshake(
+                    handshake
+                );
+                
                 let machine = Self {
                     tls_acceptor: self.tls_acceptor,
-                    inner: HandshakeMachineInner::TlsMidHandshake(handshake),
+                    inner,
                 };
 
-                (Some(Either::Right(machine)), true)
+                Err(TlsHandshakeMachineError::WouldBlock(machine))
             },
             Err(native_tls::HandshakeError::Failure(err)) => {
-                info!("tls handshake error: {}", err);
-
-                (None, false)
+                Err(TlsHandshakeMachineError::Failure(err))
             }
         }
     }
diff --git a/aquatic_http/src/lib/network/mod.rs b/aquatic_http/src/lib/network/mod.rs
index da8f4c8..7cc5c3c 100644
--- a/aquatic_http/src/lib/network/mod.rs
+++ b/aquatic_http/src/lib/network/mod.rs
@@ -5,6 +5,7 @@ use std::io::ErrorKind;
 use std::sync::Arc;
 use std::vec::Drain;
 
+use either::Either;
 use hashbrown::HashMap;
 use log::{info, debug, error};
 use native_tls::TlsAcceptor;
@@ -247,19 +248,31 @@ pub fn run_handshakes_and_read_requests(
         } else if let Some(handshake_machine) = connections.remove(&poll_token)
             .and_then(|c| c.inner.right())
         {
-            let (opt_inner, would_block) = handshake_machine.advance();
+            match handshake_machine.establish_tls(){
+                Ok(established) => {
+                    let connection = Connection {
+                        valid_until,
+                        inner: Either::Left(established)
+                    };
 
-            if let Some(inner) = opt_inner {
-                let connection = Connection {
-                    valid_until,
-                    inner
-                };
+                    connections.insert(poll_token, connection);
+                },
+                Err(TlsHandshakeMachineError::WouldBlock(machine)) => {
+                    let connection = Connection {
+                        valid_until,
+                        inner: Either::Right(machine)
+                    };
 
-                connections.insert(poll_token, connection);
-            }
+                    connections.insert(poll_token, connection);
 
-            if would_block {
-                break;
+                    break
+                },
+                Err(TlsHandshakeMachineError::Failure(err)) => {
+                    info!("tls handshake error: {}", err);
+
+                    // TLS negotiation error occured
+                    break
+                }
             }
         }
     }