let plugins override query responses

7 files changed, 92 insertions, 20 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 000da56b..1cb70cc2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -157,6 +157,7 @@ dependencies = [
  "azalea-auth",
  "azalea-block",
  "azalea-brigadier",
+ "azalea-buf",
  "azalea-chat",
  "azalea-client",
  "azalea-core",
diff --git a/azalea-auth/src/sessionserver.rs b/azalea-auth/src/sessionserver.rs
index e6469cef..9c73fbf0 100755
--- a/azalea-auth/src/sessionserver.rs
+++ b/azalea-auth/src/sessionserver.rs
@@ -69,6 +69,15 @@ pub async fn join(
         private_key,
     ));
 
+    join_with_server_id_hash(&client, access_token, uuid, &server_hash).await
+}
+
+pub async fn join_with_server_id_hash(
+    client: &reqwest::Client,
+    access_token: &str,
+    uuid: &Uuid,
+    server_hash: &str,
+) -> Result<(), ClientSessionServerError> {
     let mut encode_buffer = Uuid::encode_buffer();
     let undashed_uuid = uuid.simple().encode_lower(&mut encode_buffer);
 
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))
diff --git a/azalea/Cargo.toml b/azalea/Cargo.toml
index dad1043b..0c02207e 100644
--- a/azalea/Cargo.toml
+++ b/azalea/Cargo.toml
@@ -25,6 +25,7 @@ azalea-registry = { version = "0.8.0", path = "../azalea-registry" }
 azalea-world = { version = "0.8.0", path = "../azalea-world" }
 azalea-auth = { version = "0.8.0", path = "../azalea-auth" }
 azalea-brigadier = { version = "0.8.0", path = "../azalea-brigadier" }
+azalea-buf = { version = "0.8.0", path = "../azalea-buf" }
 bevy_app = "0.12.0"
 bevy_ecs = "0.12.0"
 bevy_tasks = { version = "0.12.0", features = ["multi-threaded"] }
diff --git a/azalea/src/lib.rs b/azalea/src/lib.rs
index 8f86f97f..18a3c314 100644
--- a/azalea/src/lib.rs
+++ b/azalea/src/lib.rs
@@ -18,6 +18,7 @@ use app::{App, Plugins};
 pub use azalea_auth as auth;
 pub use azalea_block as blocks;
 pub use azalea_brigadier as brigadier;
+pub use azalea_buf as buf;
 pub use azalea_chat::FormattedText;
 pub use azalea_client::*;
 pub use azalea_core as core;