From 8837d1c6c337b244c2df5e14f12ab1c91c330053 Mon Sep 17 00:00:00 2001
From: Sofia <sofia@teascade.net>
Date: Tue, 4 Aug 2026 20:48:45 +0300
Subject: [PATCH] Add close reasons

---
 src/connections.rs | 53 +++++++++++++++++++++++++++++-----------------
 src/lib.rs         |  3 ++-
 src/package.rs     | 12 ++++++++++-
 3 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/src/connections.rs b/src/connections.rs
index a2062cf..af8417a 100644
--- a/src/connections.rs
+++ b/src/connections.rs
@@ -12,7 +12,7 @@ use thiserror::Error;
 use crate::{
     PeerConfig, PeerMessage,
     listener::{DATAGRAM_SIZE, ListenerError},
-    package::{Message, Messages, Package},
+    package::{CloseReason, Message, Messages, Package},
     stats::NetStats,
 };
 
@@ -71,7 +71,7 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
     pub fn close(&mut self) {
         self.closing_since = Some(Instant::now());
         for (addr, _) in self.connections.clone() {
-            self.close_connection(&addr);
+            self.close_connection(&addr, CloseReason::ShuttingDown);
         }
     }
 
@@ -131,8 +131,9 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
                                 }
                             }
                         }
-                        ConnectionState::Closing | ConnectionState::ReceivingClosing => {
-                            match self.udp.send_to(*addr, Package::<T>::Close) {
+                        ConnectionState::Closing(reason)
+                        | ConnectionState::ReceivingClosing(reason) => {
+                            match self.udp.send_to(*addr, Package::<T>::Close(reason)) {
                                 Ok(bytes) => {
                                     conn.bytes_tx += bytes;
                                     conn.last_sent_ping = now;
@@ -223,11 +224,11 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
     }
 
     /// Close the connection to another peer at the given address
-    pub fn close_connection(&mut self, addr: &SocketAddr) {
+    pub fn close_connection(&mut self, addr: &SocketAddr, reason: CloseReason) {
         if let Some(conn) = self.connections.get_mut(addr) {
-            if conn.state != ConnectionState::Closing {
+            if !matches!(conn.state, ConnectionState::Closing(_)) {
                 conn.closing_since = Instant::now();
-                conn.state = ConnectionState::Closing;
+                conn.state = ConnectionState::Closing(reason);
             }
         }
     }
@@ -265,8 +266,13 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
                     }
                 } else {
                     if !self.accepting_connections || self.config.identifier != identifier {
-                        self.connections
-                            .insert(*addr, Connection::from(*addr, ConnectionState::Closing));
+                        self.connections.insert(
+                            *addr,
+                            Connection::from(
+                                *addr,
+                                ConnectionState::Closing(CloseReason::NotAcceptingConnections),
+                            ),
+                        );
                     } else {
                         self.connections.insert(
                             *addr,
@@ -306,22 +312,25 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
                     }
                 }
             }
-            Package::Close => {
+            Package::Close(reason) => {
                 let remove = if let Some(conn) = self.connections.get_mut(addr) {
                     match conn.state {
-                        ConnectionState::Closing => {
-                            messages
-                                .push(PeerMessage::Disconnected(conn.clone(), conn.error.take()));
+                        ConnectionState::Closing(_) => {
+                            messages.push(PeerMessage::Disconnected(
+                                conn.clone(),
+                                conn.error.take(),
+                                CloseReason::SelfClosed,
+                            ));
                             true
                         }
                         ConnectionState::Error => false,
-                        ConnectionState::ReceivingClosing => {
+                        ConnectionState::ReceivingClosing(_) => {
                             conn.last_recv_close = Instant::now();
                             false
                         }
                         _ => {
                             conn.last_recv_close = Instant::now();
-                            conn.state = ConnectionState::ReceivingClosing;
+                            conn.state = ConnectionState::ReceivingClosing(reason);
                             false
                         }
                     }
@@ -401,16 +410,20 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
 
         self.connections.retain(|_, conn| {
             let retain = match conn.state {
-                ConnectionState::ReceivingClosing => {
+                ConnectionState::ReceivingClosing(reason) => {
                     if (now - conn.last_recv_close) > self.config.disconnect_timeout {
-                        messages.push(PeerMessage::Disconnected(conn.clone(), None));
+                        messages.push(PeerMessage::Disconnected(conn.clone(), None, reason));
                         false
                     } else {
                         true
                     }
                 }
                 ConnectionState::Error => {
-                    messages.push(PeerMessage::Disconnected(conn.clone(), conn.error.take()));
+                    messages.push(PeerMessage::Disconnected(
+                        conn.clone(),
+                        conn.error.take(),
+                        CloseReason::Error,
+                    ));
                     false
                 }
                 _ => true,
@@ -584,9 +597,9 @@ pub enum ConnectionState {
     Connected,
 
     /// Connection is being closed
-    Closing,
+    Closing(CloseReason),
     /// Connection is being closed by the other peer
-    ReceivingClosing,
+    ReceivingClosing(CloseReason),
 
     /// Connection has ended up in an erronous state
     Error,
diff --git a/src/lib.rs b/src/lib.rs
index 6877026..117803a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -16,6 +16,7 @@ use thiserror::*;
 use crate::{
     connections::{Connection, ConnectionError, ConnectionManager},
     listener::{Listener, ListenerMessage},
+    package::CloseReason,
     stats::NetStats,
 };
 
@@ -38,7 +39,7 @@ pub enum PeerMessage<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'st
     /// A new connection has been connected
     NewConnection(Connection<T>),
     /// An existing connection has disconnected, with an optional error
-    Disconnected(Connection<T>, Option<ConnectionError>),
+    Disconnected(Connection<T>, Option<ConnectionError>, CloseReason),
     /// The Peer has been closed
     Closed,
     /// A single message of type T
diff --git a/src/package.rs b/src/package.rs
index 9227ebf..37f9fe6 100644
--- a/src/package.rs
+++ b/src/package.rs
@@ -5,10 +5,20 @@ pub enum Package<T> {
     Hello(String),
     Ping,
     Pong,
-    Close,
+    Close(CloseReason),
     Messages(Messages<T>),
 }
 
+#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
+pub enum CloseReason {
+    NotAcceptingConnections,
+    Kicked,
+    Error,
+    ShuttingDown,
+    /// Closed because local peer initiated close
+    SelfClosed,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Messages<T> {
     pub ack: u64,