From 1b4c1c9677a0663ebc0a154afa9601178f90a56c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ole=20Str=C3=B8hm?= <Strohm99@gmail.com>
Date: Tue, 7 Nov 2023 02:02:46 +0000
Subject: [PATCH] Renet: Allow application code to be transport agnostic (#119)

---
 bevy_renet/examples/simple.rs  |   4 +-
 bevy_renet/src/lib.rs          |  59 +++++++++++++++++++
 bevy_renet/src/steam.rs        |  51 ++--------------
 bevy_renet/src/transport.rs    |  52 ++---------------
 demo_bevy/src/bin/client.rs    |   6 +-
 demo_chat/src/client.rs        |   2 +-
 renet/examples/echo.rs         |   2 +-
 renet/src/error.rs             |   2 +-
 renet/src/lib.rs               |   2 +-
 renet/src/remote_connection.rs | 103 ++++++++++++++++++++++++---------
 renet/src/server.rs            |  10 +---
 renet/src/transport/client.rs  |  22 +++----
 renet_steam/examples/echo.rs   |   2 +-
 renet_steam/src/client.rs      |  18 +++---
 14 files changed, 178 insertions(+), 157 deletions(-)

diff --git a/bevy_renet/examples/simple.rs b/bevy_renet/examples/simple.rs
index 39bca086..0384b768 100644
--- a/bevy_renet/examples/simple.rs
+++ b/bevy_renet/examples/simple.rs
@@ -5,7 +5,7 @@ use bevy_renet::{
         ConnectionConfig, DefaultChannel, RenetClient, RenetServer, ServerEvent,
     },
     transport::{NetcodeClientPlugin, NetcodeServerPlugin},
-    RenetClientPlugin, RenetServerPlugin,
+    RenetClientPlugin, RenetServerPlugin, client_connected,
 };
 use renet::{
     transport::{NetcodeClientTransport, NetcodeServerTransport, NetcodeTransportError},
@@ -118,7 +118,7 @@ fn main() {
 
         app.add_systems(
             Update,
-            (player_input, client_send_input, client_sync_players).run_if(bevy_renet::transport::client_connected()),
+            (player_input, client_send_input, client_sync_players).run_if(client_connected()),
         );
     }
 
diff --git a/bevy_renet/src/lib.rs b/bevy_renet/src/lib.rs
index 3fe2831b..dda6b137 100644
--- a/bevy_renet/src/lib.rs
+++ b/bevy_renet/src/lib.rs
@@ -10,6 +10,24 @@ pub mod transport;
 #[cfg(feature = "steam")]
 pub mod steam;
 
+/// This system set is where all transports receive messages
+///
+/// If you want to ensure data has arrived in the [`RenetClient`] or [`RenetServer`], then schedule your
+/// system after this set.
+///
+/// This system set runs in PreUpdate.
+#[derive(Debug, SystemSet, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct RenetReceive;
+
+/// This system set is where all transports send messages
+///
+/// If you want to ensure your packets have been registered by the [`RenetClient`] or [`RenetServer`], then
+/// schedule your system before this set.
+///
+/// This system set runs in PostUpdate.
+#[derive(Debug, SystemSet, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct RenetSend;
+
 pub struct RenetServerPlugin;
 
 pub struct RenetClientPlugin;
@@ -42,3 +60,44 @@ impl RenetClientPlugin {
         client.update(time.delta());
     }
 }
+
+pub fn client_connected() -> impl FnMut(Option<Res<RenetClient>>) -> bool {
+    |client| match client {
+        Some(client) => client.is_connected(),
+        None => false,
+    }
+}
+
+pub fn client_disconnected() -> impl FnMut(Option<Res<RenetClient>>) -> bool {
+    |client| match client {
+        Some(client) => client.is_disconnected(),
+        None => true,
+    }
+}
+
+pub fn client_connecting() -> impl FnMut(Option<Res<RenetClient>>) -> bool {
+    |client| match client {
+        Some(client) => client.is_connecting(),
+        None => false,
+    }
+}
+
+pub fn client_just_connected() -> impl FnMut(Local<bool>, Option<Res<RenetClient>>) -> bool {
+    |mut last_connected: Local<bool>, client| {
+        let connected = client.map(|client| client.is_connected()).unwrap_or(false);
+
+        let just_connected = !*last_connected && connected;
+        *last_connected = connected;
+        just_connected
+    }
+}
+
+pub fn client_just_disconnected() -> impl FnMut(Local<bool>, Option<Res<RenetClient>>) -> bool {
+    |mut last_connected: Local<bool>, client| {
+        let disconnected = client.map(|client| client.is_disconnected()).unwrap_or(true);
+
+        let just_disconnected = *last_connected && disconnected;
+        *last_connected = !disconnected;
+        just_disconnected
+    }
+}
diff --git a/bevy_renet/src/steam.rs b/bevy_renet/src/steam.rs
index d02d8c74..b39c0c97 100644
--- a/bevy_renet/src/steam.rs
+++ b/bevy_renet/src/steam.rs
@@ -2,7 +2,7 @@ use bevy::{app::AppExit, prelude::*};
 use renet::{RenetClient, RenetServer};
 use renet_steam::steamworks::SteamError;
 
-use crate::{RenetClientPlugin, RenetServerPlugin};
+use crate::{RenetClientPlugin, RenetReceive, RenetSend, RenetServerPlugin};
 
 pub use renet_steam::{AccessPermission, SteamClientTransport, SteamServerConfig, SteamServerTransport};
 
@@ -24,25 +24,26 @@ impl Plugin for SteamServerPlugin {
         app.add_systems(
             PreUpdate,
             Self::update_system
+                .in_set(RenetReceive)
                 .run_if(resource_exists::<RenetServer>())
                 .after(RenetServerPlugin::update_system),
         );
 
         app.add_systems(
             PostUpdate,
-            (Self::send_packets, Self::disconnect_on_exit).run_if(resource_exists::<RenetServer>()),
+            (Self::send_packets.in_set(RenetSend), Self::disconnect_on_exit).run_if(resource_exists::<RenetServer>()),
         );
     }
 }
 
 impl SteamServerPlugin {
-    pub fn update_system(mut transport: Option<NonSendMut<SteamServerTransport>>, mut server: ResMut<RenetServer>) {
+    fn update_system(mut transport: Option<NonSendMut<SteamServerTransport>>, mut server: ResMut<RenetServer>) {
         if let Some(transport) = transport.as_mut() {
             transport.update(&mut server);
         }
     }
 
-    pub fn send_packets(mut transport: Option<NonSendMut<SteamServerTransport>>, mut server: ResMut<RenetServer>) {
+    fn send_packets(mut transport: Option<NonSendMut<SteamServerTransport>>, mut server: ResMut<RenetServer>) {
         if let Some(transport) = transport.as_mut() {
             transport.send_packets(&mut server);
         }
@@ -97,48 +98,8 @@ impl SteamClientPlugin {
     }
 
     fn disconnect_on_exit(exit: EventReader<AppExit>, mut transport: ResMut<SteamClientTransport>) {
-        if !exit.is_empty() && !transport.is_disconnected() {
+        if !exit.is_empty() {
             transport.disconnect();
         }
     }
 }
-
-pub fn client_connected() -> impl FnMut(Option<Res<SteamClientTransport>>) -> bool {
-    |transport| match transport {
-        Some(transport) => transport.is_connected(),
-        None => false,
-    }
-}
-
-pub fn client_disconnected() -> impl FnMut(Option<Res<SteamClientTransport>>) -> bool {
-    |transport| match transport {
-        Some(transport) => transport.is_disconnected(),
-        None => true,
-    }
-}
-
-pub fn client_connecting() -> impl FnMut(Option<Res<SteamClientTransport>>) -> bool {
-    |transport| match transport {
-        Some(transport) => transport.is_connecting(),
-        None => false,
-    }
-}
-
-pub fn client_just_connected() -> impl FnMut(Local<bool>, Option<Res<SteamClientTransport>>) -> bool {
-    |mut last_connected: Local<bool>, transport| {
-        let connected = transport.map(|transport| transport.is_connected()).unwrap_or(false);
-        let just_connected = !*last_connected && connected;
-        *last_connected = connected;
-        just_connected
-    }
-}
-
-pub fn client_just_disconnected() -> impl FnMut(Local<bool>, Option<Res<SteamClientTransport>>) -> bool {
-    |mut last_connected: Local<bool>, transport| {
-        let disconnected = transport.map(|transport| transport.is_disconnected()).unwrap_or(true);
-
-        let just_disconnected = *last_connected && disconnected;
-        *last_connected = !disconnected;
-        just_disconnected
-    }
-}
diff --git a/bevy_renet/src/transport.rs b/bevy_renet/src/transport.rs
index 2b77dfa8..bb6b79a9 100644
--- a/bevy_renet/src/transport.rs
+++ b/bevy_renet/src/transport.rs
@@ -5,7 +5,7 @@ use renet::{
 
 use bevy::{app::AppExit, prelude::*};
 
-use crate::{RenetClientPlugin, RenetServerPlugin};
+use crate::{RenetClientPlugin, RenetReceive, RenetSend, RenetServerPlugin};
 
 pub struct NetcodeServerPlugin;
 
@@ -18,6 +18,7 @@ impl Plugin for NetcodeServerPlugin {
         app.add_systems(
             PreUpdate,
             Self::update_system
+                .in_set(RenetReceive)
                 .run_if(resource_exists::<NetcodeServerTransport>())
                 .run_if(resource_exists::<RenetServer>())
                 .after(RenetServerPlugin::update_system),
@@ -25,7 +26,7 @@ impl Plugin for NetcodeServerPlugin {
 
         app.add_systems(
             PostUpdate,
-            (Self::send_packets, Self::disconnect_on_exit)
+            (Self::send_packets.in_set(RenetSend), Self::disconnect_on_exit)
                 .run_if(resource_exists::<NetcodeServerTransport>())
                 .run_if(resource_exists::<RenetServer>()),
         );
@@ -33,7 +34,7 @@ impl Plugin for NetcodeServerPlugin {
 }
 
 impl NetcodeServerPlugin {
-    pub fn update_system(
+    fn update_system(
         mut transport: ResMut<NetcodeServerTransport>,
         mut server: ResMut<RenetServer>,
         time: Res<Time>,
@@ -44,7 +45,7 @@ impl NetcodeServerPlugin {
         }
     }
 
-    pub fn send_packets(mut transport: ResMut<NetcodeServerTransport>, mut server: ResMut<RenetServer>) {
+    fn send_packets(mut transport: ResMut<NetcodeServerTransport>, mut server: ResMut<RenetServer>) {
         transport.send_packets(&mut server);
     }
 
@@ -98,49 +99,8 @@ impl NetcodeClientPlugin {
     }
 
     fn disconnect_on_exit(exit: EventReader<AppExit>, mut transport: ResMut<NetcodeClientTransport>) {
-        if !exit.is_empty() && !transport.is_disconnected() {
+        if !exit.is_empty() {
             transport.disconnect();
         }
     }
 }
-
-pub fn client_connected() -> impl FnMut(Option<Res<NetcodeClientTransport>>) -> bool {
-    |transport| match transport {
-        Some(transport) => transport.is_connected(),
-        None => false,
-    }
-}
-
-pub fn client_disconnected() -> impl FnMut(Option<Res<NetcodeClientTransport>>) -> bool {
-    |transport| match transport {
-        Some(transport) => transport.is_disconnected(),
-        None => true,
-    }
-}
-
-pub fn client_connecting() -> impl FnMut(Option<Res<NetcodeClientTransport>>) -> bool {
-    |transport| match transport {
-        Some(transport) => transport.is_connecting(),
-        None => false,
-    }
-}
-
-pub fn client_just_connected() -> impl FnMut(Local<bool>, Option<Res<NetcodeClientTransport>>) -> bool {
-    |mut last_connected: Local<bool>, transport| {
-        let connected = transport.map(|transport| transport.is_connected()).unwrap_or(false);
-
-        let just_connected = !*last_connected && connected;
-        *last_connected = connected;
-        just_connected
-    }
-}
-
-pub fn client_just_disconnected() -> impl FnMut(Local<bool>, Option<Res<NetcodeClientTransport>>) -> bool {
-    |mut last_connected: Local<bool>, transport| {
-        let disconnected = transport.map(|transport| transport.is_disconnected()).unwrap_or(true);
-
-        let just_disconnected = *last_connected && disconnected;
-        *last_connected = !disconnected;
-        just_disconnected
-    }
-}
diff --git a/demo_bevy/src/bin/client.rs b/demo_bevy/src/bin/client.rs
index 8d9b7463..360b0680 100644
--- a/demo_bevy/src/bin/client.rs
+++ b/demo_bevy/src/bin/client.rs
@@ -7,6 +7,7 @@ use bevy::{
 };
 use bevy_egui::{EguiContexts, EguiPlugin};
 use bevy_renet::{
+    client_connected,
     renet::{ClientId, RenetClient},
     RenetClientPlugin,
 };
@@ -46,7 +47,8 @@ fn add_netcode_network(app: &mut App) {
     use std::{net::UdpSocket, time::SystemTime};
 
     app.add_plugins(bevy_renet::transport::NetcodeClientPlugin);
-    app.configure_sets(Update, Connected.run_if(bevy_renet::transport::client_connected()));
+
+    app.configure_sets(Update, Connected.run_if(client_connected()));
 
     let client = RenetClient::new(connection_config());
 
@@ -98,7 +100,7 @@ fn add_steam_network(app: &mut App) {
     app.insert_resource(transport);
     app.insert_resource(CurrentClientId(steam_client.user().steam_id().raw()));
 
-    app.configure_sets(Update, Connected.run_if(bevy_renet::steam::client_connected()));
+    app.configure_sets(Update, Connected.run_if(client_connected()));
 
     app.insert_non_send_resource(single);
     fn steam_callbacks(client: NonSend<SingleClient>) {
diff --git a/demo_chat/src/client.rs b/demo_chat/src/client.rs
index 020c2f19..8643cd85 100644
--- a/demo_chat/src/client.rs
+++ b/demo_chat/src/client.rs
@@ -57,7 +57,7 @@ impl ChatApp {
             AppState::MainScreen => {
                 draw_main_screen(&mut self.ui_state, &mut self.state, ctx);
             }
-            AppState::ClientChat { transport, .. } if transport.is_connecting() => draw_loader(ctx),
+            AppState::ClientChat { client, .. } if client.is_connecting() => draw_loader(ctx),
             AppState::ClientChat { usernames, .. } => {
                 let usernames = usernames.clone();
                 draw_chat(&mut self.ui_state, &mut self.state, usernames, ctx);
diff --git a/renet/examples/echo.rs b/renet/examples/echo.rs
index 366a391c..ca3d58ce 100644
--- a/renet/examples/echo.rs
+++ b/renet/examples/echo.rs
@@ -153,7 +153,7 @@ fn client(server_addr: SocketAddr, username: Username) {
         client.update(duration);
         transport.update(duration, &mut client).unwrap();
 
-        if transport.is_connected() {
+        if client.is_connected() {
             match stdin_channel.try_recv() {
                 Ok(text) => client.send_message(DefaultChannel::ReliableOrdered, text.as_bytes().to_vec()),
                 Err(TryRecvError::Empty) => {}
diff --git a/renet/src/error.rs b/renet/src/error.rs
index a45e4b95..f419fefe 100644
--- a/renet/src/error.rs
+++ b/renet/src/error.rs
@@ -2,7 +2,7 @@ use std::fmt;
 
 use crate::packet::SerializationError;
 
-/// Possibles reasons for a disconnection.
+/// Possible reasons for a disconnection.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum DisconnectReason {
     /// Connection was terminated by the transport layer
diff --git a/renet/src/lib.rs b/renet/src/lib.rs
index 3a150c2c..74600c79 100644
--- a/renet/src/lib.rs
+++ b/renet/src/lib.rs
@@ -10,7 +10,7 @@ pub mod transport;
 
 pub use channel::{ChannelConfig, DefaultChannel, SendType};
 pub use error::{ChannelError, ClientNotFound, DisconnectReason};
-pub use remote_connection::{ConnectionConfig, NetworkInfo, RenetClient};
+pub use remote_connection::{ConnectionConfig, RenetConnectionStatus, NetworkInfo, RenetClient};
 pub use server::{RenetServer, ServerEvent};
 
 pub use bytes::Bytes;
diff --git a/renet/src/remote_connection.rs b/renet/src/remote_connection.rs
index ef066602..082436d7 100644
--- a/renet/src/remote_connection.rs
+++ b/renet/src/remote_connection.rs
@@ -70,6 +70,14 @@ pub struct NetworkInfo {
     pub bytes_received_per_second: f64,
 }
 
+/// The connection status of a [`RenetClient`].
+#[derive(Debug)]
+pub enum RenetConnectionStatus {
+    Connected,
+    Connecting,
+    Disconnected { reason: DisconnectReason },
+}
+
 #[derive(Debug)]
 #[cfg_attr(feature = "bevy", derive(bevy_ecs::system::Resource))]
 pub struct RenetClient {
@@ -84,7 +92,7 @@ pub struct RenetClient {
     receive_reliable_channels: HashMap<u8, ReceiveChannelReliable>,
     stats: ConnectionStats,
     available_bytes_per_tick: u64,
-    pub(crate) disconnect_reason: Option<DisconnectReason>,
+    connection_status: RenetConnectionStatus,
     rtt: f64,
 }
 
@@ -180,7 +188,7 @@ impl RenetClient {
             stats: ConnectionStats::new(),
             rtt: 0.0,
             available_bytes_per_tick,
-            disconnect_reason: None,
+            connection_status: RenetConnectionStatus::Connecting,
         }
     }
 
@@ -214,41 +222,74 @@ impl RenetClient {
         }
     }
 
-    /// Returns if the client is disconnected.
-    ///
-    /// Note: to check if a client is connecting you need to use the transport layer [NetcodeClientTransport::is_connecting][crate::transport::NetcodeClientTransport::is_connecting].
-    /// You can't use `!client.is_disconnected()` for that.
+    /// Returns whether the client is connected.
+    #[inline]
+    pub fn is_connected(&self) -> bool {
+        matches!(self.connection_status, RenetConnectionStatus::Connected)
+    }
+
+    /// Returns whether the client is connecting.
+    #[inline]
+    pub fn is_connecting(&self) -> bool {
+        matches!(self.connection_status, RenetConnectionStatus::Connecting)
+    }
+
+    /// Returns whether the client is disconnected.
     #[inline]
     pub fn is_disconnected(&self) -> bool {
-        self.disconnect_reason.is_some()
+        matches!(self.connection_status, RenetConnectionStatus::Disconnected { .. })
     }
 
     /// Returns the disconneect reason if the client is disconnected.
     pub fn disconnect_reason(&self) -> Option<DisconnectReason> {
-        self.disconnect_reason
+        if let RenetConnectionStatus::Disconnected { reason } = self.connection_status {
+            Some(reason)
+        } else {
+            None
+        }
+    }
+
+    /// Set the client connection status to connected.
+    ///
+    /// Does nothing if the client is disconnected. A disconnected client must be reconstructed.
+    ///
+    /// <p style="background:rgba(77,220,255,0.16);padding:0.5em;">
+    /// <strong>Note:</strong> This should only be called by the transport layer.
+    /// </p>
+    pub fn set_connected(&mut self) {
+        if !self.is_disconnected() {
+            self.connection_status = RenetConnectionStatus::Connected;
+        }
+    }
+
+    /// Set the client connection status to connecting.
+    ///
+    /// Does nothing if the client is disconnected. A disconnected client must be reconstructed.
+    ///
+    /// <p style="background:rgba(77,220,255,0.16);padding:0.5em;">
+    /// <strong>Note:</strong> This should only be called by the transport layer.
+    /// </p>
+    pub fn set_connecting(&mut self) {
+        if !self.is_disconnected() {
+            self.connection_status = RenetConnectionStatus::Connecting;
+        }
     }
 
     /// Disconnect the client.
+    ///
     /// If the client is already disconnected, it does nothing.
     pub fn disconnect(&mut self) {
-        if self.disconnect_reason.is_some() {
-            return;
-        }
-
-        self.disconnect_reason = Some(DisconnectReason::DisconnectedByClient);
+        self.disconnect_with_reason(DisconnectReason::DisconnectedByClient);
     }
 
     /// Disconnect the client because an error occurred in the transport layer.
+    ///
     /// If the client is already disconnected, it does nothing.
     /// <p style="background:rgba(77,220,255,0.16);padding:0.5em;">
     /// <strong>Note:</strong> This should only be called by the transport layer.
     /// </p>
     pub fn disconnect_due_to_transport(&mut self) {
-        if self.disconnect_reason.is_some() {
-            return;
-        }
-
-        self.disconnect_reason = Some(DisconnectReason::Transport);
+        self.disconnect_with_reason(DisconnectReason::Transport);
     }
 
     /// Returns the available memory in bytes for the given channel.
@@ -284,7 +325,7 @@ impl RenetClient {
         let channel_id = channel_id.into();
         if let Some(reliable_channel) = self.send_reliable_channels.get_mut(&channel_id) {
             if let Err(error) = reliable_channel.send_message(message.into()) {
-                self.disconnect_reason = Some(DisconnectReason::SendChannelError { channel_id, error });
+                self.disconnect_with_reason(DisconnectReason::SendChannelError { channel_id, error });
             }
         } else if let Some(unreliable_channel) = self.send_unreliable_channels.get_mut(&channel_id) {
             unreliable_channel.send_message(message.into());
@@ -350,7 +391,7 @@ impl RenetClient {
         let mut octets = octets::Octets::with_slice(packet);
         let packet = match Packet::from_bytes(&mut octets) {
             Err(err) => {
-                self.disconnect_reason = Some(DisconnectReason::PacketDeserialization(err));
+                self.disconnect_with_reason(DisconnectReason::PacketDeserialization(err));
                 return;
             }
             Ok(packet) => packet,
@@ -361,20 +402,20 @@ impl RenetClient {
         match packet {
             Packet::SmallReliable { channel_id, messages, .. } => {
                 let Some(channel) = self.receive_reliable_channels.get_mut(&channel_id) else {
-                    self.disconnect_reason = Some(DisconnectReason::ReceivedInvalidChannelId(channel_id));
+                    self.disconnect_with_reason(DisconnectReason::ReceivedInvalidChannelId(channel_id));
                     return;
                 };
 
                 for (message_id, message) in messages {
                     if let Err(error) = channel.process_message(message, message_id) {
-                        self.disconnect_reason = Some(DisconnectReason::ReceiveChannelError { channel_id, error });
+                        self.disconnect_with_reason(DisconnectReason::ReceiveChannelError { channel_id, error });
                         return;
                     }
                 }
             }
             Packet::SmallUnreliable { channel_id, messages, .. } => {
                 let Some(channel) = self.receive_unreliable_channels.get_mut(&channel_id) else {
-                    self.disconnect_reason = Some(DisconnectReason::ReceivedInvalidChannelId(channel_id));
+                    self.disconnect_with_reason(DisconnectReason::ReceivedInvalidChannelId(channel_id));
                     return;
                 };
 
@@ -384,22 +425,22 @@ impl RenetClient {
             }
             Packet::ReliableSlice { channel_id, slice, .. } => {
                 let Some(channel) = self.receive_reliable_channels.get_mut(&channel_id) else {
-                    self.disconnect_reason = Some(DisconnectReason::ReceivedInvalidChannelId(channel_id));
+                    self.disconnect_with_reason(DisconnectReason::ReceivedInvalidChannelId(channel_id));
                     return;
                 };
 
                 if let Err(error) = channel.process_slice(slice) {
-                    self.disconnect_reason = Some(DisconnectReason::ReceiveChannelError { channel_id, error });
+                    self.disconnect_with_reason(DisconnectReason::ReceiveChannelError { channel_id, error });
                 }
             }
             Packet::UnreliableSlice { channel_id, slice, .. } => {
                 let Some(channel) = self.receive_unreliable_channels.get_mut(&channel_id) else {
-                    self.disconnect_reason = Some(DisconnectReason::ReceivedInvalidChannelId(channel_id));
+                    self.disconnect_with_reason(DisconnectReason::ReceivedInvalidChannelId(channel_id));
                     return;
                 };
 
                 if let Err(error) = channel.process_slice(slice, self.current_time) {
-                    self.disconnect_reason = Some(DisconnectReason::ReceiveChannelError { channel_id, error });
+                    self.disconnect_with_reason(DisconnectReason::ReceiveChannelError { channel_id, error });
                 }
             }
             Packet::Ack { ack_ranges, .. } => {
@@ -557,7 +598,7 @@ impl RenetClient {
             let mut oct = OctetsMut::with_slice(&mut buffer);
             let len = match packet.to_bytes(&mut oct) {
                 Err(err) => {
-                    self.disconnect_reason = Some(DisconnectReason::PacketSerialization(err));
+                    self.disconnect_with_reason(DisconnectReason::PacketSerialization(err));
                     return vec![];
                 }
                 Ok(len) => len,
@@ -645,6 +686,12 @@ impl RenetClient {
             return;
         }
     }
+
+    pub(crate) fn disconnect_with_reason(&mut self, reason: DisconnectReason) {
+        if !self.is_disconnected() {
+            self.connection_status = RenetConnectionStatus::Disconnected { reason };
+        }
+    }
 }
 
 #[cfg(test)]
diff --git a/renet/src/server.rs b/renet/src/server.rs
index 1baada3c..f463a8b6 100644
--- a/renet/src/server.rs
+++ b/renet/src/server.rs
@@ -133,21 +133,17 @@ impl RenetServer {
         }
     }
 
-    /// Disconnects a client, it does nothing if the client does not exits.
+    /// Disconnects a client, it does nothing if the client does not exist.
     pub fn disconnect(&mut self, client_id: ClientId) {
         if let Some(connection) = self.connections.get_mut(&client_id) {
-            if !connection.is_disconnected() {
-                connection.disconnect_reason = Some(DisconnectReason::DisconnectedByServer);
-            }
+            connection.disconnect_with_reason(DisconnectReason::DisconnectedByServer)
         }
     }
 
     /// Disconnects all client.
     pub fn disconnect_all(&mut self) {
         for connection in self.connections.values_mut() {
-            if !connection.is_disconnected() {
-                connection.disconnect_reason = Some(DisconnectReason::DisconnectedByServer);
-            }
+            connection.disconnect_with_reason(DisconnectReason::DisconnectedByServer)
         }
     }
 
diff --git a/renet/src/transport/client.rs b/renet/src/transport/client.rs
index 1987bda3..0595cc80 100644
--- a/renet/src/transport/client.rs
+++ b/renet/src/transport/client.rs
@@ -38,18 +38,6 @@ impl NetcodeClientTransport {
         self.netcode_client.client_id()
     }
 
-    pub fn is_connecting(&self) -> bool {
-        self.netcode_client.is_connecting()
-    }
-
-    pub fn is_connected(&self) -> bool {
-        self.netcode_client.is_connected()
-    }
-
-    pub fn is_disconnected(&self) -> bool {
-        self.netcode_client.is_disconnected()
-    }
-
     /// Returns the duration since the client last received a packet.
     /// Usefull to detect timeouts.
     pub fn time_since_last_received_packet(&self) -> Duration {
@@ -99,9 +87,7 @@ impl NetcodeClientTransport {
     pub fn update(&mut self, duration: Duration, client: &mut RenetClient) -> Result<(), NetcodeTransportError> {
         if let Some(reason) = self.netcode_client.disconnect_reason() {
             // Mark the client as disconnected if an error occured in the transport layer
-            if !client.is_disconnected() {
-                client.disconnect_due_to_transport();
-            }
+            client.disconnect_due_to_transport();
 
             return Err(NetcodeError::Disconnected(reason).into());
         }
@@ -112,6 +98,12 @@ impl NetcodeClientTransport {
             return Err(error.into());
         }
 
+        if self.netcode_client.is_connected() {
+            client.set_connected();
+        } else if self.netcode_client.is_connecting() {
+            client.set_connecting();
+        }
+
         loop {
             let packet = match self.socket.recv_from(&mut self.buffer) {
                 Ok((len, addr)) => {
diff --git a/renet_steam/examples/echo.rs b/renet_steam/examples/echo.rs
index ba951d21..be7f81ac 100644
--- a/renet_steam/examples/echo.rs
+++ b/renet_steam/examples/echo.rs
@@ -157,7 +157,7 @@ fn run_client(steam_client: Client<ClientManager>, single: SingleClient, server_
         client.update(duration);
         transport.update(&mut client);
 
-        if transport.is_connected() {
+        if client.is_connected() {
             match stdin_channel.try_recv() {
                 Ok(text) => {
                     client.send_message(DefaultChannel::ReliableOrdered, text.as_bytes().to_vec());
diff --git a/renet_steam/src/client.rs b/renet_steam/src/client.rs
index 9f13a811..2e9f58ec 100644
--- a/renet_steam/src/client.rs
+++ b/renet_steam/src/client.rs
@@ -31,25 +31,25 @@ impl SteamClientTransport {
         })
     }
 
-    pub fn is_connected(&self) -> bool {
+    fn is_connected(&self) -> bool {
         let status = self.connection_state();
 
         status == NetworkingConnectionState::Connected
     }
 
-    pub fn is_disconnected(&self) -> bool {
+    fn is_disconnected(&self) -> bool {
         let status = self.connection_state();
         status == NetworkingConnectionState::ClosedByPeer
             || status == NetworkingConnectionState::ProblemDetectedLocally
             || status == NetworkingConnectionState::None
     }
 
-    pub fn is_connecting(&self) -> bool {
+    fn is_connecting(&self) -> bool {
         let status = self.connection_state();
         status == NetworkingConnectionState::Connecting || status == NetworkingConnectionState::FindingRoute
     }
 
-    pub fn connection_state(&self) -> NetworkingConnectionState {
+    fn connection_state(&self) -> NetworkingConnectionState {
         let connection = match &self.state {
             ConnectionState::Connected { connection } => connection,
             ConnectionState::Disconnected { .. } => {
@@ -103,9 +103,7 @@ impl SteamClientTransport {
     pub fn update(&mut self, client: &mut RenetClient) {
         if self.is_disconnected() {
             // Mark the client as disconnected if an error occured in the transport layer
-            if !client.is_disconnected() {
-                client.disconnect_due_to_transport();
-            }
+            client.disconnect_due_to_transport();
 
             if let ConnectionState::Connected { connection } = &self.state {
                 let end_reason = self
@@ -121,6 +119,12 @@ impl SteamClientTransport {
             return;
         };
 
+        if self.is_connected() {
+            client.set_connected();
+        } else if self.is_connecting() {
+            client.set_connecting();
+        }
+
         let ConnectionState::Connected { connection } = &mut self.state else {
             unreachable!()
         };