aquatic_ws: don't call unwrap on stream.get_peer_addr

This commit is contained in:
Joakim Frostegård 2021-10-16 00:27:52 +02:00
parent 2e68155bf4
commit d1ee18394d

View file

@ -21,10 +21,10 @@ pub enum Stream {
impl Stream { impl Stream {
#[inline] #[inline]
pub fn get_peer_addr(&self) -> SocketAddr { pub fn get_peer_addr(&self) -> ::std::io::Result<SocketAddr> {
match self { match self {
Self::TcpStream(stream) => stream.peer_addr().unwrap(), Self::TcpStream(stream) => stream.peer_addr(),
Self::TlsStream(stream) => stream.get_ref().peer_addr().unwrap(), Self::TlsStream(stream) => stream.get_ref().peer_addr(),
} }
} }
@ -162,9 +162,8 @@ impl HandshakeMachine {
result: Result<WebSocket<Stream>, HandshakeError<ServerHandshake<Stream, NoCallback>>>, result: Result<WebSocket<Stream>, HandshakeError<ServerHandshake<Stream, NoCallback>>>,
) -> (Option<Either<EstablishedWs, Self>>, bool) { ) -> (Option<Either<EstablishedWs, Self>>, bool) {
match result { match result {
Ok(mut ws) => { Ok(mut ws) => match ws.get_mut().get_peer_addr() {
let peer_addr = ws.get_mut().get_peer_addr(); Ok(peer_addr) => {
::log::trace!( ::log::trace!(
"established ws handshake with peer with addr: {:?}", "established ws handshake with peer with addr: {:?}",
peer_addr peer_addr
@ -174,6 +173,15 @@ impl HandshakeMachine {
(Some(Either::Left(established_ws)), false) (Some(Either::Left(established_ws)), false)
} }
Err(err) => {
::log::info!(
"get_peer_addr failed during handshake, removing connection: {:?}",
err
);
(None, false)
}
},
Err(HandshakeError::Interrupted(handshake)) => ( Err(HandshakeError::Interrupted(handshake)) => (
Some(Either::Right(HandshakeMachine::WsMidHandshake(handshake))), Some(Either::Right(HandshakeMachine::WsMidHandshake(handshake))),
true, true,