From 89bc5ca91ebfe585c8df8891d94856068c7ad2be Mon Sep 17 00:00:00 2001
From: mat <git@matdoes.dev>
Date: Fri, 25 Apr 2025 09:56:57 +1000
Subject: [PATCH] update RawConnection::state when start_configuration is
 received

---
 azalea-client/src/plugins/connection.rs       |  1 -
 .../src/plugins/packet/config/mod.rs          |  2 +-
 azalea-client/src/plugins/packet/game/mod.rs  | 37 ++++++++++++-------
 3 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/azalea-client/src/plugins/connection.rs b/azalea-client/src/plugins/connection.rs
index 30ac5d20..36aa2ee7 100644
--- a/azalea-client/src/plugins/connection.rs
+++ b/azalea-client/src/plugins/connection.rs
@@ -34,7 +34,6 @@ impl Plugin for ConnectionPlugin {
 }
 
 pub fn read_packets(ecs: &mut World) {
-    // receive_game_packet_events: EventWriter<ReceiveGamePacketEvent>,
     let mut entity_and_conn_query = ecs.query::<(Entity, &mut RawConnection)>();
     let mut conn_query = ecs.query::<&mut RawConnection>();
 
diff --git a/azalea-client/src/plugins/packet/config/mod.rs b/azalea-client/src/plugins/packet/config/mod.rs
index 910019a6..9560703b 100644
--- a/azalea-client/src/plugins/packet/config/mod.rs
+++ b/azalea-client/src/plugins/packet/config/mod.rs
@@ -96,12 +96,12 @@ impl ConfigPacketHandler<'_> {
             self.ecs,
             |(mut commands, mut query)| {
                 let mut raw_conn = query.get_mut(self.player).unwrap();
+                raw_conn.state = ConnectionProtocol::Game;
 
                 commands.trigger(SendConfigPacketEvent::new(
                     self.player,
                     ServerboundFinishConfiguration,
                 ));
-                raw_conn.state = ConnectionProtocol::Game;
 
                 // these components are added now that we're going to be in the Game state
                 commands
diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index beec9219..6235eafd 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -13,7 +13,7 @@ use azalea_entity::{
     indexing::{EntityIdIndex, EntityUuidIndex},
     metadata::{Health, apply_metadata},
 };
-use azalea_protocol::packets::game::*;
+use azalea_protocol::packets::{ConnectionProtocol, game::*};
 use azalea_world::{InstanceContainer, InstanceName, MinecraftEntityId, PartialInstance};
 use bevy_ecs::{prelude::*, system::SystemState};
 pub use events::*;
@@ -22,7 +22,9 @@ use tracing::{debug, error, trace, warn};
 use crate::{
     ClientInformation, PlayerInfo,
     chat::{ChatPacket, ChatReceivedEvent},
-    chunks, declare_packet_handlers,
+    chunks,
+    connection::RawConnection,
+    declare_packet_handlers,
     disconnect::DisconnectEvent,
     inventory::{
         ClientSideCloseContainerEvent, Inventory, MenuOpenedEvent, SetContainerContentEvent,
@@ -1485,18 +1487,27 @@ impl GamePacketHandler<'_> {
     pub fn start_configuration(&mut self, _p: &ClientboundStartConfiguration) {
         debug!("Got start configuration packet");
 
-        as_system::<Commands>(self.ecs, |mut commands| {
-            commands.trigger(SendPacketEvent::new(
-                self.player,
-                ServerboundConfigurationAcknowledged,
-            ));
+        as_system::<(Commands, Query<&mut RawConnection>)>(
+            self.ecs,
+            |(mut commands, mut query)| {
+                let Some(mut raw_conn) = query.get_mut(self.player).ok() else {
+                    warn!("Got start configuration packet but player doesn't have a RawConnection");
+                    return;
+                };
+                raw_conn.state = ConnectionProtocol::Configuration;
 
-            commands
-                .entity(self.player)
-                .insert(crate::client::InConfigState)
-                .remove::<crate::JoinedClientBundle>()
-                .remove::<EntityBundle>();
-        });
+                commands.trigger(SendPacketEvent::new(
+                    self.player,
+                    ServerboundConfigurationAcknowledged,
+                ));
+
+                commands
+                    .entity(self.player)
+                    .insert(crate::client::InConfigState)
+                    .remove::<crate::JoinedClientBundle>()
+                    .remove::<EntityBundle>();
+            },
+        );
     }
 
     pub fn entity_position_sync(&mut self, p: &ClientboundEntityPositionSync) {