PongPlugin

9 files changed, 113 insertions, 17 deletions

diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 0d36fe56..72175dab 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -84,6 +84,7 @@ use crate::{
         login::{self, InLoginState, LoginSendPacketQueue},
     },
     player::retroactively_add_game_profile_component,
+    pong::PongPlugin,
     raw_connection::RawConnection,
     respawn::RespawnPlugin,
     task_pool::TaskPoolPlugin,
@@ -1035,7 +1036,8 @@ impl PluginGroup for DefaultPlugins {
             .add(ChunksPlugin)
             .add(TickEndPlugin)
             .add(BrandPlugin)
-            .add(TickBroadcastPlugin);
+            .add(TickBroadcastPlugin)
+            .add(PongPlugin);
         #[cfg(feature = "log")]
         {
             group = group.add(bevy_log::LogPlugin::default());
diff --git a/azalea-client/src/plugins/events.rs b/azalea-client/src/plugins/events.rs
index 92da08be..ac26aaf5 100644
--- a/azalea-client/src/plugins/events.rs
+++ b/azalea-client/src/plugins/events.rs
@@ -191,10 +191,9 @@ pub fn spawn_listener(
 
 pub fn chat_listener(query: Query<&LocalPlayerEvents>, mut events: EventReader<ChatReceivedEvent>) {
     for event in events.read() {
-        let local_player_events = query
-            .get(event.entity)
-            .expect("Non-local entities shouldn't be able to receive chat events");
-        let _ = local_player_events.send(Event::Chat(event.packet.clone()));
+        if let Ok(local_player_events) = query.get(event.entity) {
+            let _ = local_player_events.send(Event::Chat(event.packet.clone()));
+        }
     }
 }
 
@@ -210,10 +209,9 @@ pub fn packet_listener(
     mut events: EventReader<ReceivePacketEvent>,
 ) {
     for event in events.read() {
-        let local_player_events = query
-            .get(event.entity)
-            .expect("Non-local entities shouldn't be able to receive packet events");
-        let _ = local_player_events.send(Event::Packet(event.packet.clone()));
+        if let Ok(local_player_events) = query.get(event.entity) {
+            let _ = local_player_events.send(Event::Packet(event.packet.clone()));
+        }
     }
 }
 
diff --git a/azalea-client/src/plugins/mod.rs b/azalea-client/src/plugins/mod.rs
index 11794fb3..b5005b22 100644
--- a/azalea-client/src/plugins/mod.rs
+++ b/azalea-client/src/plugins/mod.rs
@@ -9,6 +9,7 @@ pub mod inventory;
 pub mod mining;
 pub mod movement;
 pub mod packet;
+pub mod pong;
 pub mod respawn;
 pub mod task_pool;
 pub mod tick_end;
diff --git a/azalea-client/src/plugins/packet/config/events.rs b/azalea-client/src/plugins/packet/config/events.rs
index 9ed6c097..d0a7f3be 100644
--- a/azalea-client/src/plugins/packet/config/events.rs
+++ b/azalea-client/src/plugins/packet/config/events.rs
@@ -109,3 +109,13 @@ fn packet_interrupts(packet: &ClientboundConfigPacket) -> bool {
             | ClientboundConfigPacket::Transfer(_)
     )
 }
+
+/// A Bevy trigger that's sent when our client receives a [`ClientboundPing`]
+/// packet in the config state.
+///
+/// See [`PingEvent`] for more information.
+///
+/// [`ClientboundPing`]: azalea_protocol::packets::config::ClientboundPing
+/// [`PingEvent`]: crate::packet::game::PingEvent
+#[derive(Event, Debug, Clone)]
+pub struct ConfigPingEvent(pub azalea_protocol::packets::config::ClientboundPing);
diff --git a/azalea-client/src/plugins/packet/config/mod.rs b/azalea-client/src/plugins/packet/config/mod.rs
index c9b84eac..ae601793 100644
--- a/azalea-client/src/plugins/packet/config/mod.rs
+++ b/azalea-client/src/plugins/packet/config/mod.rs
@@ -142,10 +142,8 @@ impl ConfigPacketHandler<'_> {
     pub fn ping(&mut self, p: ClientboundPing) {
         debug!("Got ping packet (in configuration) {p:?}");
 
-        as_system::<Query<&RawConnection>>(self.ecs, |query| {
-            let raw_conn = query.get(self.player).unwrap();
-
-            raw_conn.write_packet(ServerboundPong { id: p.id }).unwrap();
+        as_system::<Commands>(self.ecs, |mut commands| {
+            commands.trigger_targets(ConfigPingEvent(p), self.player);
         });
     }
 
diff --git a/azalea-client/src/plugins/packet/game/events.rs b/azalea-client/src/plugins/packet/game/events.rs
index 8a5aca3c..ad81f9bd 100644
--- a/azalea-client/src/plugins/packet/game/events.rs
+++ b/azalea-client/src/plugins/packet/game/events.rs
@@ -212,3 +212,21 @@ pub struct InstanceLoadedEvent {
     pub name: ResourceLocation,
     pub instance: Weak<RwLock<Instance>>,
 }
+
+/// A Bevy trigger that's sent when our client receives a [`ClientboundPing`]
+/// packet in the game state.
+///
+/// Also see [`ConfigPingEvent`] which is used for the config state.
+///
+/// This is not an event and can't be listened to from a normal system,
+///so  `EventReader<PingEvent>` will not work.
+///
+/// To use it, add your "system" with `add_observer` instead of `add_systems`
+/// and use `Trigger<PingEvent>` instead of `EventReader`.
+///
+/// The client Entity that received the packet will be attached to the trigger.
+///
+/// [`ClientboundPing`]: azalea_protocol::packets::game::ClientboundPing
+/// [`ConfigPingEvent`]: crate::packet::config::ConfigPingEvent
+#[derive(Event, Debug, Clone)]
+pub struct PingEvent(pub azalea_protocol::packets::game::ClientboundPing);
diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index 1e9232bf..cc67bcd7 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -1363,10 +1363,7 @@ impl GamePacketHandler<'_> {
         debug!("Got ping packet {p:?}");
 
         as_system::<Commands>(self.ecs, |mut commands| {
-            commands.trigger(SendPacketEvent::new(
-                self.player,
-                ServerboundPong { id: p.id },
-            ));
+            commands.trigger_targets(PingEvent(p.clone()), self.player);
         });
     }
 
diff --git a/azalea-client/src/plugins/pong.rs b/azalea-client/src/plugins/pong.rs
new file mode 100644
index 00000000..827ddfb1
--- /dev/null
+++ b/azalea-client/src/plugins/pong.rs
@@ -0,0 +1,37 @@
+use bevy_app::{App, Plugin};
+use bevy_ecs::prelude::*;
+
+use super::packet::{
+    config::{ConfigPingEvent, SendConfigPacketEvent},
+    game::PingEvent,
+};
+use crate::packet::game::SendPacketEvent;
+
+/// A plugin that replies to [`ClientboundPing`] packets with
+/// [`ServerboundPong`].
+///
+/// This works in both the `game` and `config` states.
+///
+/// [`ClientboundPing`]: azalea_protocol::packets::game::ClientboundPing
+/// [`ServerboundPong`]: azalea_protocol::packets::game::ServerboundPong
+pub struct PongPlugin;
+impl Plugin for PongPlugin {
+    fn build(&self, app: &mut App) {
+        app.add_observer(reply_to_game_ping)
+            .add_observer(reply_to_config_ping);
+    }
+}
+
+pub fn reply_to_game_ping(trigger: Trigger<PingEvent>, mut commands: Commands) {
+    commands.trigger(SendPacketEvent::new(
+        trigger.entity(),
+        azalea_protocol::packets::game::ServerboundPong { id: trigger.0.id },
+    ));
+}
+
+pub fn reply_to_config_ping(trigger: Trigger<ConfigPingEvent>, mut commands: Commands) {
+    commands.trigger(SendConfigPacketEvent::new(
+        trigger.entity(),
+        azalea_protocol::packets::config::ServerboundPong { id: trigger.0.id },
+    ));
+}
diff --git a/azalea-client/tests/reply_to_ping_with_pong.rs b/azalea-client/tests/reply_to_ping_with_pong.rs
new file mode 100644
index 00000000..4ef5b2cc
--- /dev/null
+++ b/azalea-client/tests/reply_to_ping_with_pong.rs
@@ -0,0 +1,35 @@
+use std::sync::Arc;
+
+use azalea_client::{packet::game::SendPacketEvent, test_simulation::*};
+use azalea_protocol::packets::{
+    ConnectionProtocol,
+    game::{ClientboundPing, ServerboundGamePacket},
+};
+use bevy_ecs::observer::Trigger;
+use bevy_log::tracing_subscriber;
+use parking_lot::Mutex;
+
+#[test]
+fn reply_to_ping_with_pong() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let mut simulation = Simulation::new(ConnectionProtocol::Game);
+    let reply_count = Arc::new(Mutex::new(0));
+    let reply_count_clone = reply_count.clone();
+    simulation
+        .app
+        .add_observer(move |trigger: Trigger<SendPacketEvent>| {
+            if trigger.sent_by == simulation.entity {
+                if let ServerboundGamePacket::Pong(packet) = &trigger.packet {
+                    assert_eq!(packet.id, 123);
+                    *reply_count_clone.lock() += 1;
+                }
+            }
+        });
+
+    simulation.tick();
+    simulation.receive_packet(ClientboundPing { id: 123 });
+    simulation.tick();
+
+    assert_eq!(*reply_count.lock(), 1);
+}