implement cookie packets to fix betteranticheat support

5 files changed, 158 insertions, 29 deletions

diff --git a/azalea-client/src/plugins/cookies.rs b/azalea-client/src/plugins/cookies.rs
new file mode 100644
index 00000000..adb3f285
--- /dev/null
+++ b/azalea-client/src/plugins/cookies.rs
@@ -0,0 +1,112 @@
+//! Arbitrary data sent by the server that gets temporarily stored on the
+//! client.
+
+use std::collections::HashMap;
+
+use azalea_protocol::packets::{
+    config,
+    game::{self},
+    login,
+};
+use azalea_registry::identifier::Identifier;
+use bevy_app::{App, Plugin};
+use bevy_ecs::{
+    component::Component,
+    entity::Entity,
+    event::EntityEvent,
+    observer::On,
+    system::{Commands, Query},
+};
+use tracing::warn;
+
+use crate::{
+    InConfigState, InGameState,
+    packet::{
+        config::SendConfigPacketEvent,
+        game::SendGamePacketEvent,
+        login::{InLoginState, SendLoginPacketEvent},
+    },
+};
+
+pub struct CookiesPlugin;
+impl Plugin for CookiesPlugin {
+    fn build(&self, app: &mut App) {
+        app.add_observer(handle_request_cookie)
+            .add_observer(handle_store_cookie);
+    }
+}
+
+/// A component that holds arbitrary data sent by the server, that our client
+/// temporarily stores and persists across transfers.
+#[derive(Component, Default)]
+pub struct ServerCookies {
+    pub map: HashMap<Identifier, Vec<u8>>,
+}
+
+#[derive(EntityEvent)]
+pub struct RequestCookieEvent {
+    pub entity: Entity,
+    pub key: Identifier,
+}
+#[derive(EntityEvent)]
+pub struct StoreCookieEvent {
+    pub entity: Entity,
+    pub key: Identifier,
+    pub payload: Vec<u8>,
+}
+
+pub fn handle_request_cookie(
+    request_cookie: On<RequestCookieEvent>,
+    mut commands: Commands,
+    query: Query<(
+        Option<&ServerCookies>,
+        Option<&InGameState>,
+        Option<&InConfigState>,
+        Option<&InLoginState>,
+    )>,
+) {
+    let Ok((server_cookies, in_game_state, in_config_state, in_login_state)) =
+        query.get(request_cookie.entity)
+    else {
+        return;
+    };
+
+    let key = request_cookie.key.clone();
+    let payload = server_cookies.and_then(|c| c.map.get(&key)).cloned();
+
+    if in_game_state.is_some() {
+        commands.trigger(SendGamePacketEvent::new(
+            request_cookie.entity,
+            game::ServerboundCookieResponse { key, payload },
+        ));
+    } else if in_config_state.is_some() {
+        commands.trigger(SendConfigPacketEvent::new(
+            request_cookie.entity,
+            config::ServerboundCookieResponse { key, payload },
+        ));
+    } else if in_login_state.is_some() {
+        commands.trigger(SendLoginPacketEvent::new(
+            request_cookie.entity,
+            login::ServerboundCookieResponse { key, payload },
+        ));
+    } else {
+        warn!("got RequestCookieEvent while in an unknown state")
+    }
+}
+pub fn handle_store_cookie(
+    store_cookie: On<StoreCookieEvent>,
+    mut commands: Commands,
+    mut query: Query<&mut ServerCookies>,
+) {
+    if let Ok(mut server_cookies) = query.get_mut(store_cookie.entity) {
+        server_cookies
+            .map
+            .insert(store_cookie.key.clone(), store_cookie.payload.clone());
+    } else {
+        commands.entity(store_cookie.entity).insert(ServerCookies {
+            map: [(store_cookie.key.clone(), store_cookie.payload.clone())]
+                .into_iter()
+                .collect(),
+        });
+    }
+}
diff --git a/azalea-client/src/plugins/mod.rs b/azalea-client/src/plugins/mod.rs
index 54f577e3..2ccbb3cc 100644
--- a/azalea-client/src/plugins/mod.rs
+++ b/azalea-client/src/plugins/mod.rs
@@ -10,6 +10,7 @@ pub mod chat_signing;
 pub mod chunks;
 pub mod client_information;
 pub mod connection;
+pub mod cookies;
 pub mod disconnect;
 pub mod events;
 pub mod interact;
@@ -63,7 +64,8 @@ impl PluginGroup for DefaultPlugins {
             .add(connection::ConnectionPlugin)
             .add(login::LoginPlugin)
             .add(join::JoinPlugin)
-            .add(auto_reconnect::AutoReconnectPlugin);
+            .add(auto_reconnect::AutoReconnectPlugin)
+            .add(cookies::CookiesPlugin);
         #[cfg(feature = "online-mode")]
         {
             group = group.add(chat_signing::ChatSigningPlugin);
diff --git a/azalea-client/src/plugins/packet/config/mod.rs b/azalea-client/src/plugins/packet/config/mod.rs
index 629f9f71..27dae837 100644
--- a/azalea-client/src/plugins/packet/config/mod.rs
+++ b/azalea-client/src/plugins/packet/config/mod.rs
@@ -15,6 +15,7 @@ use super::{as_system, declare_packet_handlers};
 use crate::{
     client::InConfigState,
     connection::RawConnection,
+    cookies::{RequestCookieEvent, StoreCookieEvent},
     disconnect::DisconnectEvent,
     local_player::InstanceHolder,
     packet::game::{KeepAliveEvent, ResourcePackEvent},
@@ -179,16 +180,21 @@ impl ConfigPacketHandler<'_> {
 
     pub fn cookie_request(&mut self, p: &ClientboundCookieRequest) {
         debug!("Got cookie request packet {p:?}");
-
         as_system::<Commands>(self.ecs, |mut commands| {
-            commands.trigger(SendConfigPacketEvent::new(
-                self.player,
-                ServerboundCookieResponse {
-                    key: p.key.clone(),
-                    // cookies aren't implemented
-                    payload: None,
-                },
-            ));
+            commands.trigger(RequestCookieEvent {
+                entity: self.player,
+                key: p.key.clone(),
+            });
+        });
+    }
+    pub fn store_cookie(&mut self, p: &ClientboundStoreCookie) {
+        debug!("Got store cookie packet {p:?}");
+        as_system::<Commands>(self.ecs, |mut commands| {
+            commands.trigger(StoreCookieEvent {
+                entity: self.player,
+                key: p.key.clone(),
+                payload: p.payload.clone(),
+            });
         });
     }
 
@@ -196,10 +202,6 @@ impl ConfigPacketHandler<'_> {
         debug!("Got reset chat packet {p:?}");
     }
 
-    pub fn store_cookie(&mut self, p: &ClientboundStoreCookie) {
-        debug!("Got store cookie packet {p:?}");
-    }
-
     pub fn transfer(&mut self, p: &ClientboundTransfer) {
         debug!("Got transfer packet {p:?}");
     }
diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index 89b33274..f39a9f8f 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -29,6 +29,7 @@ use crate::{
     chat::{ChatPacket, ChatReceivedEvent},
     chunks,
     connection::RawConnection,
+    cookies::{RequestCookieEvent, StoreCookieEvent},
     disconnect::DisconnectEvent,
     interact::BlockStatePredictionHandler,
     inventory::{ClientsideCloseContainerEvent, MenuOpenedEvent, SetContainerContentEvent},
@@ -1599,10 +1600,27 @@ impl GamePacketHandler<'_> {
     pub fn ticking_state(&mut self, _p: &ClientboundTickingState) {}
     pub fn ticking_step(&mut self, _p: &ClientboundTickingStep) {}
     pub fn reset_score(&mut self, _p: &ClientboundResetScore) {}
-    pub fn cookie_request(&mut self, _p: &ClientboundCookieRequest) {}
+    pub fn cookie_request(&mut self, p: &ClientboundCookieRequest) {
+        debug!("Got cookie request packet {p:?}");
+        as_system::<Commands>(self.ecs, |mut commands| {
+            commands.trigger(RequestCookieEvent {
+                entity: self.player,
+                key: p.key.clone(),
+            });
+        });
+    }
+    pub fn store_cookie(&mut self, p: &ClientboundStoreCookie) {
+        debug!("Got store cookie packet {p:?}");
+        as_system::<Commands>(self.ecs, |mut commands| {
+            commands.trigger(StoreCookieEvent {
+                entity: self.player,
+                key: p.key.clone(),
+                payload: p.payload.clone(),
+            });
+        });
+    }
     pub fn debug_sample(&mut self, _p: &ClientboundDebugSample) {}
     pub fn pong_response(&mut self, _p: &ClientboundPongResponse) {}
-    pub fn store_cookie(&mut self, _p: &ClientboundStoreCookie) {}
     pub fn transfer(&mut self, _p: &ClientboundTransfer) {}
     pub fn move_minecart_along_track(&mut self, _p: &ClientboundMoveMinecartAlongTrack) {}
     pub fn set_held_slot(&mut self, p: &ClientboundSetHeldSlot) {
diff --git a/azalea-client/src/plugins/packet/login/mod.rs b/azalea-client/src/plugins/packet/login/mod.rs
index 8fe1e22a..de44309c 100644
--- a/azalea-client/src/plugins/packet/login/mod.rs
+++ b/azalea-client/src/plugins/packet/login/mod.rs
@@ -8,7 +8,7 @@ use azalea_protocol::packets::{
     login::{
         ClientboundCookieRequest, ClientboundCustomQuery, ClientboundHello,
         ClientboundLoginCompression, ClientboundLoginDisconnect, ClientboundLoginFinished,
-        ClientboundLoginPacket, ServerboundCookieResponse, ServerboundLoginAcknowledged,
+        ClientboundLoginPacket, ServerboundLoginAcknowledged,
     },
 };
 use bevy_ecs::prelude::*;
@@ -17,8 +17,8 @@ use tracing::{debug, error};
 
 use super::as_system;
 use crate::{
-    Account, InConfigState, connection::RawConnection, disconnect::DisconnectEvent,
-    packet::declare_packet_handlers, player::GameProfileComponent,
+    Account, InConfigState, connection::RawConnection, cookies::RequestCookieEvent,
+    disconnect::DisconnectEvent, packet::declare_packet_handlers, player::GameProfileComponent,
 };
 
 pub fn process_packet(ecs: &mut World, player: Entity, packet: &ClientboundLoginPacket) {
@@ -127,17 +127,12 @@ impl LoginPacketHandler<'_> {
         });
     }
     pub fn cookie_request(&mut self, p: &ClientboundCookieRequest) {
-        debug!("Got cookie request {p:?}");
-
+        debug!("Got cookie request packet {p:?}");
         as_system::<Commands>(self.ecs, |mut commands| {
-            commands.trigger(SendLoginPacketEvent::new(
-                self.player,
-                ServerboundCookieResponse {
-                    key: p.key.clone(),
-                    // cookies aren't implemented
-                    payload: None,
-                },
-            ));
+            commands.trigger(RequestCookieEvent {
+                entity: self.player,
+                key: p.key.clone(),
+            });
         });
     }
 }