WIP: aquatic_ws; work on simplifying network code

aquatic_ws/src/lib/network/common.rs

@@ -1,4 +1,5 @@
 use std::net::{SocketAddr};
+use std::io::{Read, Write};
 
 use hashbrown::HashMap;
 use mio::Token;
@@ -27,7 +28,47 @@ impl ::tungstenite::handshake::server::Callback for DebugCallback {
 }
 
 
-pub type Stream = TlsStream<TcpStream>;
+pub enum Stream {
+    TcpStream(TcpStream),
+    TlsStream(TlsStream<TcpStream>),
+}
+
+
+impl Stream {
+    pub fn get_peer_addr(&self) -> SocketAddr {
+        match self {
+            Self::TcpStream(stream) => stream.peer_addr().unwrap(),
+            Self::TlsStream(stream) => stream.get_ref().peer_addr().unwrap(),
+        }
+    }
+}
+
+
+impl Read for Stream {
+    fn read(&mut self, buf: &mut [u8]) -> Result<usize, ::std::io::Error> {
+        match self {
+            Self::TcpStream(stream) => stream.read(buf),
+            Self::TlsStream(stream) => stream.read(buf),
+        }
+    }
+}
+
+
+impl Write for Stream {
+    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
+        match self {
+            Self::TcpStream(stream) => stream.write(buf),
+            Self::TlsStream(stream) => stream.write(buf),
+        }
+    }
+
+    fn flush(&mut self) -> ::std::io::Result<()> {
+        match self {
+            Self::TcpStream(stream) => stream.flush(),
+            Self::TlsStream(stream) => stream.flush(),
+        }
+    }
+}
 
 
 pub struct EstablishedWs<S> {
@@ -37,20 +78,17 @@ pub struct EstablishedWs<S> {
 
 
 pub enum ConnectionStage {
-    TcpStream(TcpStream),
+    Stream(Stream),
     TlsMidHandshake(native_tls::MidHandshakeTlsStream<TcpStream>),
-    TlsStream(Stream),
-    WsHandshakeNoTls(MidHandshake<ServerHandshake<TcpStream, DebugCallback>>),
-    WsHandshakeTls(MidHandshake<ServerHandshake<Stream, DebugCallback>>),
-    EstablishedWsNoTls(EstablishedWs<TcpStream>),
-    EstablishedWsTls(EstablishedWs<Stream>),
+    WsHandshake(MidHandshake<ServerHandshake<Stream, DebugCallback>>),
+    EstablishedWs(EstablishedWs<Stream>),
 }
 
 
 impl ConnectionStage {
     pub fn is_established(&self) -> bool {
         match self {
-            Self::EstablishedWsTls(_) | Self::EstablishedWsNoTls(_) => true,
+            Self::EstablishedWs(_) => true,
             _ => false,
         }
     }