aquatic_ws: move more code into Connection impl

aquatic_ws/src/lib/network/common.rs

@@ -73,12 +73,6 @@ impl Write for Stream {
 }
 
 
-pub struct EstablishedWs {
-    pub ws: WebSocket<Stream>,
-    pub peer_addr: SocketAddr,
-}
-
-
 pub enum HandshakeMachine {
     TcpStream(TcpStream),
     TlsStream(TlsStream<TcpStream>),
@@ -181,27 +175,31 @@ impl HandshakeMachine {
 }
 
 
+pub struct EstablishedWs {
+    pub ws: WebSocket<Stream>,
+    pub peer_addr: SocketAddr,
+}
+
+
 pub struct Connection {
     pub valid_until: ValidUntil,
     inner: Either<EstablishedWs, HandshakeMachine>,
 }
 
 
+/// Create from TcpStream. Run `advance_handshakes` until `get_established_ws`
+/// returns Some(EstablishedWs).
 impl Connection {
     pub fn new(
         valid_until: ValidUntil,
-        inner: Either<EstablishedWs, HandshakeMachine>
+        tcp_stream: TcpStream,
     ) -> Self {
         Self {
             valid_until,
-            inner
+            inner: Either::Right(HandshakeMachine::TcpStream(tcp_stream))
         }
     }
 
-    pub fn is_established(&self) -> bool {
-        self.inner.is_left()
-    }
-
     pub fn get_established_ws<'a>(&mut self) -> Option<&mut EstablishedWs> {
         match self.inner {
             Either::Left(ref mut ews) => Some(ews),
@@ -209,10 +207,23 @@ impl Connection {
         }
     }
 
-    pub fn get_machine(self) -> Option<HandshakeMachine> {
+    pub fn advance_handshakes(
+        self,
+        opt_tls_acceptor: &Option<TlsAcceptor>,
+        valid_until: ValidUntil,
+    ) -> (Option<Self>, bool) {
         match self.inner {
-            Either::Left(_) => None,
-            Either::Right(machine) => Some(machine),
+            Either::Left(_) => (Some(self), false),
+            Either::Right(machine) => {
+                let (opt_inner, stop_loop) = machine.advance(opt_tls_acceptor);
+
+                let opt_new_self = opt_inner.map(|inner| Self {
+                    valid_until,
+                    inner
+                });
+
+                (opt_new_self, stop_loop)
+            }
         }
     }
 

aquatic_ws/src/lib/network/mod.rs

@@ -116,10 +116,7 @@ fn accept_new_streams(
                     .register(&mut stream, token, Interest::READABLE)
                     .unwrap();
 
-                let connection = Connection::new(
-                    valid_until,
-                    Either::Right(HandshakeMachine::new(stream))
-                );
+                let connection = Connection::new(valid_until, stream);
 
                 connections.insert(token, connection);
             },
@@ -178,15 +175,13 @@ pub fn run_handshakes_and_read_messages(
                     break;
                 }
             }
-        } else if let Some(machine) = connections.remove(&poll_token)
-            .and_then(Connection::get_machine)
-        {
-            let (result, stop_loop) = machine
-                .advance(opt_tls_acceptor);
-
-            if let Some(inner) = result {
-                let connection = Connection::new(valid_until, inner);
+        } else if let Some(connection) = connections.remove(&poll_token){
+            let (opt_new_connection, stop_loop) = connection.advance_handshakes(
+                opt_tls_acceptor,
+                valid_until
+            );
 
+            if let Some(connection) = opt_new_connection {
                 connections.insert(poll_token, connection);
             }