WIP: aquatic_ws: continue work on simplifying network code

This commit is contained in:
Joakim Frostegård 2020-05-13 14:46:38 +02:00
parent 91590858b9
commit 05844e9645
3 changed files with 66 additions and 93 deletions

View file

@ -4,7 +4,7 @@ use std::io::{Read, Write};
use hashbrown::HashMap; use hashbrown::HashMap;
use mio::Token; use mio::Token;
use mio::net::TcpStream; use mio::net::TcpStream;
use native_tls::TlsStream; use native_tls::{TlsStream, MidHandshakeTlsStream};
use tungstenite::WebSocket; use tungstenite::WebSocket;
use tungstenite::handshake::{MidHandshake, server::ServerHandshake}; use tungstenite::handshake::{MidHandshake, server::ServerHandshake};
@ -78,9 +78,10 @@ pub struct EstablishedWs<S> {
pub enum ConnectionStage { pub enum ConnectionStage {
Stream(Stream), TcpStream(TcpStream),
TlsMidHandshake(native_tls::MidHandshakeTlsStream<TcpStream>), TlsStream(TlsStream<TcpStream>),
WsHandshake(MidHandshake<ServerHandshake<Stream, DebugCallback>>), TlsMidHandshake(MidHandshakeTlsStream<TcpStream>),
WsMidHandshake(MidHandshake<ServerHandshake<Stream, DebugCallback>>),
EstablishedWs(EstablishedWs<Stream>), EstablishedWs(EstablishedWs<Stream>),
} }

View file

@ -4,10 +4,10 @@ use std::io::ErrorKind;
use tungstenite::WebSocket; use tungstenite::WebSocket;
use tungstenite::handshake::{HandshakeError, server::ServerHandshake}; use tungstenite::handshake::{HandshakeError, server::ServerHandshake};
use hashbrown::HashMap; use hashbrown::HashMap;
use native_tls::{TlsAcceptor, TlsStream}; use native_tls::TlsAcceptor;
use mio::{Events, Poll, Interest, Token}; use mio::{Events, Poll, Interest, Token};
use mio::net::{TcpListener, TcpStream}; use mio::net::TcpListener;
use crate::common::*; use crate::common::*;
use crate::config::Config; use crate::config::Config;
@ -119,12 +119,10 @@ fn accept_new_streams(
poll.registry() poll.registry()
.register(&mut stream, token, Interest::READABLE) .register(&mut stream, token, Interest::READABLE)
.unwrap(); .unwrap();
let stream = Stream::TcpStream(stream);
let connection = Connection { let connection = Connection {
valid_until, valid_until,
stage: ConnectionStage::Stream(stream) stage: ConnectionStage::TcpStream(stream)
}; };
connections.insert(token, connection); connections.insert(token, connection);
@ -172,7 +170,7 @@ pub fn handle_ws_handshake_result(
let connection = Connection { let connection = Connection {
valid_until, valid_until,
stage: ConnectionStage::WsHandshake(handshake), stage: ConnectionStage::WsMidHandshake(handshake),
}; };
connections.insert(poll_token, connection); connections.insert(poll_token, connection);
@ -188,64 +186,6 @@ pub fn handle_ws_handshake_result(
} }
// Macro hack to not have to write the following twice in
// `run_handshakes_and_read_messages` (putting it in a function causes error
// because of multiple mutable references)
macro_rules! read_ws_messages {
(
$socket_worker_index: ident,
$in_message_sender: ident,
$poll: ident,
$connections: ident,
$poll_token: ident,
$established_ws: ident
) => {
println!("conn established");
match $established_ws.ws.read_message(){
Ok(ws_message) => {
dbg!(ws_message.clone());
if let Some(in_message) = InMessage::from_ws_message(ws_message){
dbg!(in_message.clone());
let meta = ConnectionMeta {
worker_index: $socket_worker_index,
poll_token: $poll_token,
peer_addr: $established_ws.peer_addr
};
$in_message_sender.send((meta, in_message));
}
},
Err(tungstenite::Error::Io(err)) => {
if err.kind() == ErrorKind::WouldBlock {
break;
}
remove_connection_if_exists($poll, $connections, $poll_token);
eprint!("{}", err);
break;
},
Err(tungstenite::Error::ConnectionClosed) => {
remove_connection_if_exists($poll, $connections, $poll_token);
break;
},
Err(err) => {
dbg!(err);
remove_connection_if_exists($poll, $connections, $poll_token);
break;
}
}
};
}
/// Get TLS (if requested) and tungstenite up and running, then read messages /// Get TLS (if requested) and tungstenite up and running, then read messages
pub fn run_handshakes_and_read_messages( pub fn run_handshakes_and_read_messages(
socket_worker_index: usize, socket_worker_index: usize,
@ -259,16 +199,18 @@ pub fn run_handshakes_and_read_messages(
println!("poll_token: {}", poll_token.0); println!("poll_token: {}", poll_token.0);
loop { loop {
let established = match connections.get(&poll_token).map(|c| &c.stage){ let established = if let Some(c) = connections.get(&poll_token){
Some(stage) => stage.is_established(), c.stage.is_established()
None => break, } else {
// Connection is not present, so it is closed. Stop processing
break;
}; };
if !established { if !established {
let conn = connections.remove(&poll_token).unwrap(); let conn = connections.remove(&poll_token).unwrap();
match conn.stage { match conn.stage {
ConnectionStage::Stream(Stream::TcpStream(stream)) => { ConnectionStage::TcpStream(stream) => {
if let Some(tls_acceptor) = opt_tls_acceptor { if let Some(tls_acceptor) = opt_tls_acceptor {
match tls_acceptor.accept(stream){ match tls_acceptor.accept(stream){
Ok(stream) => { Ok(stream) => {
@ -276,7 +218,7 @@ pub fn run_handshakes_and_read_messages(
let connection = Connection { let connection = Connection {
valid_until, valid_until,
stage: ConnectionStage::Stream(Stream::TlsStream(stream)) stage: ConnectionStage::TlsStream(stream)
}; };
connections.insert(poll_token, connection); connections.insert(poll_token, connection);
@ -315,7 +257,7 @@ pub fn run_handshakes_and_read_messages(
} }
} }
}, },
ConnectionStage::Stream(Stream::TlsStream(stream)) => { ConnectionStage::TlsStream(stream) => {
let handshake_result = ::tungstenite::server::accept_hdr( let handshake_result = ::tungstenite::server::accept_hdr(
Stream::TlsStream(stream), Stream::TlsStream(stream),
DebugCallback DebugCallback
@ -339,7 +281,7 @@ pub fn run_handshakes_and_read_messages(
let connection = Connection { let connection = Connection {
valid_until, valid_until,
stage: ConnectionStage::Stream(Stream::TlsStream(stream)) stage: ConnectionStage::TlsStream(stream)
}; };
connections.insert(poll_token, connection); connections.insert(poll_token, connection);
@ -361,7 +303,7 @@ pub fn run_handshakes_and_read_messages(
} }
} }
}, },
ConnectionStage::WsHandshake(handshake) => { ConnectionStage::WsMidHandshake(handshake) => {
let stop_loop = handle_ws_handshake_result( let stop_loop = handle_ws_handshake_result(
connections, connections,
poll_token, poll_token,
@ -375,22 +317,49 @@ pub fn run_handshakes_and_read_messages(
}, },
ConnectionStage::EstablishedWs(_) => unreachable!(), ConnectionStage::EstablishedWs(_) => unreachable!(),
} }
} else { } else if let Some(Connection {
match connections.get_mut(&poll_token){ stage: ConnectionStage::EstablishedWs(established_ws),
Some(Connection{ ..
stage: ConnectionStage::EstablishedWs(established_ws), }) = connections.get_mut(&poll_token){
.. match established_ws.ws.read_message(){
}) => { Ok(ws_message) => {
read_ws_messages!( dbg!(ws_message.clone());
socket_worker_index,
in_message_sender, if let Some(in_message) = InMessage::from_ws_message(ws_message){
poll, dbg!(in_message.clone());
connections,
poll_token, let meta = ConnectionMeta {
established_ws worker_index: socket_worker_index,
); poll_token: poll_token,
peer_addr: established_ws.peer_addr
};
in_message_sender.send((meta, in_message));
}
}, },
_ => () Err(tungstenite::Error::Io(err)) => {
if err.kind() == ErrorKind::WouldBlock {
break;
}
remove_connection_if_exists(poll, connections, poll_token);
eprint!("{}", err);
break;
},
Err(tungstenite::Error::ConnectionClosed) => {
remove_connection_if_exists(poll, connections, poll_token);
break;
},
Err(err) => {
dbg!(err);
remove_connection_if_exists(poll, connections, poll_token);
break;
}
} }
} }
} }

View file

@ -63,13 +63,16 @@ pub fn close_and_deregister_connection(
connection: &mut Connection, connection: &mut Connection,
){ ){
match connection.stage { match connection.stage {
ConnectionStage::Stream(ref mut stream) => { ConnectionStage::TcpStream(ref mut stream) => {
/* /*
poll.registry() poll.registry()
.deregister(stream) .deregister(stream)
.unwrap(); .unwrap();
*/ */
}, },
ConnectionStage::TlsStream(ref mut stream) => {
}
ConnectionStage::TlsMidHandshake(ref mut handshake) => { ConnectionStage::TlsMidHandshake(ref mut handshake) => {
/* /*
poll.registry() poll.registry()
@ -77,7 +80,7 @@ pub fn close_and_deregister_connection(
.unwrap(); .unwrap();
*/ */
}, },
ConnectionStage::WsHandshake(ref mut handshake) => { ConnectionStage::WsMidHandshake(ref mut handshake) => {
/* /*
poll.registry() poll.registry()
.deregister(handshake.get_mut().get_mut()) .deregister(handshake.get_mut().get_mut())