let plugins override query responses

3 files changed, 80 insertions, 20 deletions

diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 8757bc39..567d7f47 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -48,7 +48,6 @@ use azalea_protocol::{
             ServerboundHandshakePacket,
         },
         login::{
-            serverbound_custom_query_answer_packet::ServerboundCustomQueryAnswerPacket,
             serverbound_hello_packet::ServerboundHelloPacket,
             serverbound_key_packet::ServerboundKeyPacket,
             serverbound_login_acknowledged_packet::ServerboundLoginAcknowledgedPacket,
@@ -218,7 +217,7 @@ impl Client {
 
             let entity_uuid_index = ecs.resource::<EntityUuidIndex>();
             let uuid = account.uuid_or_offline();
-            if let Some(entity) = entity_uuid_index.get(&account.uuid_or_offline()) {
+            let entity = if let Some(entity) = entity_uuid_index.get(&account.uuid_or_offline()) {
                 debug!("Reusing entity {entity:?} for client");
                 entity
             } else {
@@ -228,7 +227,12 @@ impl Client {
                 let mut entity_uuid_index = ecs.resource_mut::<EntityUuidIndex>();
                 entity_uuid_index.insert(uuid, entity);
                 entity
-            }
+            };
+
+            // add the Account to the entity now so plugins can access it earlier
+            ecs.entity_mut(entity).insert(account.to_owned());
+
+            entity
         };
 
         let conn = Connection::new(resolved_address).await?;
@@ -279,7 +283,6 @@ impl Client {
                 local_player_events: LocalPlayerEvents(tx),
                 game_profile: GameProfileComponent(game_profile),
                 client_information: crate::ClientInformation::default(),
-                account: account.to_owned(),
             },
             InConfigurationState,
         ));
@@ -321,10 +324,10 @@ impl Client {
         // this makes it so plugins can send an `SendLoginPacketEvent` event to the ecs
         // and we'll send it to the server
         let (ecs_packets_tx, mut ecs_packets_rx) = mpsc::unbounded_channel();
-        ecs_lock
-            .lock()
-            .entity_mut(entity)
-            .insert(LoginSendPacketQueue { tx: ecs_packets_tx });
+        ecs_lock.lock().entity_mut(entity).insert((
+            LoginSendPacketQueue { tx: ecs_packets_tx },
+            login::IgnoreQueryIds::default(),
+        ));
 
         // login
         conn.write(
@@ -343,6 +346,7 @@ impl Client {
                 packet = conn.read() => packet?,
                 Some(packet) = ecs_packets_rx.recv() => {
                     // write this packet to the server
+                    println!("wrote packet to server {:?}", packet);
                     conn.write(packet).await?;
                     continue;
                 }
@@ -425,18 +429,18 @@ impl Client {
                 }
                 ClientboundLoginPacket::CustomQuery(p) => {
                     debug!("Got custom query {:?}", p);
-                    conn.write(
-                        ServerboundCustomQueryAnswerPacket {
-                            transaction_id: p.transaction_id,
-                            data: None,
-                        }
-                        .get(),
-                    )
-                    .await?;
+                    // replying to custom query is done in
+                    // packet_handling::login::process_packet_events
                 }
             }
         };
 
+        ecs_lock
+            .lock()
+            .entity_mut(entity)
+            .remove::<login::IgnoreQueryIds>()
+            .remove::<LoginSendPacketQueue>();
+
         Ok((conn, profile))
     }
 
@@ -628,7 +632,6 @@ pub struct LocalPlayerBundle {
     pub local_player_events: LocalPlayerEvents,
     pub game_profile: GameProfileComponent,
     pub client_information: ClientInformation,
-    pub account: Account,
 }
 
 /// A bundle for the components that are present on a local player that is
diff --git a/azalea-client/src/packet_handling/login.rs b/azalea-client/src/packet_handling/login.rs
index 472ac850..7d71b440 100644
--- a/azalea-client/src/packet_handling/login.rs
+++ b/azalea-client/src/packet_handling/login.rs
@@ -1,16 +1,24 @@
 // login packets aren't actually handled here because compression/encryption
 // would make packet handling a lot messier
 
-use std::sync::Arc;
+use std::{collections::HashSet, sync::Arc};
 
-use azalea_protocol::packets::login::{ClientboundLoginPacket, ServerboundLoginPacket};
-use bevy_ecs::prelude::*;
+use azalea_protocol::packets::login::{
+    serverbound_custom_query_answer_packet::ServerboundCustomQueryAnswerPacket,
+    ClientboundLoginPacket, ServerboundLoginPacket,
+};
+use bevy_ecs::{prelude::*, system::SystemState};
+use derive_more::{Deref, DerefMut};
 use tokio::sync::mpsc;
 use tracing::error;
 
 // this struct is defined here anyways though so it's consistent with the other
 // ones
 
+/// An event that's sent when we receive a login packet from the server. Note
+/// that if you want to handle this in a system, you must add
+/// `.before(azalea::packet_handling::login::process_packet_events)` to it
+/// because that system clears the events.
 #[derive(Event, Debug, Clone)]
 pub struct LoginPacketEvent {
     /// The client entity that received the packet.
@@ -43,3 +51,51 @@ pub fn handle_send_packet_event(
         }
     }
 }
+
+/// 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 fn process_packet_events(ecs: &mut World) {
+    let mut events_owned = Vec::new();
+    let mut system_state: SystemState<ResMut<Events<LoginPacketEvent>>> = SystemState::new(ecs);
+    let mut events = system_state.get_mut(ecs);
+    for LoginPacketEvent {
+        entity: player_entity,
+        packet,
+    } in events.drain()
+    {
+        // we do this so `ecs` isn't borrowed for the whole loop
+        events_owned.push((player_entity, packet));
+    }
+    for (player_entity, packet) in events_owned {
+        #[allow(clippy::single_match)]
+        match packet.as_ref() {
+            ClientboundLoginPacket::CustomQuery(p) => {
+                let mut system_state: SystemState<(
+                    EventWriter<SendLoginPacketEvent>,
+                    Query<&IgnoreQueryIds>,
+                )> = SystemState::new(ecs);
+                let (mut send_packet_events, query) = system_state.get_mut(ecs);
+
+                let ignore_query_ids = query.get(player_entity).ok().map(|x| x.0.clone());
+                if let Some(ignore_query_ids) = ignore_query_ids {
+                    if ignore_query_ids.contains(&p.transaction_id) {
+                        continue;
+                    }
+                }
+
+                send_packet_events.send(SendLoginPacketEvent {
+                    entity: player_entity,
+                    packet: ServerboundCustomQueryAnswerPacket {
+                        transaction_id: p.transaction_id,
+                        data: None,
+                    }
+                    .get(),
+                });
+            }
+            _ => {}
+        }
+    }
+}
diff --git a/azalea-client/src/packet_handling/mod.rs b/azalea-client/src/packet_handling/mod.rs
index 7f797a5c..9823035d 100644
--- a/azalea-client/src/packet_handling/mod.rs
+++ b/azalea-client/src/packet_handling/mod.rs
@@ -46,6 +46,7 @@ impl Plugin for PacketHandlerPlugin {
                     .before(EntityUpdateSet::Deindex),
                 configuration::process_packet_events,
                 login::handle_send_packet_event,
+                login::process_packet_events,
             ),
         )
         .add_systems(Update, death_event_on_0_health.before(death_listener))