From e4cc7cb24ad570293ef7345212090569f32c4f01 Mon Sep 17 00:00:00 2001
From: Hamcha <hamcha@crunchy.rocks>
Date: Mon, 5 Oct 2020 17:17:44 +0200
Subject: [PATCH] Less fragility

---
 src/network.rs | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/src/network.rs b/src/network.rs
index 364a323..913766b 100644
--- a/src/network.rs
+++ b/src/network.rs
@@ -49,13 +49,11 @@ impl NetworkManager {
 			.connections
 			.get_mut(&message.conn_id)
 			.expect("cant send message to an unregistered connection");
-		match message.data {
-			NetworkPayload::Message(msg) => conn
-				.stream
-				.send(msg.into())
-				.await
-				.expect("could not send message to connection"),
-			_ => (),
+		if let NetworkPayload::Message(msg) = message.data {
+			let res = conn.stream.send(msg.into()).await;
+			if let Err(err) = res {
+				log::error!("Error sending message over to game instance: {:?}", err)
+			}
 		}
 	}
 }
@@ -74,8 +72,12 @@ pub async fn listen(
 			.peer_addr()
 			.expect("connected streams should have a peer address");
 		log::info!("Peer address: {}", peer);
-		let conn = accept_connection(peer, stream, incoming.clone()).await;
-		network.lock().await.register(conn);
+		match accept_connection(peer, stream, incoming.clone()).await {
+			Ok(conn) => {
+				network.lock().await.register(conn);
+			}
+			Err(err) => log::warn!("failed connection from {:?}: {:?}", peer, err),
+		}
 	}
 }
 
@@ -83,12 +85,12 @@ async fn accept_connection(
 	peer: SocketAddr,
 	stream: TcpStream,
 	incoming: Sender<NetworkMessage<ClientMessage>>,
-) -> Connection {
+) -> Result<Connection, tungstenite::Error> {
 	// Generate random connection id
 	let conn_id = rand::random();
 
 	// Get streams
-	let ws_stream = accept_async(stream).await.expect("could not accept");
+	let ws_stream = accept_async(stream).await?;
 	let (ws_sender, ws_receiver) = ws_stream.split();
 
 	// Send event to game instance
@@ -102,11 +104,11 @@ async fn accept_connection(
 	task::spawn(read_loop(ws_receiver, incoming, conn_id));
 
 	// Return connection instance for network manager
-	Connection {
+	Ok(Connection {
 		stream: ws_sender,
 		address: peer,
 		conn_id,
-	}
+	})
 }
 
 async fn read_loop(
@@ -159,7 +161,14 @@ async fn read_loop(
 
 async fn write_loop(net: Arc<Mutex<NetworkManager>>, ch: Receiver<NetworkMessage<ServerMessage>>) {
 	loop {
-		let msg = ch.recv().await.expect("failed to receive outgoing message");
-		net.lock().await.send(msg).await;
+		match ch.recv().await {
+			Ok(msg) => {
+				net.lock().await.send(msg).await;
+			}
+			Err(err) => {
+				log::warn!("write loop received error: {:?}", err);
+				break;
+			}
+		}
 	}
 }