aquatic_ws: network: streamline handshake handlers

This commit is contained in:
Joakim Frostegård 2020-05-13 18:44:10 +02:00
parent 6d06a43db8
commit 2967129c1f

View file

@ -137,49 +137,29 @@ fn accept_new_streams(
pub fn handle_tls_handshake_result( pub fn handle_tls_handshake_result(
connections: &mut ConnectionMap,
poll_token: Token,
valid_until: ValidUntil,
result: Result<TlsStream<TcpStream>, ::native_tls::HandshakeError<TcpStream>>, result: Result<TlsStream<TcpStream>, ::native_tls::HandshakeError<TcpStream>>,
) -> bool { ) -> (Option<ConnectionStage>, bool) {
match result { match result {
Ok(stream) => { Ok(stream) => {
println!("handshake established"); println!("handshake established");
let connection = Connection { (Some(ConnectionStage::TlsStream(stream)), false)
valid_until,
stage: ConnectionStage::TlsStream(stream)
};
connections.insert(poll_token, connection);
}, },
Err(native_tls::HandshakeError::WouldBlock(handshake)) => { Err(native_tls::HandshakeError::WouldBlock(handshake)) => {
println!("interrupted"); println!("interrupted");
let connection = Connection { (Some(ConnectionStage::TlsMidHandshake(handshake)), true)
valid_until,
stage: ConnectionStage::TlsMidHandshake(handshake),
};
connections.insert(poll_token, connection);
return true;
}, },
Err(native_tls::HandshakeError::Failure(err)) => { Err(native_tls::HandshakeError::Failure(err)) => {
dbg!(err); (None, false)
} }
} }
false
} }
pub fn handle_ws_handshake_result( pub fn handle_ws_handshake_result(
connections: &mut ConnectionMap,
poll_token: Token,
valid_until: ValidUntil,
result: Result<WebSocket<Stream>, HandshakeError<ServerHandshake<Stream, DebugCallback>>> , result: Result<WebSocket<Stream>, HandshakeError<ServerHandshake<Stream, DebugCallback>>> ,
) -> bool { ) -> (Option<ConnectionStage>, bool) {
match result { match result {
Ok(mut ws) => { Ok(mut ws) => {
println!("handshake established"); println!("handshake established");
@ -191,31 +171,19 @@ pub fn handle_ws_handshake_result(
peer_addr, peer_addr,
}; };
let connection = Connection { (Some(ConnectionStage::EstablishedWs(established_ws)), false)
valid_until,
stage: ConnectionStage::EstablishedWs(established_ws)
};
connections.insert(poll_token, connection);
}, },
Err(HandshakeError::Interrupted(handshake)) => { Err(HandshakeError::Interrupted(handshake)) => {
println!("interrupted"); println!("interrupted");
let connection = Connection { (Some(ConnectionStage::WsMidHandshake(handshake)), true)
valid_until,
stage: ConnectionStage::WsMidHandshake(handshake),
};
connections.insert(poll_token, connection);
return true
}, },
Err(HandshakeError::Failure(err)) => { Err(HandshakeError::Failure(err)) => {
dbg!(err); dbg!(err);
(None, false)
} }
} }
false
} }
@ -264,13 +232,10 @@ pub fn run_handshakes_and_read_messages(
} }
} }
} else if let Some(connection) = connections.remove(&poll_token) { } else if let Some(connection) = connections.remove(&poll_token) {
let stop_loop = match connection.stage { let (opt_new_stage, stop_loop) = match connection.stage {
ConnectionStage::TcpStream(stream) => { ConnectionStage::TcpStream(stream) => {
if let Some(tls_acceptor) = opt_tls_acceptor { if let Some(tls_acceptor) = opt_tls_acceptor {
handle_tls_handshake_result( handle_tls_handshake_result(
connections,
poll_token,
valid_until,
tls_acceptor.accept(stream) tls_acceptor.accept(stream)
) )
} else { } else {
@ -279,12 +244,7 @@ pub fn run_handshakes_and_read_messages(
DebugCallback DebugCallback
); );
handle_ws_handshake_result( handle_ws_handshake_result(handshake_result)
connections,
poll_token,
valid_until,
handshake_result
)
} }
}, },
ConnectionStage::TlsStream(stream) => { ConnectionStage::TlsStream(stream) => {
@ -293,32 +253,26 @@ pub fn run_handshakes_and_read_messages(
DebugCallback DebugCallback
); );
handle_ws_handshake_result( handle_ws_handshake_result(handshake_result)
connections,
poll_token,
valid_until,
handshake_result
)
}, },
ConnectionStage::TlsMidHandshake(handshake) => { ConnectionStage::TlsMidHandshake(handshake) => {
handle_tls_handshake_result( handle_tls_handshake_result(handshake.handshake())
connections,
poll_token,
valid_until,
handshake.handshake()
)
}, },
ConnectionStage::WsMidHandshake(handshake) => { ConnectionStage::WsMidHandshake(handshake) => {
handle_ws_handshake_result( handle_ws_handshake_result(handshake.handshake())
connections,
poll_token,
valid_until,
handshake.handshake()
)
}, },
ConnectionStage::EstablishedWs(_) => unreachable!(), ConnectionStage::EstablishedWs(_) => unreachable!(),
}; };
if let Some(stage) = opt_new_stage {
let connection = Connection {
valid_until,
stage,
};
connections.insert(poll_token, connection);
}
if stop_loop { if stop_loop {
break; break;
} }