make it so plugins can send and receive packets during the login state

4 files changed, 94 insertions, 16 deletions

diff --git a/azalea-client/src/packet_handling/configuration.rs b/azalea-client/src/packet_handling/configuration.rs
index 0285884e..e7d3ffd8 100644
--- a/azalea-client/src/packet_handling/configuration.rs
+++ b/azalea-client/src/packet_handling/configuration.rs
@@ -22,7 +22,7 @@ use crate::packet_handling::game::KeepAliveEvent;
 use crate::raw_connection::RawConnection;
 
 #[derive(Event, Debug, Clone)]
-pub struct PacketEvent {
+pub struct ConfigurationPacketEvent {
     /// The client entity that received the packet.
     pub entity: Entity,
     /// The packet that was actually received.
@@ -31,7 +31,7 @@ pub struct PacketEvent {
 
 pub fn send_packet_events(
     query: Query<(Entity, &RawConnection), With<InConfigurationState>>,
-    mut packet_events: ResMut<Events<PacketEvent>>,
+    mut packet_events: ResMut<Events<ConfigurationPacketEvent>>,
 ) {
     // 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
@@ -51,7 +51,7 @@ pub fn send_packet_events(
                         continue;
                     }
                 };
-                packet_events.send(PacketEvent {
+                packet_events.send(ConfigurationPacketEvent {
                     entity: player_entity,
                     packet,
                 });
@@ -64,9 +64,10 @@ pub fn send_packet_events(
 
 pub fn process_packet_events(ecs: &mut World) {
     let mut events_owned = Vec::new();
-    let mut system_state: SystemState<EventReader<PacketEvent>> = SystemState::new(ecs);
+    let mut system_state: SystemState<EventReader<ConfigurationPacketEvent>> =
+        SystemState::new(ecs);
     let mut events = system_state.get_mut(ecs);
-    for PacketEvent {
+    for ConfigurationPacketEvent {
         entity: player_entity,
         packet,
     } in events.read()
diff --git a/azalea-client/src/packet_handling/game.rs b/azalea-client/src/packet_handling/game.rs
index cffa0d2b..20facadb 100644
--- a/azalea-client/src/packet_handling/game.rs
+++ b/azalea-client/src/packet_handling/game.rs
@@ -23,6 +23,7 @@ use azalea_protocol::{
         serverbound_keep_alive_packet::ServerboundKeepAlivePacket,
         serverbound_move_player_pos_rot_packet::ServerboundMovePlayerPosRotPacket,
         serverbound_pong_packet::ServerboundPongPacket, ClientboundGamePacket,
+        ServerboundGamePacket,
     },
     read::deserialize_packet,
 };
@@ -40,8 +41,7 @@ use crate::{
         SetContainerContentEvent,
     },
     local_player::{
-        GameProfileComponent, Hunger, InstanceHolder, LocalGameMode, PlayerAbilities,
-        SendPacketEvent, TabList,
+        GameProfileComponent, Hunger, InstanceHolder, LocalGameMode, PlayerAbilities, TabList,
     },
     movement::{KnockbackEvent, KnockbackType},
     raw_connection::RawConnection,
@@ -1391,3 +1391,24 @@ pub fn process_packet_events(ecs: &mut World) {
         }
     }
 }
+
+/// An event for sending a packet to the server while we're in the `game` state.
+#[derive(Event)]
+pub struct SendPacketEvent {
+    pub entity: Entity,
+    pub packet: ServerboundGamePacket,
+}
+
+pub fn handle_send_packet_event(
+    mut send_packet_events: EventReader<SendPacketEvent>,
+    mut query: Query<&mut RawConnection>,
+) {
+    for event in send_packet_events.read() {
+        if let Ok(raw_connection) = query.get_mut(event.entity) {
+            // debug!("Sending packet: {:?}", event.packet);
+            if let Err(e) = raw_connection.write_packet(event.packet.clone()) {
+                error!("Failed to send packet: {e}");
+            }
+        }
+    }
+}
diff --git a/azalea-client/src/packet_handling/login.rs b/azalea-client/src/packet_handling/login.rs
new file mode 100644
index 00000000..83a6a98d
--- /dev/null
+++ b/azalea-client/src/packet_handling/login.rs
@@ -0,0 +1,49 @@
+// login packets aren't actually handled here because compression/encryption
+// would make packet handling a lot messier
+
+use std::sync::Arc;
+
+use azalea_protocol::packets::login::{ClientboundLoginPacket, ServerboundLoginPacket};
+use bevy_ecs::prelude::*;
+use tokio::sync::mpsc;
+use tracing::error;
+
+use crate::raw_connection::RawConnection;
+
+use super::game::SendPacketEvent;
+
+// this struct is defined here anyways though so it's consistent with the other
+// ones
+
+#[derive(Event, Debug, Clone)]
+pub struct LoginPacketEvent {
+    /// The client entity that received the packet.
+    pub entity: Entity,
+    /// The packet that was actually received.
+    pub packet: Arc<ClientboundLoginPacket>,
+}
+
+/// Event for sending a login packet to the server.
+#[derive(Event)]
+pub struct SendLoginPacketEvent {
+    pub entity: Entity,
+    pub packet: ServerboundLoginPacket,
+}
+
+#[derive(Component)]
+pub struct LoginSendPacketQueue {
+    pub tx: mpsc::UnboundedSender<ServerboundLoginPacket>,
+}
+
+pub fn handle_send_packet_event(
+    mut send_packet_events: EventReader<SendLoginPacketEvent>,
+    mut query: Query<&mut LoginSendPacketQueue>,
+) {
+    for event in send_packet_events.read() {
+        if let Ok(queue) = query.get_mut(event.entity) {
+            let _ = queue.tx.send(event.packet.clone());
+        } else {
+            error!("Sent SendPacketEvent for entity that doesn't have a LoginSendPacketQueue");
+        }
+    }
+}
diff --git a/azalea-client/src/packet_handling/mod.rs b/azalea-client/src/packet_handling/mod.rs
index 35bdfc04..7f797a5c 100644
--- a/azalea-client/src/packet_handling/mod.rs
+++ b/azalea-client/src/packet_handling/mod.rs
@@ -4,13 +4,17 @@ use bevy_ecs::prelude::*;
 
 use crate::{chat::ChatReceivedEvent, events::death_listener};
 
-use self::game::{
-    AddPlayerEvent, DeathEvent, InstanceLoadedEvent, KeepAliveEvent, RemovePlayerEvent,
-    ResourcePackEvent, UpdatePlayerEvent,
+use self::{
+    game::{
+        AddPlayerEvent, DeathEvent, InstanceLoadedEvent, KeepAliveEvent, RemovePlayerEvent,
+        ResourcePackEvent, UpdatePlayerEvent,
+    },
+    login::{LoginPacketEvent, SendLoginPacketEvent},
 };
 
 pub mod configuration;
 pub mod game;
+pub mod login;
 
 pub struct PacketHandlerPlugin;
 
@@ -37,16 +41,17 @@ impl Plugin for PacketHandlerPlugin {
         .add_systems(
             PreUpdate,
             (
-                game::process_packet_events,
+                game::process_packet_events
+                    // we want to index and deindex right after
+                    .before(EntityUpdateSet::Deindex),
                 configuration::process_packet_events,
-            )
-                // we want to index and deindex right after
-                .before(EntityUpdateSet::Deindex),
+                login::handle_send_packet_event,
+            ),
         )
         .add_systems(Update, death_event_on_0_health.before(death_listener))
         // we do this instead of add_event so we can handle the events ourselves
         .init_resource::<Events<game::PacketEvent>>()
-        .init_resource::<Events<configuration::PacketEvent>>()
+        .init_resource::<Events<configuration::ConfigurationPacketEvent>>()
         .add_event::<AddPlayerEvent>()
         .add_event::<RemovePlayerEvent>()
         .add_event::<UpdatePlayerEvent>()
@@ -54,6 +59,8 @@ impl Plugin for PacketHandlerPlugin {
         .add_event::<DeathEvent>()
         .add_event::<KeepAliveEvent>()
         .add_event::<ResourcePackEvent>()
-        .add_event::<InstanceLoadedEvent>();
+        .add_event::<InstanceLoadedEvent>()
+        .add_event::<LoginPacketEvent>()
+        .add_event::<SendLoginPacketEvent>();
     }
 }