Move login state to the ECS (#213)

* use packet handlers code for login custom_query

* initial broken implementation for ecs-only login

* fixes

* run Update schedule 60 times per second and delete code related to run_schedule_sender

* fix tests

* fix online-mode

* reply to query packets in a separate system and make it easier for plugins to disable individual replies

* remove unused imports

2 files changed, 96 insertions, 160 deletions

diff --git a/azalea-client/src/plugins/packet/config/events.rs b/azalea-client/src/plugins/packet/config/events.rs
index 24a1157b..a9237e75 100644
--- a/azalea-client/src/plugins/packet/config/events.rs
+++ b/azalea-client/src/plugins/packet/config/events.rs
@@ -1,23 +1,20 @@
-use std::io::Cursor;
+use std::sync::Arc;
 
-use azalea_protocol::{
-    packets::{
-        Packet,
-        config::{ClientboundConfigPacket, ServerboundConfigPacket},
-    },
-    read::deserialize_packet,
+use azalea_protocol::packets::{
+    Packet,
+    config::{ClientboundConfigPacket, ServerboundConfigPacket},
 };
 use bevy_ecs::prelude::*;
 use tracing::{debug, error};
 
-use crate::{InConfigState, raw_connection::RawConnection};
+use crate::{InConfigState, connection::RawConnection};
 
 #[derive(Event, Debug, Clone)]
 pub struct ReceiveConfigPacketEvent {
     /// The client entity that received the packet.
     pub entity: Entity,
     /// The packet that was actually received.
-    pub packet: ClientboundConfigPacket,
+    pub packet: Arc<ClientboundConfigPacket>,
 }
 
 /// An event for sending a packet to the server while we're in the
@@ -39,7 +36,7 @@ pub fn handle_outgoing_packets_observer(
     mut query: Query<(&mut RawConnection, Option<&InConfigState>)>,
 ) {
     let event = trigger.event();
-    if let Ok((raw_conn, in_configuration_state)) = query.get_mut(event.sent_by) {
+    if let Ok((mut raw_conn, in_configuration_state)) = query.get_mut(event.sent_by) {
         if in_configuration_state.is_none() {
             error!(
                 "Tried to send a configuration packet {:?} while not in configuration state",
@@ -47,8 +44,8 @@ pub fn handle_outgoing_packets_observer(
             );
             return;
         }
-        debug!("Sending packet: {:?}", event.packet);
-        if let Err(e) = raw_conn.write_packet(event.packet.clone()) {
+        debug!("Sending config packet: {:?}", event.packet);
+        if let Err(e) = raw_conn.write(event.packet.clone()) {
             error!("Failed to send packet: {e}");
         }
     }
@@ -64,61 +61,6 @@ pub fn handle_outgoing_packets(
     }
 }
 
-pub fn emit_receive_config_packet_events(
-    query: Query<(Entity, &RawConnection), With<InConfigState>>,
-    mut packet_events: ResMut<Events<ReceiveConfigPacketEvent>>,
-) {
-    // we manually clear and send the events at the beginning of each update
-    // since otherwise it'd cause issues with events in process_packet_events
-    // running twice
-    packet_events.clear();
-    for (player_entity, raw_conn) in &query {
-        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,
-                )) {
-                    Ok(packet) => packet,
-                    Err(err) => {
-                        error!("failed to read packet: {err:?}");
-                        debug!("packet bytes: {raw_packet:?}");
-                        continue;
-                    }
-                };
-
-                let should_interrupt = packet_interrupts(&packet);
-
-                packet_events.send(ReceiveConfigPacketEvent {
-                    entity: player_entity,
-                    packet,
-                });
-
-                if should_interrupt {
-                    break;
-                }
-            }
-            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(_)
-    )
-}
-
 /// A Bevy trigger that's sent when our client receives a [`ClientboundPing`]
 /// packet in the config state.
 ///
diff --git a/azalea-client/src/plugins/packet/config/mod.rs b/azalea-client/src/plugins/packet/config/mod.rs
index ae601793..910019a6 100644
--- a/azalea-client/src/plugins/packet/config/mod.rs
+++ b/azalea-client/src/plugins/packet/config/mod.rs
@@ -1,65 +1,61 @@
 mod events;
 
+use std::io::Cursor;
+
 use azalea_entity::LocalEntity;
 use azalea_protocol::packets::ConnectionProtocol;
 use azalea_protocol::packets::config::*;
+use azalea_protocol::read::ReadPacketError;
+use azalea_protocol::read::deserialize_packet;
 use bevy_ecs::prelude::*;
-use bevy_ecs::system::SystemState;
 pub use events::*;
 use tracing::{debug, warn};
 
 use super::as_system;
 use crate::client::InConfigState;
+use crate::connection::RawConnection;
 use crate::disconnect::DisconnectEvent;
 use crate::packet::game::KeepAliveEvent;
 use crate::packet::game::ResourcePackEvent;
-use crate::raw_connection::RawConnection;
 use crate::{InstanceHolder, declare_packet_handlers};
 
-pub fn process_packet_events(ecs: &mut World) {
-    let mut events_owned = Vec::new();
-    let mut system_state: SystemState<EventReader<ReceiveConfigPacketEvent>> =
-        SystemState::new(ecs);
-    let mut events = system_state.get_mut(ecs);
-    for ReceiveConfigPacketEvent {
-        entity: player_entity,
+pub fn process_raw_packet(
+    ecs: &mut World,
+    player: Entity,
+    raw_packet: &[u8],
+) -> Result<(), Box<ReadPacketError>> {
+    let packet = deserialize_packet(&mut Cursor::new(raw_packet))?;
+    process_packet(ecs, player, &packet);
+    Ok(())
+}
+
+pub fn process_packet(ecs: &mut World, player: Entity, packet: &ClientboundConfigPacket) {
+    let mut handler = ConfigPacketHandler { player, ecs };
+
+    declare_packet_handlers!(
+        ClientboundConfigPacket,
         packet,
-    } in events.read()
-    {
-        // we do this so `ecs` isn't borrowed for the whole loop
-        events_owned.push((*player_entity, packet.clone()));
-    }
-    for (player_entity, packet) in events_owned {
-        let mut handler = ConfigPacketHandler {
-            player: player_entity,
-            ecs,
-        };
-
-        declare_packet_handlers!(
-            ClientboundConfigPacket,
-            packet,
-            handler,
-            [
-                cookie_request,
-                custom_payload,
-                disconnect,
-                finish_configuration,
-                keep_alive,
-                ping,
-                reset_chat,
-                registry_data,
-                resource_pack_pop,
-                resource_pack_push,
-                store_cookie,
-                transfer,
-                update_enabled_features,
-                update_tags,
-                select_known_packs,
-                custom_report_details,
-                server_links,
-            ]
-        );
-    }
+        handler,
+        [
+            cookie_request,
+            custom_payload,
+            disconnect,
+            finish_configuration,
+            keep_alive,
+            ping,
+            reset_chat,
+            registry_data,
+            resource_pack_pop,
+            resource_pack_push,
+            store_cookie,
+            transfer,
+            update_enabled_features,
+            update_tags,
+            select_known_packs,
+            custom_report_details,
+            server_links,
+        ]
+    );
 }
 
 pub struct ConfigPacketHandler<'a> {
@@ -67,44 +63,45 @@ pub struct ConfigPacketHandler<'a> {
     pub player: Entity,
 }
 impl ConfigPacketHandler<'_> {
-    pub fn registry_data(&mut self, p: ClientboundRegistryData) {
+    pub fn registry_data(&mut self, p: &ClientboundRegistryData) {
         as_system::<Query<&mut InstanceHolder>>(self.ecs, |mut query| {
             let instance_holder = query.get_mut(self.player).unwrap();
             let mut instance = instance_holder.instance.write();
 
             // add the new registry data
-            instance.registries.append(p.registry_id, p.entries);
+            instance
+                .registries
+                .append(p.registry_id.clone(), p.entries.clone());
         });
     }
 
-    pub fn custom_payload(&mut self, p: ClientboundCustomPayload) {
+    pub fn custom_payload(&mut self, p: &ClientboundCustomPayload) {
         debug!("Got custom payload packet {p:?}");
     }
 
-    pub fn disconnect(&mut self, p: ClientboundDisconnect) {
+    pub fn disconnect(&mut self, p: &ClientboundDisconnect) {
         warn!("Got disconnect packet {p:?}");
         as_system::<EventWriter<_>>(self.ecs, |mut events| {
             events.send(DisconnectEvent {
                 entity: self.player,
-                reason: Some(p.reason),
+                reason: Some(p.reason.clone()),
             });
         });
     }
 
-    pub fn finish_configuration(&mut self, p: ClientboundFinishConfiguration) {
-        debug!("got FinishConfiguration packet: {p:?}");
+    pub fn finish_configuration(&mut self, _p: &ClientboundFinishConfiguration) {
+        debug!("got FinishConfiguration packet");
 
         as_system::<(Commands, Query<&mut RawConnection>)>(
             self.ecs,
             |(mut commands, mut query)| {
                 let mut raw_conn = query.get_mut(self.player).unwrap();
 
-                raw_conn
-                    .write_packet(ServerboundFinishConfiguration)
-                    .expect(
-                        "we should be in the right state and encoding this packet shouldn't fail",
-                    );
-                raw_conn.set_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
@@ -120,34 +117,33 @@ impl ConfigPacketHandler<'_> {
         );
     }
 
-    pub fn keep_alive(&mut self, p: ClientboundKeepAlive) {
+    pub fn keep_alive(&mut self, p: &ClientboundKeepAlive) {
         debug!(
             "Got keep alive packet (in configuration) {p:?} for {:?}",
             self.player
         );
 
-        as_system::<(Query<&RawConnection>, EventWriter<_>)>(self.ecs, |(query, mut events)| {
-            let raw_conn = query.get(self.player).unwrap();
-
+        as_system::<(Commands, EventWriter<_>)>(self.ecs, |(mut commands, mut events)| {
             events.send(KeepAliveEvent {
                 entity: self.player,
                 id: p.id,
             });
-            raw_conn
-                .write_packet(ServerboundKeepAlive { id: p.id })
-                .unwrap();
+            commands.trigger(SendConfigPacketEvent::new(
+                self.player,
+                ServerboundKeepAlive { id: p.id },
+            ));
         });
     }
 
-    pub fn ping(&mut self, p: ClientboundPing) {
+    pub fn ping(&mut self, p: &ClientboundPing) {
         debug!("Got ping packet (in configuration) {p:?}");
 
         as_system::<Commands>(self.ecs, |mut commands| {
-            commands.trigger_targets(ConfigPingEvent(p), self.player);
+            commands.trigger_targets(ConfigPingEvent(p.clone()), self.player);
         });
     }
 
-    pub fn resource_pack_push(&mut self, p: ClientboundResourcePackPush) {
+    pub fn resource_pack_push(&mut self, p: &ClientboundResourcePackPush) {
         debug!("Got resource pack push packet {p:?}");
 
         as_system::<EventWriter<_>>(self.ecs, |mut events| {
@@ -162,66 +158,64 @@ impl ConfigPacketHandler<'_> {
         });
     }
 
-    pub fn resource_pack_pop(&mut self, p: ClientboundResourcePackPop) {
+    pub fn resource_pack_pop(&mut self, p: &ClientboundResourcePackPop) {
         debug!("Got resource pack pop packet {p:?}");
     }
 
-    pub fn update_enabled_features(&mut self, p: ClientboundUpdateEnabledFeatures) {
+    pub fn update_enabled_features(&mut self, p: &ClientboundUpdateEnabledFeatures) {
         debug!("Got update enabled features packet {p:?}");
     }
 
-    pub fn update_tags(&mut self, _p: ClientboundUpdateTags) {
+    pub fn update_tags(&mut self, _p: &ClientboundUpdateTags) {
         debug!("Got update tags packet");
     }
 
-    pub fn cookie_request(&mut self, p: ClientboundCookieRequest) {
+    pub fn cookie_request(&mut self, p: &ClientboundCookieRequest) {
         debug!("Got cookie request packet {p:?}");
 
-        as_system::<Query<&RawConnection>>(self.ecs, |query| {
-            let raw_conn = query.get(self.player).unwrap();
-
-            raw_conn
-                .write_packet(ServerboundCookieResponse {
-                    key: p.key,
+        as_system::<Commands>(self.ecs, |mut commands| {
+            commands.trigger(SendConfigPacketEvent::new(
+                self.player,
+                ServerboundCookieResponse {
+                    key: p.key.clone(),
                     // cookies aren't implemented
                     payload: None,
-                })
-                .unwrap();
+                },
+            ));
         });
     }
 
-    pub fn reset_chat(&mut self, p: ClientboundResetChat) {
+    pub fn reset_chat(&mut self, p: &ClientboundResetChat) {
         debug!("Got reset chat packet {p:?}");
     }
 
-    pub fn store_cookie(&mut self, p: ClientboundStoreCookie) {
+    pub fn store_cookie(&mut self, p: &ClientboundStoreCookie) {
         debug!("Got store cookie packet {p:?}");
     }
 
-    pub fn transfer(&mut self, p: ClientboundTransfer) {
+    pub fn transfer(&mut self, p: &ClientboundTransfer) {
         debug!("Got transfer packet {p:?}");
     }
 
-    pub fn select_known_packs(&mut self, p: ClientboundSelectKnownPacks) {
+    pub fn select_known_packs(&mut self, p: &ClientboundSelectKnownPacks) {
         debug!("Got select known packs packet {p:?}");
 
-        as_system::<Query<&RawConnection>>(self.ecs, |query| {
-            let raw_conn = query.get(self.player).unwrap();
-
+        as_system::<Commands>(self.ecs, |mut commands| {
             // resource pack management isn't implemented
-            raw_conn
-                .write_packet(ServerboundSelectKnownPacks {
+            commands.trigger(SendConfigPacketEvent::new(
+                self.player,
+                ServerboundSelectKnownPacks {
                     known_packs: vec![],
-                })
-                .unwrap();
+                },
+            ));
         });
     }
 
-    pub fn server_links(&mut self, p: ClientboundServerLinks) {
+    pub fn server_links(&mut self, p: &ClientboundServerLinks) {
         debug!("Got server links packet {p:?}");
     }
 
-    pub fn custom_report_details(&mut self, p: ClientboundCustomReportDetails) {
+    pub fn custom_report_details(&mut self, p: &ClientboundCustomReportDetails) {
         debug!("Got custom report details packet {p:?}");
     }
 }