From e9b3128103d203ad4902a40124e4d22a012c196a Mon Sep 17 00:00:00 2001
From: mat <git@matdoes.dev>
Date: Fri, 9 May 2025 14:12:51 -1030
Subject: don't send chat signing certs on offline-mode servers

---
 azalea-client/src/plugins/chat_signing.rs |  4 ++--
 azalea-client/src/plugins/disconnect.rs   | 34 +++++++++++++++++++++----------
 azalea-client/src/plugins/login.rs        | 10 ++++++++-
 3 files changed, 34 insertions(+), 14 deletions(-)

(limited to 'azalea-client/src/plugins')

diff --git a/azalea-client/src/plugins/chat_signing.rs b/azalea-client/src/plugins/chat_signing.rs
index 03dd1a21..d65f4eb3 100644
--- a/azalea-client/src/plugins/chat_signing.rs
+++ b/azalea-client/src/plugins/chat_signing.rs
@@ -12,7 +12,7 @@ use chrono::Utc;
 use tracing::{debug, error};
 use uuid::Uuid;
 
-use super::{chat, packet::game::SendPacketEvent};
+use super::{chat, login::IsAuthenticated, packet::game::SendPacketEvent};
 use crate::{Account, InGameState};
 
 pub struct ChatSigningPlugin;
@@ -149,7 +149,7 @@ pub struct QueuedCertsToSend {
 
 pub fn handle_queued_certs_to_send(
     mut commands: Commands,
-    query: Query<(Entity, &QueuedCertsToSend)>,
+    query: Query<(Entity, &QueuedCertsToSend), With<IsAuthenticated>>,
 ) {
     for (entity, queued_certs) in &query {
         let certs = &queued_certs.certs;
diff --git a/azalea-client/src/plugins/disconnect.rs b/azalea-client/src/plugins/disconnect.rs
index 3e27eba3..c1c962d5 100644
--- a/azalea-client/src/plugins/disconnect.rs
+++ b/azalea-client/src/plugins/disconnect.rs
@@ -7,6 +7,7 @@ use bevy_ecs::prelude::*;
 use derive_more::Deref;
 use tracing::info;
 
+use super::login::IsAuthenticated;
 use crate::{InstanceHolder, chat_signing, client::JoinedClientBundle, connection::RawConnection};
 
 pub struct DisconnectPlugin;
@@ -44,6 +45,27 @@ pub struct DisconnectEvent {
     pub reason: Option<FormattedText>,
 }
 
+/// A bundle of components that are removed when a client disconnects.
+///
+/// This shouldn't be used for inserts because not all of the components should
+/// always be present.
+#[derive(Bundle)]
+pub struct RemoveOnDisconnectBundle {
+    pub joined_client: JoinedClientBundle,
+    pub entity: EntityBundle,
+    pub instance_holder: InstanceHolder,
+    pub player_metadata: PlayerMetadataBundle,
+    pub in_loaded_chunk: InLoadedChunk,
+    //// This makes it close the TCP connection.
+    pub raw_connection: RawConnection,
+    /// This makes it not send [`DisconnectEvent`] again.
+    pub is_connection_alive: IsConnectionAlive,
+    /// Resend our chat signing certs next time.
+    pub chat_signing_session: chat_signing::ChatSigningSession,
+    /// They're not authenticated anymore if they disconnected.
+    pub is_authenticated: IsAuthenticated,
+}
+
 /// A system that removes the several components from our clients when they get
 /// a [`DisconnectEvent`].
 pub fn remove_components_from_disconnected_players(
@@ -62,17 +84,7 @@ pub fn remove_components_from_disconnected_players(
         );
         commands
             .entity(*entity)
-            .remove::<JoinedClientBundle>()
-            .remove::<EntityBundle>()
-            .remove::<InstanceHolder>()
-            .remove::<PlayerMetadataBundle>()
-            .remove::<InLoadedChunk>()
-            // this makes it close the tcp connection
-            .remove::<RawConnection>()
-            // this makes it not send DisconnectEvent again
-            .remove::<IsConnectionAlive>()
-            // resend our chat signing certs next time
-            .remove::<chat_signing::ChatSigningSession>();
+            .remove::<RemoveOnDisconnectBundle>();
         // note that we don't remove the client from the ECS, so if they decide
         // to reconnect they'll keep their state
 
diff --git a/azalea-client/src/plugins/login.rs b/azalea-client/src/plugins/login.rs
index 357769e9..ebba5905 100644
--- a/azalea-client/src/plugins/login.rs
+++ b/azalea-client/src/plugins/login.rs
@@ -33,6 +33,11 @@ fn handle_receive_hello_event(trigger: Trigger<ReceiveHelloEvent>, mut commands:
     commands.entity(player).insert(AuthTask(task));
 }
 
+/// A marker component on our clients that indicates that the server is
+/// online-mode and the client has authenticated their join with Mojang.
+#[derive(Component)]
+pub struct IsAuthenticated;
+
 pub fn poll_auth_task(
     mut commands: Commands,
     mut query: Query<(Entity, &mut AuthTask, &mut RawConnection)>,
@@ -40,7 +45,10 @@ pub fn poll_auth_task(
     for (entity, mut auth_task, mut raw_conn) in query.iter_mut() {
         if let Some(poll_res) = future::block_on(future::poll_once(&mut auth_task.0)) {
             debug!("Finished auth");
-            commands.entity(entity).remove::<AuthTask>();
+            commands
+                .entity(entity)
+                .remove::<AuthTask>()
+                .insert(IsAuthenticated);
             match poll_res {
                 Ok((packet, private_key)) => {
                     // we use this instead of SendLoginPacketEvent to ensure that it's sent right
-- 
cgit v1.2.3