fix errors when switching to Game state and add fast_login test

6 files changed, 101 insertions, 20 deletions

diff --git a/azalea-client/src/lib.rs b/azalea-client/src/lib.rs
index bcea00bc..e8f9db01 100644
--- a/azalea-client/src/lib.rs
+++ b/azalea-client/src/lib.rs
@@ -23,8 +23,8 @@ pub mod test_simulation;
 pub use account::{Account, AccountOpts};
 pub use azalea_protocol::common::client_information::ClientInformation;
 pub use client::{
-    Client, DefaultPlugins, InConfigState, JoinError, JoinedClientBundle, LocalPlayerBundle,
-    StartClientOpts, TickBroadcast, start_ecs_runner,
+    Client, DefaultPlugins, InConfigState, InGameState, JoinError, JoinedClientBundle,
+    LocalPlayerBundle, StartClientOpts, TickBroadcast, start_ecs_runner,
 };
 pub use events::Event;
 pub use local_player::{GameProfileComponent, Hunger, InstanceHolder, TabList};
diff --git a/azalea-client/src/plugins/packet/config/events.rs b/azalea-client/src/plugins/packet/config/events.rs
index a9b2179a..eddcf72f 100644
--- a/azalea-client/src/plugins/packet/config/events.rs
+++ b/azalea-client/src/plugins/packet/config/events.rs
@@ -55,7 +55,7 @@ pub fn handle_send_packet_event(
     }
 }
 
-pub fn send_packet_events(
+pub fn emit_receive_config_packet_events(
     query: Query<(Entity, &RawConnection), With<InConfigState>>,
     mut packet_events: ResMut<Events<ReceiveConfigPacketEvent>>,
 ) {
@@ -67,7 +67,9 @@ pub fn send_packet_events(
         let packets_lock = raw_conn.incoming_packet_queue();
         let mut packets = packets_lock.lock();
         if !packets.is_empty() {
+            let mut packets_read = 0;
             for raw_packet in packets.iter() {
+                packets_read += 1;
                 let packet = match deserialize_packet::<ClientboundConfigPacket>(&mut Cursor::new(
                     raw_packet,
                 )) {
@@ -78,13 +80,32 @@ pub fn send_packet_events(
                         continue;
                     }
                 };
+
+                let should_interrupt = packet_interrupts(&packet);
+
                 packet_events.send(ReceiveConfigPacketEvent {
                     entity: player_entity,
                     packet,
                 });
+
+                if should_interrupt {
+                    break;
+                }
             }
-            // clear the packets right after we read them
-            packets.clear();
+            packets.drain(0..packets_read);
         }
     }
 }
+
+/// Whether the given packet should make us stop deserializing the received
+/// packets until next update.
+///
+/// This is used for packets that can switch the client state.
+fn packet_interrupts(packet: &ClientboundConfigPacket) -> bool {
+    matches!(
+        packet,
+        ClientboundConfigPacket::FinishConfiguration(_)
+            | ClientboundConfigPacket::Disconnect(_)
+            | ClientboundConfigPacket::Transfer(_)
+    )
+}
diff --git a/azalea-client/src/plugins/packet/game/events.rs b/azalea-client/src/plugins/packet/game/events.rs
index 9e7f88fa..e5c87971 100644
--- a/azalea-client/src/plugins/packet/game/events.rs
+++ b/azalea-client/src/plugins/packet/game/events.rs
@@ -75,7 +75,7 @@ pub fn handle_outgoing_packets(
     }
 }
 
-pub fn send_receivepacketevent(
+pub fn emit_receive_packet_events(
     query: Query<(Entity, &RawConnection), With<InGameState>>,
     mut packet_events: ResMut<Events<ReceivePacketEvent>>,
 ) {
@@ -87,7 +87,9 @@ pub fn send_receivepacketevent(
         let packets_lock = raw_connection.incoming_packet_queue();
         let mut packets = packets_lock.lock();
         if !packets.is_empty() {
+            let mut packets_read = 0;
             for raw_packet in packets.iter() {
+                packets_read += 1;
                 let packet =
                     match deserialize_packet::<ClientboundGamePacket>(&mut Cursor::new(raw_packet))
                     {
@@ -98,17 +100,36 @@ pub fn send_receivepacketevent(
                             continue;
                         }
                     };
+
+                let should_interrupt = packet_interrupts(&packet);
+
                 packet_events.send(ReceivePacketEvent {
                     entity: player_entity,
                     packet: Arc::new(packet),
                 });
+
+                if should_interrupt {
+                    break;
+                }
             }
-            // clear the packets right after we read them
-            packets.clear();
+            packets.drain(0..packets_read);
         }
     }
 }
 
+/// Whether the given packet should make us stop deserializing the received
+/// packets until next update.
+///
+/// This is used for packets that can switch the client state.
+fn packet_interrupts(packet: &ClientboundGamePacket) -> bool {
+    matches!(
+        packet,
+        ClientboundGamePacket::StartConfiguration(_)
+            | ClientboundGamePacket::Disconnect(_)
+            | ClientboundGamePacket::Transfer(_)
+    )
+}
+
 /// A player joined the game (or more specifically, was added to the tab
 /// list of a local player).
 #[derive(Event, Debug, Clone)]
diff --git a/azalea-client/src/plugins/packet/mod.rs b/azalea-client/src/plugins/packet/mod.rs
index cbd8a175..96a396d9 100644
--- a/azalea-client/src/plugins/packet/mod.rs
+++ b/azalea-client/src/plugins/packet/mod.rs
@@ -38,7 +38,10 @@ impl Plugin for PacketPlugin {
     fn build(&self, app: &mut App) {
         app.add_systems(
             First,
-            (game::send_receivepacketevent, config::send_packet_events),
+            (
+                game::emit_receive_packet_events,
+                config::emit_receive_config_packet_events,
+            ),
         )
         .add_systems(
             PreUpdate,
diff --git a/azalea-client/tests/change_dimension_to_nether_and_back.rs b/azalea-client/tests/change_dimension_to_nether_and_back.rs
index 16febca0..748ea713 100644
--- a/azalea-client/tests/change_dimension_to_nether_and_back.rs
+++ b/azalea-client/tests/change_dimension_to_nether_and_back.rs
@@ -1,10 +1,9 @@
-use azalea_client::{InConfigState, test_simulation::*};
+use azalea_client::{InConfigState, InGameState, test_simulation::*};
 use azalea_core::{position::ChunkPos, resource_location::ResourceLocation};
-use azalea_entity::{LocalEntity, metadata::Health};
+use azalea_entity::LocalEntity;
 use azalea_protocol::packets::{
     ConnectionProtocol,
     config::{ClientboundFinishConfiguration, ClientboundRegistryData},
-    game::ClientboundSetHealth,
 };
 use azalea_registry::DimensionType;
 use azalea_world::InstanceName;
@@ -17,6 +16,7 @@ fn test_change_dimension_to_nether_and_back() {
 
     let mut simulation = Simulation::new(ConnectionProtocol::Configuration);
     assert!(simulation.has_component::<InConfigState>());
+    assert!(!simulation.has_component::<InGameState>());
 
     simulation.receive_packet(ClientboundRegistryData {
         registry_id: ResourceLocation::new("minecraft:dimension_type"),
@@ -53,16 +53,9 @@ fn test_change_dimension_to_nether_and_back() {
     simulation.tick();
 
     assert!(!simulation.has_component::<InConfigState>());
+    assert!(simulation.has_component::<InGameState>());
     assert!(simulation.has_component::<LocalEntity>());
 
-    simulation.receive_packet(ClientboundSetHealth {
-        health: 15.,
-        food: 20,
-        saturation: 20.,
-    });
-    simulation.tick();
-    assert_eq!(*simulation.component::<Health>(), 15.);
-
     //
     // OVERWORLD
     //
diff --git a/azalea-client/tests/fast_login.rs b/azalea-client/tests/fast_login.rs
new file mode 100644
index 00000000..bc26079a
--- /dev/null
+++ b/azalea-client/tests/fast_login.rs
@@ -0,0 +1,43 @@
+use azalea_client::{InConfigState, test_simulation::*};
+use azalea_core::resource_location::ResourceLocation;
+use azalea_entity::metadata::Health;
+use azalea_protocol::packets::{
+    ConnectionProtocol,
+    config::{ClientboundFinishConfiguration, ClientboundRegistryData},
+    game::ClientboundSetHealth,
+};
+use bevy_log::tracing_subscriber;
+use simdnbt::owned::{NbtCompound, NbtTag};
+
+#[test]
+fn test_fast_login() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let mut simulation = Simulation::new(ConnectionProtocol::Configuration);
+    assert!(simulation.has_component::<InConfigState>());
+
+    simulation.receive_packet(ClientboundRegistryData {
+        registry_id: ResourceLocation::new("minecraft:dimension_type"),
+        entries: vec![(
+            ResourceLocation::new("minecraft:overworld"),
+            Some(NbtCompound::from_values(vec![
+                ("height".into(), NbtTag::Int(384)),
+                ("min_y".into(), NbtTag::Int(-64)),
+            ])),
+        )]
+        .into_iter()
+        .collect(),
+    });
+
+    simulation.receive_packet(ClientboundFinishConfiguration);
+    // note that there's no simulation tick here
+    simulation.receive_packet(ClientboundSetHealth {
+        health: 15.,
+        food: 20,
+        saturation: 20.,
+    });
+    simulation.tick();
+    // we need a second tick to handle the state switch properly
+    simulation.tick();
+    assert_eq!(*simulation.component::<Health>(), 15.);
+}