From 19bc67a78f8af7f9a5ed3dffde1f63ee20cd1380 Mon Sep 17 00:00:00 2001
From: griffi-gh <prasol258@gmail.com>
Date: Fri, 3 Feb 2023 19:29:38 +0100
Subject: [PATCH] disconnect reason

---
 kubi-udp/Cargo.toml    |  2 ++
 kubi-udp/src/client.rs | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/kubi-udp/Cargo.toml b/kubi-udp/Cargo.toml
index 8a64827..f99b6ad 100644
--- a/kubi-udp/Cargo.toml
+++ b/kubi-udp/Cargo.toml
@@ -8,3 +8,5 @@ edition = "2021"
 [dependencies]
 bincode = "2.0.0-rc"
 anyhow = "1.0"
+nohash-hasher = "0.2.0"
+log = "0.4"
diff --git a/kubi-udp/src/client.rs b/kubi-udp/src/client.rs
index 5ac8c5d..0538c43 100644
--- a/kubi-udp/src/client.rs
+++ b/kubi-udp/src/client.rs
@@ -18,6 +18,19 @@ pub struct ClientConfig {
   pub heartbeat_interval: Duration,
 }
 
+pub type ClientId = u8;
+
+#[derive(Default, Encode, Decode)]
+#[repr(u8)]
+pub enum DisconnectReason {
+  #[default]
+  NotConnected,
+  ClientDisconnected,
+  KickedByServer,
+  ClientTimeout,
+  ServerTimeout,
+}
+
 pub struct Client<S, R> where S: Encode + Decode, R: Encode + Decode {
   addr: SocketAddr,
   config: ClientConfig,
@@ -25,6 +38,8 @@ pub struct Client<S, R> where S: Encode + Decode, R: Encode + Decode {
   status: ClientStatus,
   timeout: Instant,
   last_heartbeat: Instant,
+  client_id: Option<ClientId>,
+  disconnect_reason: DisconnectReason,
   _s: PhantomData<*const S>,
   _r: PhantomData<*const R>,
 }
@@ -40,6 +55,8 @@ impl<S, R> Client<S, R> where S: Encode + Decode, R: Encode + Decode {
       status: ClientStatus::Disconnected,
       timeout: Instant::now(),
       last_heartbeat: Instant::now(),
+      client_id: None,
+      disconnect_reason: DisconnectReason::default(),
       _s: PhantomData,
       _r: PhantomData,
     })
@@ -64,12 +81,17 @@ impl<S, R> Client<S, R> where S: Encode + Decode, R: Encode + Decode {
     self.send_raw_packet(&ClientPacket::Data(message))?;
     Ok(())
   }
+  fn disconnect_inner(&mut self, reason: DisconnectReason) -> anyhow::Result<()> {
+    self.send_raw_packet(&ClientPacket::Disconnect)?;
+    self.status = ClientStatus::Disconnected;
+    self.disconnect_reason = DisconnectReason::ClientDisconnected;
+    Ok(())
+  }
   pub fn disconnect(&mut self) -> anyhow::Result<()> {
     if self.status != ClientStatus::Connected {
       anyhow::bail!("Not Connected");
     }
-    self.send_raw_packet(&ClientPacket::Disconnect)?;
-    self.status = ClientStatus::Disconnected;
+    self.disconnect_inner(DisconnectReason::ClientDisconnected)?;
     Ok(())
   }
   pub fn update(&mut self) -> anyhow::Result<()> {
@@ -77,12 +99,15 @@ impl<S, R> Client<S, R> where S: Encode + Decode, R: Encode + Decode {
       return Ok(())
     }
     if self.timeout.elapsed() > self.config.timeout {
-      //We don't care if this packet actually gets sent becauseserver is likely dead
-      let _ = self.send_raw_packet(&ClientPacket::Disconnect);
-      self.status = ClientStatus::Disconnected;
+      log::warn!("Client timed out");
+      //We don't care if this packet actually gets sent because the server is likely dead
+      let _ = self.disconnect_inner(DisconnectReason::ClientDisconnected).map_err(|_| {
+        log::warn!("Failed to send disconnect packet");
+      });
       return Ok(())
     }
     if self.last_heartbeat.elapsed() > self.config.heartbeat_interval {
+      log::trace!("Sending heartbeat packet");
       self.send_raw_packet(&ClientPacket::Heartbeat)?;
       self.last_heartbeat = Instant::now();
     }