Fix errors on switching dimensions (#204)

* Fix errors on switching dimensions

* fix other tests

* clippy

* fix log feature in test_simulation

* fix chunks oops

5 files changed, 382 insertions, 161 deletions

diff --git a/azalea-client/src/lib.rs b/azalea-client/src/lib.rs
index a61c5d32..d07e323d 100644
--- a/azalea-client/src/lib.rs
+++ b/azalea-client/src/lib.rs
@@ -29,6 +29,9 @@ pub mod respawn;
 pub mod send_client_end;
 pub mod task_pool;
 
+#[doc(hidden)]
+pub mod test_simulation;
+
 pub use account::{Account, AccountOpts};
 pub use azalea_protocol::common::client_information::ClientInformation;
 pub use client::{
diff --git a/azalea-client/src/packet_handling/game.rs b/azalea-client/src/packet_handling/game.rs
index 207cdd91..73442804 100644
--- a/azalea-client/src/packet_handling/game.rs
+++ b/azalea-client/src/packet_handling/game.rs
@@ -246,17 +246,10 @@ pub fn process_packet_events(ecs: &mut World) {
                             .insert(InstanceName(new_instance_name.clone()));
                     }
 
-                    let Some(dimension_type_element) =
-                        instance_holder.instance.read().registries.dimension_type()
+                    let Some((_dimension_type, dimension_data)) = p
+                        .common
+                        .dimension_type(&instance_holder.instance.read().registries)
                     else {
-                        error!("Server didn't send dimension type registry, can't log in");
-                        continue;
-                    };
-
-                    let dimension_name = ResourceLocation::new(&p.common.dimension.to_string());
-
-                    let Some(dimension) = dimension_type_element.map.get(&dimension_name) else {
-                        error!("No dimension_type with name {dimension_name}");
                         continue;
                     };
 
@@ -264,8 +257,9 @@ pub fn process_packet_events(ecs: &mut World) {
                     // there)
                     let weak_instance = instance_container.insert(
                         new_instance_name.clone(),
-                        dimension.height,
-                        dimension.min_y,
+                        dimension_data.height,
+                        dimension_data.min_y,
+                        &instance_holder.instance.read().registries,
                     );
                     instance_loaded_events.send(InstanceLoadedEvent {
                         entity: player_entity,
@@ -1387,17 +1381,10 @@ pub fn process_packet_events(ecs: &mut World) {
                 {
                     let new_instance_name = p.common.dimension.clone();
 
-                    let Some(dimension_type_element) =
-                        instance_holder.instance.read().registries.dimension_type()
+                    let Some((_dimension_type, dimension_data)) = p
+                        .common
+                        .dimension_type(&instance_holder.instance.read().registries)
                     else {
-                        error!("Server didn't send dimension type registry, can't log in.");
-                        continue;
-                    };
-
-                    let dimension_name = ResourceLocation::new(&p.common.dimension.to_string());
-
-                    let Some(dimension) = dimension_type_element.map.get(&dimension_name) else {
-                        error!("No dimension_type with name {dimension_name}");
                         continue;
                     };
 
@@ -1405,8 +1392,9 @@ pub fn process_packet_events(ecs: &mut World) {
                     // there)
                     let weak_instance = instance_container.insert(
                         new_instance_name.clone(),
-                        dimension.height,
-                        dimension.min_y,
+                        dimension_data.height,
+                        dimension_data.min_y,
+                        &instance_holder.instance.read().registries,
                     );
                     instance_loaded_events.send(InstanceLoadedEvent {
                         entity: player_entity,
diff --git a/azalea-client/tests/simulation.rs b/azalea-client/src/test_simulation.rs
index 593622aa..5afd8b00 100644
--- a/azalea-client/tests/simulation.rs
+++ b/azalea-client/src/test_simulation.rs
@@ -1,102 +1,106 @@
 use std::{fmt::Debug, sync::Arc, time::Duration};
 
 use azalea_auth::game_profile::GameProfile;
-use azalea_client::{
-    events::LocalPlayerEvents,
-    raw_connection::{RawConnection, RawConnectionReader, RawConnectionWriter},
-    ClientInformation, GameProfileComponent, InConfigState, InstanceHolder, LocalPlayerBundle,
-};
-use azalea_core::{
-    game_type::{GameMode, OptionalGameType},
-    resource_location::ResourceLocation,
-    tick::GameTick,
-};
-use azalea_entity::{
-    metadata::{Health, PlayerMetadataBundle},
-    LocalEntity,
-};
-use azalea_protocol::packets::{
-    common::CommonPlayerSpawnInfo,
-    config::{ClientboundFinishConfiguration, ClientboundRegistryData},
-    game::{ClientboundLogin, ClientboundSetHealth},
-    ConnectionProtocol, Packet, ProtocolPacket,
+use azalea_buf::AzaleaWrite;
+use azalea_core::game_type::{GameMode, OptionalGameType};
+use azalea_core::position::ChunkPos;
+use azalea_core::resource_location::ResourceLocation;
+use azalea_core::tick::GameTick;
+use azalea_entity::metadata::PlayerMetadataBundle;
+use azalea_protocol::packets::common::CommonPlayerSpawnInfo;
+use azalea_protocol::packets::game::c_level_chunk_with_light::ClientboundLevelChunkPacketData;
+use azalea_protocol::packets::game::c_light_update::ClientboundLightUpdatePacketData;
+use azalea_protocol::packets::game::{
+    ClientboundLevelChunkWithLight, ClientboundLogin, ClientboundRespawn,
 };
+use azalea_protocol::packets::{ConnectionProtocol, Packet, ProtocolPacket};
 use azalea_registry::DimensionType;
-use azalea_world::{Instance, MinecraftEntityId};
+use azalea_world::palette::{PalettedContainer, PalettedContainerKind};
+use azalea_world::{Chunk, Instance, MinecraftEntityId, Section};
 use bevy_app::App;
-use bevy_app::PluginGroup;
 use bevy_ecs::{prelude::*, schedule::ExecutorKind};
-use bevy_log::{tracing_subscriber, LogPlugin};
 use parking_lot::{Mutex, RwLock};
-use simdnbt::owned::{NbtCompound, NbtTag};
+use simdnbt::owned::Nbt;
 use tokio::{sync::mpsc, time::sleep};
 use uuid::Uuid;
 
-#[test]
-fn test_set_health_before_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.tick();
-    simulation.receive_packet(ClientboundFinishConfiguration);
-    simulation.tick();
+use crate::{
+    events::LocalPlayerEvents,
+    raw_connection::{RawConnection, RawConnectionReader, RawConnectionWriter},
+    ClientInformation, GameProfileComponent, InConfigState, InstanceHolder, LocalPlayerBundle,
+};
 
-    assert!(!simulation.has_component::<InConfigState>());
-    assert!(simulation.has_component::<LocalEntity>());
+/// A way to simulate a client in a server, used for some internal tests.
+pub struct Simulation {
+    pub app: App,
+    pub entity: Entity,
 
-    simulation.receive_packet(ClientboundSetHealth {
-        health: 15.,
-        food: 20,
-        saturation: 20.,
-    });
-    simulation.tick();
-    assert_eq!(*simulation.component::<Health>(), 15.);
+    // the runtime needs to be kept around for the tasks to be considered alive
+    pub rt: tokio::runtime::Runtime,
 
-    simulation.receive_packet(ClientboundLogin {
-        player_id: MinecraftEntityId(0),
-        hardcore: false,
-        levels: vec![],
-        max_players: 20,
-        chunk_radius: 8,
-        simulation_distance: 8,
-        reduced_debug_info: false,
-        show_death_screen: true,
-        do_limited_crafting: false,
-        common: CommonPlayerSpawnInfo {
-            dimension_type: DimensionType::Overworld,
-            dimension: ResourceLocation::new("minecraft:overworld"),
-            seed: 0,
-            game_type: GameMode::Survival,
-            previous_game_type: OptionalGameType(None),
-            is_debug: false,
-            is_flat: false,
-            last_death_location: None,
-            portal_cooldown: 0,
-            sea_level: 63,
-        },
-        enforces_secure_chat: false,
-    });
-    simulation.tick();
+    pub incoming_packet_queue: Arc<Mutex<Vec<Box<[u8]>>>>,
+    pub outgoing_packets_receiver: mpsc::UnboundedReceiver<Box<[u8]>>,
+}
+
+impl Simulation {
+    pub fn new(initial_connection_protocol: ConnectionProtocol) -> Self {
+        let mut app = create_simulation_app();
+        let mut entity = app.world_mut().spawn_empty();
+        let (player, outgoing_packets_receiver, incoming_packet_queue, rt) =
+            create_local_player_bundle(entity.id(), initial_connection_protocol);
+        entity.insert(player);
+
+        let entity = entity.id();
 
-    // health should stay the same
-    assert_eq!(*simulation.component::<Health>(), 15.);
+        tick_app(&mut app);
+
+        #[allow(clippy::single_match)]
+        match initial_connection_protocol {
+            ConnectionProtocol::Configuration => {
+                app.world_mut().entity_mut(entity).insert(InConfigState);
+                tick_app(&mut app);
+            }
+            _ => {}
+        }
+
+        Self {
+            app,
+            entity,
+            rt,
+            incoming_packet_queue,
+            outgoing_packets_receiver,
+        }
+    }
+
+    pub fn receive_packet<P: ProtocolPacket + Debug>(&mut self, packet: impl Packet<P>) {
+        let buf = azalea_protocol::write::serialize_packet(&packet.into_variant()).unwrap();
+        self.incoming_packet_queue.lock().push(buf);
+    }
+
+    pub fn tick(&mut self) {
+        tick_app(&mut self.app);
+    }
+    pub fn component<T: Component + Clone>(&self) -> T {
+        self.app.world().get::<T>(self.entity).unwrap().clone()
+    }
+    pub fn get_component<T: Component + Clone>(&self) -> Option<T> {
+        self.app.world().get::<T>(self.entity).cloned()
+    }
+    pub fn has_component<T: Component>(&self) -> bool {
+        self.app.world().get::<T>(self.entity).is_some()
+    }
+
+    pub fn chunk(&self, chunk_pos: ChunkPos) -> Option<Arc<RwLock<Chunk>>> {
+        self.component::<InstanceHolder>()
+            .instance
+            .read()
+            .chunks
+            .get(&chunk_pos)
+    }
 }
 
-pub fn create_local_player_bundle(
+#[allow(clippy::type_complexity)]
+fn create_local_player_bundle(
     entity: Entity,
     connection_protocol: ConnectionProtocol,
 ) -> (
@@ -164,7 +168,12 @@ pub fn create_local_player_bundle(
 
 fn create_simulation_app() -> App {
     let mut app = App::new();
-    app.add_plugins(azalea_client::DefaultPlugins.build().disable::<LogPlugin>());
+
+    #[cfg(feature = "log")]
+    app.add_plugins(
+        bevy_app::PluginGroup::build(crate::DefaultPlugins).disable::<bevy_log::LogPlugin>(),
+    );
+
     app.edit_schedule(bevy_app::Main, |schedule| {
         // makes test results more reproducible
         schedule.set_executor_kind(ExecutorKind::SingleThreaded);
@@ -172,66 +181,85 @@ fn create_simulation_app() -> App {
     app
 }
 
-pub struct Simulation {
-    pub app: App,
-    pub entity: Entity,
-
-    // the runtime needs to be kept around for the tasks to be considered alive
-    pub rt: tokio::runtime::Runtime,
-
-    pub incoming_packet_queue: Arc<Mutex<Vec<Box<[u8]>>>>,
-    pub outgoing_packets_receiver: mpsc::UnboundedReceiver<Box<[u8]>>,
+fn tick_app(app: &mut App) {
+    app.update();
+    app.world_mut().run_schedule(GameTick);
 }
 
-impl Simulation {
-    pub fn new(initial_connection_protocol: ConnectionProtocol) -> Self {
-        let mut app = create_simulation_app();
-        let mut entity = app.world_mut().spawn_empty();
-        let (player, outgoing_packets_receiver, incoming_packet_queue, rt) =
-            create_local_player_bundle(entity.id(), initial_connection_protocol);
-        entity.insert(player);
-
-        let entity = entity.id();
-
-        tick_app(&mut app);
-
-        match initial_connection_protocol {
-            ConnectionProtocol::Configuration => {
-                app.world_mut().entity_mut(entity).insert(InConfigState);
-                tick_app(&mut app);
-            }
-            _ => {}
-        }
-
-        Self {
-            app,
-            entity,
-            rt,
-            incoming_packet_queue,
-            outgoing_packets_receiver,
-        }
+pub fn make_basic_login_packet(
+    dimension_type: DimensionType,
+    dimension: ResourceLocation,
+) -> ClientboundLogin {
+    ClientboundLogin {
+        player_id: MinecraftEntityId(0),
+        hardcore: false,
+        levels: vec![],
+        max_players: 20,
+        chunk_radius: 8,
+        simulation_distance: 8,
+        reduced_debug_info: false,
+        show_death_screen: true,
+        do_limited_crafting: false,
+        common: CommonPlayerSpawnInfo {
+            dimension_type,
+            dimension,
+            seed: 0,
+            game_type: GameMode::Survival,
+            previous_game_type: OptionalGameType(None),
+            is_debug: false,
+            is_flat: false,
+            last_death_location: None,
+            portal_cooldown: 0,
+            sea_level: 63,
+        },
+        enforces_secure_chat: false,
     }
+}
 
-    pub fn receive_packet<P: ProtocolPacket + Debug>(&mut self, packet: impl Packet<P>) {
-        let buf = azalea_protocol::write::serialize_packet(&packet.into_variant()).unwrap();
-        self.incoming_packet_queue.lock().push(buf.into());
+pub fn make_basic_respawn_packet(
+    dimension_type: DimensionType,
+    dimension: ResourceLocation,
+) -> ClientboundRespawn {
+    ClientboundRespawn {
+        common: CommonPlayerSpawnInfo {
+            dimension_type,
+            dimension,
+            seed: 0,
+            game_type: GameMode::Survival,
+            previous_game_type: OptionalGameType(None),
+            is_debug: false,
+            is_flat: false,
+            last_death_location: None,
+            portal_cooldown: 0,
+            sea_level: 63,
+        },
+        data_to_keep: 0,
     }
+}
 
-    pub fn tick(&mut self) {
-        tick_app(&mut self.app);
-    }
-    pub fn component<T: Component + Clone>(&self) -> T {
-        self.app.world().get::<T>(self.entity).unwrap().clone()
-    }
-    pub fn get_component<T: Component + Clone>(&self) -> Option<T> {
-        self.app.world().get::<T>(self.entity).cloned()
-    }
-    pub fn has_component<T: Component>(&self) -> bool {
-        self.app.world().get::<T>(self.entity).is_some()
+pub fn make_basic_empty_chunk(
+    pos: ChunkPos,
+    section_count: usize,
+) -> ClientboundLevelChunkWithLight {
+    let mut chunk_bytes = Vec::new();
+    let mut sections = Vec::new();
+    for _ in 0..section_count {
+        sections.push(Section {
+            block_count: 0,
+            states: PalettedContainer::new(PalettedContainerKind::BlockStates),
+            biomes: PalettedContainer::new(PalettedContainerKind::Biomes),
+        });
     }
-}
+    sections.azalea_write(&mut chunk_bytes).unwrap();
 
-fn tick_app(app: &mut App) {
-    app.update();
-    app.world_mut().run_schedule(GameTick);
+    ClientboundLevelChunkWithLight {
+        x: pos.x,
+        z: pos.z,
+        chunk_data: ClientboundLevelChunkPacketData {
+            heightmaps: Nbt::None,
+            data: chunk_bytes,
+            block_entities: vec![],
+        },
+        light_data: ClientboundLightUpdatePacketData::default(),
+    }
 }
diff --git a/azalea-client/tests/change_dimension_to_nether_and_back.rs b/azalea-client/tests/change_dimension_to_nether_and_back.rs
new file mode 100644
index 00000000..3ed623ef
--- /dev/null
+++ b/azalea-client/tests/change_dimension_to_nether_and_back.rs
@@ -0,0 +1,147 @@
+use azalea_client::{test_simulation::*, InConfigState};
+use azalea_core::{position::ChunkPos, resource_location::ResourceLocation};
+use azalea_entity::{metadata::Health, LocalEntity};
+use azalea_protocol::packets::{
+    config::{ClientboundFinishConfiguration, ClientboundRegistryData},
+    game::ClientboundSetHealth,
+    ConnectionProtocol,
+};
+use azalea_registry::DimensionType;
+use azalea_world::InstanceName;
+use bevy_log::tracing_subscriber;
+use simdnbt::owned::{NbtCompound, NbtTag};
+
+#[test]
+fn test_change_dimension_to_nether_and_back() {
+    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![
+            (
+                // this dimension should never be created. it just exists to make sure we're not
+                // hard-coding the dimension type id anywhere.
+                ResourceLocation::new("azalea:fakedimension"),
+                Some(NbtCompound::from_values(vec![
+                    ("height".into(), NbtTag::Int(16)),
+                    ("min_y".into(), NbtTag::Int(0)),
+                ])),
+            ),
+            (
+                ResourceLocation::new("minecraft:overworld"),
+                Some(NbtCompound::from_values(vec![
+                    ("height".into(), NbtTag::Int(384)),
+                    ("min_y".into(), NbtTag::Int(-64)),
+                ])),
+            ),
+            (
+                ResourceLocation::new("minecraft:nether"),
+                Some(NbtCompound::from_values(vec![
+                    ("height".into(), NbtTag::Int(256)),
+                    ("min_y".into(), NbtTag::Int(0)),
+                ])),
+            ),
+        ]
+        .into_iter()
+        .collect(),
+    });
+    simulation.tick();
+    simulation.receive_packet(ClientboundFinishConfiguration);
+    simulation.tick();
+
+    assert!(!simulation.has_component::<InConfigState>());
+    assert!(simulation.has_component::<LocalEntity>());
+
+    simulation.receive_packet(ClientboundSetHealth {
+        health: 15.,
+        food: 20,
+        saturation: 20.,
+    });
+    simulation.tick();
+    assert_eq!(*simulation.component::<Health>(), 15.);
+
+    //
+    // OVERWORLD
+    //
+
+    simulation.receive_packet(make_basic_login_packet(
+        DimensionType::new_raw(1), // overworld
+        ResourceLocation::new("azalea:a"),
+    ));
+    simulation.tick();
+
+    assert_eq!(
+        *simulation.component::<InstanceName>(),
+        ResourceLocation::new("azalea:a"),
+        "InstanceName should be azalea:a after setting dimension to that"
+    );
+
+    simulation.receive_packet(make_basic_empty_chunk(ChunkPos::new(0, 0), (384 + 64) / 16));
+    simulation.tick();
+    // make sure the chunk exists
+    simulation
+        .chunk(ChunkPos::new(0, 0))
+        .expect("chunk should exist");
+
+    //
+    // NETHER
+    //
+
+    simulation.receive_packet(make_basic_respawn_packet(
+        DimensionType::new_raw(2), // nether
+        ResourceLocation::new("azalea:b"),
+    ));
+    simulation.tick();
+
+    assert!(
+        simulation.chunk(ChunkPos::new(0, 0)).is_none(),
+        "chunk should not exist immediately after changing dimensions"
+    );
+    assert_eq!(
+        *simulation.component::<InstanceName>(),
+        ResourceLocation::new("azalea:b"),
+        "InstanceName should be azalea:b after changing dimensions to that"
+    );
+
+    simulation.receive_packet(make_basic_empty_chunk(ChunkPos::new(0, 0), 256 / 16));
+    simulation.tick();
+    // make sure the chunk exists
+    simulation
+        .chunk(ChunkPos::new(0, 0))
+        .expect("chunk should exist");
+    simulation.receive_packet(make_basic_respawn_packet(
+        DimensionType::new_raw(2), // nether
+        ResourceLocation::new("minecraft:nether"),
+    ));
+    simulation.tick();
+
+    //
+    // BACK TO OVERWORLD
+    //
+
+    simulation.receive_packet(make_basic_login_packet(
+        DimensionType::new_raw(1), // overworld
+        ResourceLocation::new("azalea:a"),
+    ));
+    simulation.tick();
+
+    assert_eq!(
+        *simulation.component::<InstanceName>(),
+        ResourceLocation::new("azalea:a"),
+        "InstanceName should be azalea:a after setting dimension back to that"
+    );
+    assert!(
+        simulation.chunk(ChunkPos::new(0, 0)).is_none(),
+        "chunk should not exist immediately after switching back to overworld"
+    );
+
+    simulation.receive_packet(make_basic_empty_chunk(ChunkPos::new(0, 0), (384 + 64) / 16));
+    simulation.tick();
+    // make sure the chunk exists
+    simulation
+        .chunk(ChunkPos::new(0, 0))
+        .expect("chunk should exist");
+}
diff --git a/azalea-client/tests/set_health_before_login.rs b/azalea-client/tests/set_health_before_login.rs
new file mode 100644
index 00000000..d6b6c426
--- /dev/null
+++ b/azalea-client/tests/set_health_before_login.rs
@@ -0,0 +1,55 @@
+use azalea_client::{test_simulation::*, InConfigState};
+use azalea_core::resource_location::ResourceLocation;
+use azalea_entity::{metadata::Health, LocalEntity};
+use azalea_protocol::packets::{
+    config::{ClientboundFinishConfiguration, ClientboundRegistryData},
+    game::ClientboundSetHealth,
+    ConnectionProtocol,
+};
+use azalea_registry::DimensionType;
+use bevy_log::tracing_subscriber;
+use simdnbt::owned::{NbtCompound, NbtTag};
+
+#[test]
+fn test_set_health_before_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.tick();
+    simulation.receive_packet(ClientboundFinishConfiguration);
+    simulation.tick();
+
+    assert!(!simulation.has_component::<InConfigState>());
+    assert!(simulation.has_component::<LocalEntity>());
+
+    simulation.receive_packet(ClientboundSetHealth {
+        health: 15.,
+        food: 20,
+        saturation: 20.,
+    });
+    simulation.tick();
+    assert_eq!(*simulation.component::<Health>(), 15.);
+
+    simulation.receive_packet(make_basic_login_packet(
+        DimensionType::new_raw(0), // overworld
+        ResourceLocation::new("minecraft:overworld"),
+    ));
+    simulation.tick();
+
+    // health should stay the same
+    assert_eq!(*simulation.component::<Health>(), 15.);
+}