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

azalea-client/src/client.rs

@@ -319,10 +319,7 @@ impl Client {
         }
 
         as_system::<Commands>(&mut ecs_lock.lock(), |mut commands| {
-            commands.entity(entity).insert((
-                crate::packet::login::IgnoreQueryIds::default(),
-                InLoginState,
-            ));
+            commands.entity(entity).insert((InLoginState,));
             commands.trigger(SendLoginPacketEvent::new(
                 entity,
                 ServerboundHello {

azalea-client/src/plugins/connection.rs

@@ -29,7 +29,7 @@ use crate::packet::{config, game, login};
 pub struct ConnectionPlugin;
 impl Plugin for ConnectionPlugin {
     fn build(&self, app: &mut App) {
-        app.add_systems(PreUpdate, read_packets);
+        app.add_systems(PreUpdate, (read_packets, poll_all_writer_tasks).chain());
     }
 }
 
@@ -103,18 +103,22 @@ pub fn read_packets(ecs: &mut World) {
                 }
             }
         }
-
-        let mut net_conn = conn_query.get_mut(ecs, entity).unwrap();
-        if let Some(net_conn) = &mut net_conn.network {
-            // this needs to be done at some point every update, so we do it here right
-            // after the handlers are called
-            net_conn.poll_writer();
-        }
     }
 
     queued_packet_events.send_events(ecs);
 }
 
+fn poll_all_writer_tasks(mut conn_query: Query<&mut RawConnection>) {
+    for mut conn in conn_query.iter_mut() {
+        if let Some(net_conn) = &mut conn.network {
+            // this needs to be done at some point every update to make sure packets are
+            // actually sent to the network
+
+            net_conn.poll_writer();
+        }
+    }
+}
+
 #[derive(Default)]
 pub struct QueuedPacketEvents {
     login: Vec<ReceiveLoginPacketEvent>,
@@ -326,7 +330,6 @@ async fn write_task(
     mut write_half: OwnedWriteHalf,
 ) {
     while let Some(network_packet) = network_packet_writer_rx.recv().await {
-        trace!("writing encoded raw packet");
         if let Err(e) = write_half.write_all(&network_packet).await {
             debug!("Error writing packet to server: {e}");
             break;

azalea-client/src/plugins/login.rs

@@ -1,18 +1,24 @@
 use azalea_auth::sessionserver::ClientSessionServerError;
-use azalea_protocol::packets::login::{ClientboundHello, ServerboundKey};
+use azalea_protocol::packets::login::{
+    ClientboundHello, ServerboundCustomQueryAnswer, ServerboundKey,
+};
 use bevy_app::prelude::*;
 use bevy_ecs::prelude::*;
 use bevy_tasks::{IoTaskPool, Task, futures_lite::future};
-use tracing::{debug, error};
+use tracing::{debug, error, trace};
 
-use super::{connection::RawConnection, packet::login::ReceiveHelloEvent};
+use super::{
+    connection::RawConnection,
+    packet::login::{ReceiveCustomQueryEvent, ReceiveHelloEvent, SendLoginPacketEvent},
+};
 use crate::{Account, JoinError};
 
+/// Some systems that run during the `login` state.
 pub struct LoginPlugin;
 impl Plugin for LoginPlugin {
     fn build(&self, app: &mut App) {
         app.add_observer(handle_receive_hello_event)
-            .add_systems(Update, poll_auth_task);
+            .add_systems(Update, (poll_auth_task, reply_to_custom_queries));
     }
 }
 
@@ -124,3 +130,23 @@ pub async fn auth_with_account(
 
     Ok((key_packet, private_key))
 }
+
+pub fn reply_to_custom_queries(
+    mut commands: Commands,
+    mut events: EventReader<ReceiveCustomQueryEvent>,
+) {
+    for event in events.read() {
+        trace!("Maybe replying to custom query: {event:?}");
+        if event.disabled {
+            continue;
+        }
+
+        commands.trigger(SendLoginPacketEvent::new(
+            event.entity,
+            ServerboundCustomQueryAnswer {
+                transaction_id: event.packet.transaction_id,
+                data: None,
+            },
+        ));
+    }
+}

azalea-client/src/plugins/packet/login/events.rs

@@ -2,7 +2,9 @@ use std::sync::Arc;
 
 use azalea_protocol::packets::{
     Packet,
-    login::{ClientboundHello, ClientboundLoginPacket, ServerboundLoginPacket},
+    login::{
+        ClientboundCustomQuery, ClientboundHello, ClientboundLoginPacket, ServerboundLoginPacket,
+    },
 };
 use bevy_ecs::prelude::*;
 use tracing::{debug, error};
@@ -18,14 +20,27 @@ pub struct ReceiveLoginPacketEvent {
     pub packet: Arc<ClientboundLoginPacket>,
 }
 
-#[derive(Event)]
+#[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, Clone)]
+#[derive(Event, Debug, Clone)]
 pub struct SendLoginPacketEvent {
     pub sent_by: Entity,
     pub packet: ServerboundLoginPacket,

azalea-client/src/plugins/packet/login/mod.rs

@@ -10,8 +10,7 @@ use azalea_protocol::packets::{
     login::{
         ClientboundCookieRequest, ClientboundCustomQuery, ClientboundHello,
         ClientboundLoginCompression, ClientboundLoginDisconnect, ClientboundLoginFinished,
-        ClientboundLoginPacket, ServerboundCookieResponse, ServerboundCustomQueryAnswer,
-        ServerboundLoginAcknowledged,
+        ClientboundLoginPacket, ServerboundCookieResponse, ServerboundLoginAcknowledged,
     },
 };
 use bevy_ecs::prelude::*;
@@ -48,11 +47,6 @@ pub fn process_packet(ecs: &mut World, player: Entity, packet: &ClientboundLogin
 #[derive(Component, Clone, Debug)]
 pub struct InLoginState;
 
-/// Plugins can add to this set if they want to handle a custom query packet
-/// themselves. This component removed after the login state ends.
-#[derive(Component, Default, Debug, Deref, DerefMut)]
-pub struct IgnoreQueryIds(HashSet<u32>);
-
 pub struct LoginPacketHandler<'a> {
     pub ecs: &'a mut World,
     pub player: Entity,
@@ -103,7 +97,6 @@ impl LoginPacketHandler<'_> {
 
                 commands
                     .entity(self.player)
-                    .remove::<IgnoreQueryIds>()
                     .remove::<InLoginState>()
                     .insert(InConfigState)
                     .insert(GameProfileComponent(p.game_profile.clone()));
@@ -130,21 +123,12 @@ impl LoginPacketHandler<'_> {
     pub fn custom_query(&mut self, p: &ClientboundCustomQuery) {
         debug!("Got custom query {p:?}");
 
-        as_system::<(Commands, Query<&IgnoreQueryIds>)>(self.ecs, |(mut commands, query)| {
-            let ignore_query_ids = query.get(self.player).ok().map(|x| x.0.clone());
-            if let Some(ignore_query_ids) = ignore_query_ids {
-                if ignore_query_ids.contains(&p.transaction_id) {
-                    return;
-                }
-            }
-
-            commands.trigger(SendLoginPacketEvent::new(
-                self.player,
-                ServerboundCustomQueryAnswer {
-                    transaction_id: p.transaction_id,
-                    data: None,
-                },
-            ));
+        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) {

azalea-client/src/plugins/packet/mod.rs

@@ -5,10 +5,7 @@ use bevy_ecs::{
     system::{SystemParam, SystemState},
 };
 
-use self::game::{
-    AddPlayerEvent, DeathEvent, InstanceLoadedEvent, KeepAliveEvent, RemovePlayerEvent,
-    ResourcePackEvent, UpdatePlayerEvent,
-};
+use self::game::DeathEvent;
 use crate::{chat::ChatReceivedEvent, events::death_listener};
 
 pub mod config;
@@ -56,14 +53,15 @@ impl Plugin for PacketPlugin {
             .add_event::<config::SendConfigPacketEvent>()
             .add_event::<login::SendLoginPacketEvent>()
             //
-            .add_event::<AddPlayerEvent>()
-            .add_event::<RemovePlayerEvent>()
-            .add_event::<UpdatePlayerEvent>()
+            .add_event::<game::AddPlayerEvent>()
+            .add_event::<game::RemovePlayerEvent>()
+            .add_event::<game::UpdatePlayerEvent>()
             .add_event::<ChatReceivedEvent>()
-            .add_event::<DeathEvent>()
-            .add_event::<KeepAliveEvent>()
-            .add_event::<ResourcePackEvent>()
-            .add_event::<InstanceLoadedEvent>();
+            .add_event::<game::DeathEvent>()
+            .add_event::<game::KeepAliveEvent>()
+            .add_event::<game::ResourcePackEvent>()
+            .add_event::<game::InstanceLoadedEvent>()
+            .add_event::<login::ReceiveCustomQueryEvent>();
     }
 }