From bda8e7b94cc60ac943de776e3ee08ba20f3aa250 Mon Sep 17 00:00:00 2001
From: griffi-gh <prasol258@gmail.com>
Date: Sun, 12 Feb 2023 22:38:51 +0100
Subject: [PATCH] disconnect on exit

---
 kubi-udp/src/client.rs |  6 ++++++
 kubi/src/main.rs       |  6 ++++--
 kubi/src/networking.rs | 19 +++++++++++++++++--
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/kubi-udp/src/client.rs b/kubi-udp/src/client.rs
index be9d6e1..50d02de 100644
--- a/kubi-udp/src/client.rs
+++ b/kubi-udp/src/client.rs
@@ -128,6 +128,12 @@ impl<S, R> Client<S, R> where S: Encode + Decode, R: Encode + Decode {
     Ok(())
   }
 
+  #[inline]
+  pub fn set_nonblocking(&mut self, is_nonblocking: bool) -> Result<()> {
+    self.socket.set_nonblocking(is_nonblocking)?;
+    Ok(())
+  }
+
   #[inline]
   pub fn get_status(&self) -> ClientStatus {
     self.status
diff --git a/kubi/src/main.rs b/kubi/src/main.rs
index 4cdfce5..3032daf 100644
--- a/kubi/src/main.rs
+++ b/kubi/src/main.rs
@@ -75,7 +75,7 @@ use delta_time::{DeltaTime, init_delta_time};
 use cursor_lock::{insert_lock_state, update_cursor_lock_state, lock_cursor_now};
 use control_flow::{exit_on_esc, insert_control_flow_unique, SetControlFlow};
 use state::{is_ingame, is_ingame_or_loading, is_loading, init_state, update_state};
-use networking::{update_networking, is_multiplayer};
+use networking::{update_networking, is_multiplayer, disconnect_on_exit};
 use init::initialize_from_args;
 use gui::{render_gui, init_gui, update_gui};
 use loading_screen::update_loading_screen;
@@ -104,7 +104,7 @@ fn update() -> Workload {
     update_window_size,
     update_cursor_lock_state,
     process_inputs,
-    exit_on_esc,
+    
     (
       update_networking
     ).into_workload().run_if(is_multiplayer),
@@ -124,6 +124,8 @@ fn update() -> Workload {
     compute_cameras,
     update_gui,
     update_state,
+    exit_on_esc,
+    disconnect_on_exit,
   ).into_workload()
 }
 fn render() -> Workload {
diff --git a/kubi/src/networking.rs b/kubi/src/networking.rs
index 16450b0..7ef0678 100644
--- a/kubi/src/networking.rs
+++ b/kubi/src/networking.rs
@@ -1,4 +1,5 @@
-use shipyard::{Unique, AllStoragesView, UniqueView, UniqueViewMut, Workload, IntoWorkload, WorkloadModificator, EntitiesView, EntitiesViewMut, Component, ViewMut, SystemModificator};
+use shipyard::{Unique, AllStoragesView, UniqueView, UniqueViewMut, Workload, IntoWorkload, EntitiesViewMut, Component, ViewMut, SystemModificator};
+use glium::glutin::event_loop::ControlFlow;
 use std::net::SocketAddr;
 use kubi_udp::client::{Client, ClientConfig, ClientEvent};
 use kubi_shared::networking::{
@@ -6,7 +7,7 @@ use kubi_shared::networking::{
   state::ClientJoinState
 };
 
-use crate::events::EventComponent;
+use crate::{events::EventComponent, control_flow::SetControlFlow};
 
 #[derive(Unique, Clone, Copy, PartialEq, Eq)]
 pub enum GameType {
@@ -73,6 +74,20 @@ pub fn update_networking() -> Workload {
   ).into_workload()
 }
 
+pub fn disconnect_on_exit(
+  control_flow: UniqueView<SetControlFlow>,
+  mut client: UniqueViewMut<UdpClient>,
+) {
+  if let Some(ControlFlow::ExitWithCode(_)) = control_flow.0 {
+    client.0.set_nonblocking(false).expect("Failed to switch socket to blocking mode");
+    if let Err(error) = client.0.disconnect() {
+      log::error!("failed to disconnect: {}", error);
+    } else {
+      log::info!("Client disconnected");
+    }
+  }
+}
+
 pub fn is_multiplayer(
   game_type: UniqueView<GameType>
 ) -> bool {