From d1ee18394d74c3655f6f67a5ac0ab4c53f762168 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Sat, 16 Oct 2021 00:27:52 +0200
Subject: [PATCH] aquatic_ws: don't call unwrap on stream.get_peer_addr

---
 aquatic_ws/src/lib/network/connection.rs | 32 +++++++++++++++---------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/aquatic_ws/src/lib/network/connection.rs b/aquatic_ws/src/lib/network/connection.rs
index 8873509..4769dc6 100644
--- a/aquatic_ws/src/lib/network/connection.rs
+++ b/aquatic_ws/src/lib/network/connection.rs
@@ -21,10 +21,10 @@ pub enum Stream {
 
 impl Stream {
     #[inline]
-    pub fn get_peer_addr(&self) -> SocketAddr {
+    pub fn get_peer_addr(&self) -> ::std::io::Result<SocketAddr> {
         match self {
-            Self::TcpStream(stream) => stream.peer_addr().unwrap(),
-            Self::TlsStream(stream) => stream.get_ref().peer_addr().unwrap(),
+            Self::TcpStream(stream) => stream.peer_addr(),
+            Self::TlsStream(stream) => stream.get_ref().peer_addr(),
         }
     }
 
@@ -162,18 +162,26 @@ impl HandshakeMachine {
         result: Result<WebSocket<Stream>, HandshakeError<ServerHandshake<Stream, NoCallback>>>,
     ) -> (Option<Either<EstablishedWs, Self>>, bool) {
         match result {
-            Ok(mut ws) => {
-                let peer_addr = ws.get_mut().get_peer_addr();
+            Ok(mut ws) => match ws.get_mut().get_peer_addr() {
+                Ok(peer_addr) => {
+                    ::log::trace!(
+                        "established ws handshake with peer with addr: {:?}",
+                        peer_addr
+                    );
 
-                ::log::trace!(
-                    "established ws handshake with peer with addr: {:?}",
-                    peer_addr
-                );
+                    let established_ws = EstablishedWs { ws, peer_addr };
 
-                let established_ws = EstablishedWs { ws, peer_addr };
+                    (Some(Either::Left(established_ws)), false)
+                }
+                Err(err) => {
+                    ::log::info!(
+                        "get_peer_addr failed during handshake, removing connection: {:?}",
+                        err
+                    );
 
-                (Some(Either::Left(established_ws)), false)
-            }
+                    (None, false)
+                }
+            },
             Err(HandshakeError::Interrupted(handshake)) => (
                 Some(Either::Right(HandshakeMachine::WsMidHandshake(handshake))),
                 true,