let plugins override query responses

2 files changed, 60 insertions, 3 deletions

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))