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, 231 insertions, 0 deletions

diff --git a/azalea-client/src/plugins/packet/login/events.rs b/azalea-client/src/plugins/packet/login/events.rs
new file mode 100644
index 00000000..fc7a6b22
--- /dev/null
+++ b/azalea-client/src/plugins/packet/login/events.rs
@@ -0,0 +1,86 @@
+use std::sync::Arc;
+
+use azalea_protocol::packets::{
+    Packet,
+    login::{
+        ClientboundCustomQuery, ClientboundHello, ClientboundLoginPacket, ServerboundLoginPacket,
+    },
+};
+use bevy_ecs::prelude::*;
+use tracing::{debug, error};
+
+use super::InLoginState;
+use crate::{Account, connection::RawConnection};
+
+#[derive(Event, Debug, Clone)]
+pub struct ReceiveLoginPacketEvent {
+    /// The client entity that received the packet.
+    pub entity: Entity,
+    /// The packet that was actually received.
+    pub packet: Arc<ClientboundLoginPacket>,
+}
+
+#[derive(Event, Debug, Clone)]
+pub struct ReceiveHelloEvent {
+    pub account: Account,
+    pub packet: ClientboundHello,
+}
+
+#[derive(Event, Debug, Clone)]
+pub struct ReceiveCustomQueryEvent {
+    /// The client entity that received the packet.
+    pub entity: Entity,
+    pub packet: ClientboundCustomQuery,
+    /// A system can set this to `true` to make Azalea not reply to the query.
+    /// You must make sure you modify this before the
+    /// [`reply_to_custom_queries`] system runs.
+    ///
+    /// [`reply_to_custom_queries`]: crate::login::reply_to_custom_queries
+    pub disabled: bool,
+}
+
+/// Event for sending a login packet to the server.
+#[derive(Event, Debug, Clone)]
+pub struct SendLoginPacketEvent {
+    pub sent_by: Entity,
+    pub packet: ServerboundLoginPacket,
+}
+impl SendLoginPacketEvent {
+    pub fn new(entity: Entity, packet: impl Packet<ServerboundLoginPacket>) -> Self {
+        let packet = packet.into_variant();
+        Self {
+            sent_by: entity,
+            packet,
+        }
+    }
+}
+
+pub fn handle_outgoing_packets_observer(
+    trigger: Trigger<SendLoginPacketEvent>,
+    mut query: Query<(&mut RawConnection, Option<&InLoginState>)>,
+) {
+    let event = trigger.event();
+    if let Ok((mut raw_conn, in_login_state)) = query.get_mut(event.sent_by) {
+        if in_login_state.is_none() {
+            error!(
+                "Tried to send a login packet {:?} while not in login state",
+                event.packet
+            );
+            return;
+        }
+        debug!("Sending login packet: {:?}", event.packet);
+        if let Err(e) = raw_conn.write(event.packet.clone()) {
+            error!("Failed to send packet: {e}");
+        }
+    }
+}
+/// A system that converts [`SendLoginPacketEvent`] events into triggers so
+/// they get received by [`handle_outgoing_packets_observer`].
+pub fn handle_outgoing_packets(
+    mut commands: Commands,
+    mut events: EventReader<SendLoginPacketEvent>,
+) {
+    for event in events.read() {
+        commands.trigger(event.clone());
+    }
+}
diff --git a/azalea-client/src/plugins/packet/login/mod.rs b/azalea-client/src/plugins/packet/login/mod.rs
new file mode 100644
index 00000000..d313a767
--- /dev/null
+++ b/azalea-client/src/plugins/packet/login/mod.rs
@@ -0,0 +1,145 @@
+// login packets aren't actually handled here because compression/encryption
+// would make packet handling a lot messier
+
+mod events;
+
+use azalea_protocol::packets::{
+    ConnectionProtocol,
+    login::{
+        ClientboundCookieRequest, ClientboundCustomQuery, ClientboundHello,
+        ClientboundLoginCompression, ClientboundLoginDisconnect, ClientboundLoginFinished,
+        ClientboundLoginPacket, ServerboundCookieResponse, ServerboundLoginAcknowledged,
+    },
+};
+use bevy_ecs::prelude::*;
+pub use events::*;
+use tracing::{debug, error};
+
+use super::as_system;
+use crate::{
+    Account, GameProfileComponent, InConfigState, connection::RawConnection,
+    declare_packet_handlers, disconnect::DisconnectEvent,
+};
+
+pub fn process_packet(ecs: &mut World, player: Entity, packet: &ClientboundLoginPacket) {
+    let mut handler = LoginPacketHandler { player, ecs };
+
+    declare_packet_handlers!(
+        ClientboundLoginPacket,
+        packet,
+        handler,
+        [
+            hello,
+            login_disconnect,
+            login_finished,
+            login_compression,
+            custom_query,
+            cookie_request
+        ]
+    );
+}
+
+/// A marker component for local players that are currently in the
+/// `login` state.
+#[derive(Component, Clone, Debug)]
+pub struct InLoginState;
+
+pub struct LoginPacketHandler<'a> {
+    pub ecs: &'a mut World,
+    pub player: Entity,
+}
+impl LoginPacketHandler<'_> {
+    pub fn hello(&mut self, p: &ClientboundHello) {
+        debug!("Got encryption request {p:?}");
+
+        as_system::<(Commands, Query<&Account>)>(self.ecs, |(mut commands, query)| {
+            let Ok(account) = query.get(self.player) else {
+                error!(
+                    "Expected Account component to be present on player when receiving hello packet."
+                );
+                return;
+            };
+            commands.trigger_targets(
+                ReceiveHelloEvent {
+                    account: account.clone(),
+                    packet: p.clone(),
+                },
+                self.player,
+            );
+        });
+    }
+    pub fn login_disconnect(&mut self, p: &ClientboundLoginDisconnect) {
+        debug!("Got disconnect {:?}", p);
+
+        as_system::<EventWriter<_>>(self.ecs, |mut events| {
+            events.send(DisconnectEvent {
+                entity: self.player,
+                reason: Some(p.reason.clone()),
+            });
+        });
+    }
+    pub fn login_finished(&mut self, p: &ClientboundLoginFinished) {
+        debug!(
+            "Got profile {:?}. login is finished and we're now switching to the config state",
+            p.game_profile
+        );
+
+        as_system::<(Commands, Query<&mut RawConnection>)>(
+            self.ecs,
+            |(mut commands, mut query)| {
+                commands.trigger(SendLoginPacketEvent::new(
+                    self.player,
+                    ServerboundLoginAcknowledged,
+                ));
+
+                commands
+                    .entity(self.player)
+                    .remove::<InLoginState>()
+                    .insert(InConfigState)
+                    .insert(GameProfileComponent(p.game_profile.clone()));
+
+                let mut conn = query
+                    .get_mut(self.player)
+                    .expect("RawConnection component should be present when receiving packets");
+                conn.state = ConnectionProtocol::Configuration;
+            },
+        );
+    }
+    pub fn login_compression(&mut self, p: &ClientboundLoginCompression) {
+        debug!("Got compression request {p:?}");
+
+        as_system::<Query<&mut RawConnection>>(self.ecs, |mut query| {
+            let mut conn = query
+                .get_mut(self.player)
+                .expect("RawConnection component should be present when receiving packets");
+            if let Some(net_conn) = &mut conn.net_conn() {
+                net_conn.set_compression_threshold(Some(p.compression_threshold as u32));
+            }
+        })
+    }
+    pub fn custom_query(&mut self, p: &ClientboundCustomQuery) {
+        debug!("Got custom query {p:?}");
+
+        as_system::<EventWriter<ReceiveCustomQueryEvent>>(self.ecs, |mut events| {
+            events.send(ReceiveCustomQueryEvent {
+                entity: self.player,
+                packet: p.clone(),
+                disabled: false,
+            });
+        });
+    }
+    pub fn cookie_request(&mut self, p: &ClientboundCookieRequest) {
+        debug!("Got cookie request {p:?}");
+
+        as_system::<Commands>(self.ecs, |mut commands| {
+            commands.trigger(SendLoginPacketEvent::new(
+                self.player,
+                ServerboundCookieResponse {
+                    key: p.key.clone(),
+                    // cookies aren't implemented
+                    payload: None,
+                },
+            ));
+        });
+    }
+}