feat: add replaymod compatible recorder

src/events.rs

@@ -2,9 +2,11 @@ use crate::{
     State,
     commands::CommandSource,
     http::serve,
-    lua::{self, direction::Direction, player::Player, vec3::Vec3},
+    lua::{client, direction::Direction, player::Player, vec3::Vec3},
     particle,
+    replay::Recorder,
 };
+use anyhow::Context;
 use azalea::{
     brigadier::exceptions::BuiltInExceptions::DispatcherUnknownCommand, prelude::*,
     protocol::packets::game::ClientboundGamePacket,
@@ -12,7 +14,7 @@ use azalea::{
 use hyper::{server::conn::http1, service::service_fn};
 use hyper_util::rt::TokioIo;
 use log::{debug, error, info, trace};
-use mlua::{Function, IntoLuaMulti, Table};
+use mlua::{Error, Function, IntoLuaMulti, Table};
 use ncr::utils::trim_header;
 use tokio::net::TcpListener;
 
@@ -167,9 +169,23 @@ pub async fn handle_event(client: Client, event: Event, state: State) -> anyhow:
         Event::Init => {
             debug!("received initialize event");
 
-            state.lua.globals().set(
+            let globals = state.lua.globals();
+            let lua_ecs = client.ecs.clone();
+            globals.set(
+                "finish_replay_recording",
+                state.lua.create_function_mut(move |_, (): ()| {
+                    lua_ecs
+                        .lock()
+                        .remove_resource::<Recorder>()
+                        .context("recording not active")
+                        .map_err(Error::external)?
+                        .finish()
+                        .map_err(Error::external)
+                })?,
+            )?;
+            globals.set(
                 "client",
-                lua::client::Client {
+                client::Client {
                     inner: Some(client),
                 },
             )?;

src/main.rs

@@ -7,6 +7,7 @@ mod events;
 mod http;
 mod lua;
 mod particle;
+mod replay;
 
 use arguments::Arguments;
 use azalea::{
@@ -21,7 +22,8 @@ use clap::Parser;
 use commands::{CommandSource, register};
 use futures::lock::Mutex;
 use futures_locks::RwLock;
-use mlua::{Function, Lua};
+use mlua::{Function, Lua, Table};
+use replay::{Recorder, plugin::RecordPlugin};
 use std::{
     collections::HashMap,
     env,
@@ -37,10 +39,10 @@ type ListenerMap = Arc<RwLock<HashMap<String, Vec<(String, Function)>>>>;
 
 #[derive(Default, Clone, Component)]
 pub struct State {
+    http_address: Option<SocketAddr>,
     lua: Arc<Lua>,
     event_listeners: ListenerMap,
     commands: Arc<CommandDispatcher<Mutex<CommandSource>>>,
-    http_address: Option<SocketAddr>,
 }
 
 #[tokio::main]
@@ -91,15 +93,33 @@ async fn main() -> anyhow::Result<()> {
             ..Default::default()
         })
     };
+    let record_plugin = RecordPlugin {
+        recorder: Arc::new(parking_lot::Mutex::new(
+            if let Ok(options) = globals.get::<Table>("ReplayRecordingOptions")
+                && let Ok(path) = options.get::<String>("path")
+            {
+                Some(Recorder::new(
+                    path,
+                    server.clone(),
+                    options
+                        .get::<bool>("ignore_compression")
+                        .unwrap_or_default(),
+                )?)
+            } else {
+                None
+            },
+        )),
+    };
     let Err(error) = ClientBuilder::new_without_plugins()
         .add_plugins(default_plugins)
+        .add_plugins(record_plugin)
         .add_plugins(DefaultBotPlugins)
         .set_handler(events::handle_event)
         .set_state(State {
+            http_address: args.http_address,
             lua: Arc::new(lua),
             event_listeners,
             commands: Arc::new(commands),
-            http_address: args.http_address,
         })
         .start(
             if username.contains('@') {

src/replay/mod.rs

@@ -0,0 +1,88 @@
+pub mod plugin;
+
+use crate::build_info;
+use anyhow::Result;
+use azalea::{
+    buf::AzaleaWriteVar,
+    prelude::Resource,
+    protocol::packets::{PROTOCOL_VERSION, ProtocolPacket, VERSION_NAME},
+};
+use serde_json::json;
+use std::{
+    fs::File,
+    io::Write,
+    time::{SystemTime, UNIX_EPOCH},
+};
+use zip::{ZipWriter, write::SimpleFileOptions};
+
+#[derive(Resource)]
+pub struct Recorder {
+    zip_writer: ZipWriter<File>,
+    start_time: u128,
+    server: String,
+    ignore_compression: bool,
+}
+
+impl Recorder {
+    pub fn new(path: String, server: String, ignore_compression: bool) -> Result<Self> {
+        let mut zip_writer = ZipWriter::new(
+            File::options()
+                .write(true)
+                .create(true)
+                .truncate(true)
+                .open(path)?,
+        );
+        zip_writer.start_file("recording.tmcpr", SimpleFileOptions::default())?;
+        Ok(Self {
+            zip_writer,
+            start_time: SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis(),
+            server,
+            ignore_compression,
+        })
+    }
+
+    pub fn finish(mut self) -> Result<()> {
+        self.zip_writer
+            .start_file("metaData.json", SimpleFileOptions::default())?;
+        self.zip_writer.write_all(
+            json!({
+                "singleplayer": false,
+                "serverName": self.server,
+                "duration": SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() - self.start_time,
+                "date": self.start_time,
+                "mcversion": VERSION_NAME,
+                "fileFormat": "MCPR",
+                "fileFormatVersion": 14,
+                "protocol": PROTOCOL_VERSION,
+                "generator": build_info::version_formatted(),
+            })
+            .to_string()
+            .as_bytes(),
+        )?;
+        self.zip_writer.finish()?;
+
+        Ok(())
+    }
+
+    fn get_timestamp(&self) -> Result<[u8; 4]> {
+        Ok(TryInto::<u32>::try_into(
+            SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() - self.start_time,
+        )?
+        .to_be_bytes())
+    }
+
+    fn save_raw_packet(&mut self, raw_packet: &[u8]) -> Result<()> {
+        let mut data = Vec::from(self.get_timestamp()?);
+        data.extend(TryInto::<u32>::try_into(raw_packet.len())?.to_be_bytes());
+        data.extend(raw_packet);
+        self.zip_writer.write_all(&data)?;
+        Ok(())
+    }
+
+    fn save_packet<T: ProtocolPacket>(&mut self, packet: &T) -> Result<()> {
+        let mut raw_packet = Vec::new();
+        packet.id().azalea_write_var(&mut raw_packet)?;
+        packet.write(&mut raw_packet)?;
+        self.save_raw_packet(&raw_packet)
+    }
+}

src/replay/plugin.rs

@@ -0,0 +1,76 @@
+use super::Recorder;
+use azalea::{
+    ecs::{event::EventReader, system::Query},
+    packet_handling::{
+        configuration::ConfigurationEvent,
+        game::send_packet_events,
+        login::{LoginPacketEvent, process_packet_events},
+    },
+    protocol::packets::login::ClientboundLoginPacket,
+    raw_connection::RawConnection,
+};
+use bevy_app::{First, Plugin};
+use bevy_ecs::{schedule::IntoSystemConfigs, system::ResMut};
+use log::error;
+use parking_lot::Mutex;
+use std::sync::Arc;
+
+pub struct RecordPlugin {
+    pub recorder: Arc<Mutex<Option<Recorder>>>,
+}
+
+impl Plugin for RecordPlugin {
+    fn build(&self, app: &mut bevy_app::App) {
+        if let Some(recorder) = self.recorder.lock().take() {
+            app.insert_resource(recorder);
+        }
+        app.add_systems(First, record_login_packets.before(process_packet_events))
+            .add_systems(First, record_configuration_packets)
+            .add_systems(First, record_game_packets.before(send_packet_events));
+    }
+}
+
+fn record_login_packets(
+    recorder: Option<ResMut<Recorder>>,
+    mut events: EventReader<LoginPacketEvent>,
+) {
+    if let Some(mut recorder) = recorder {
+        for event in events.read() {
+            if recorder.ignore_compression
+                && let ClientboundLoginPacket::LoginCompression(_) = *event.packet
+            {
+                continue;
+            }
+
+            if let Err(error) = recorder.save_packet(event.packet.as_ref()) {
+                error!("failed to record login packet: {error:?}");
+            }
+        }
+    }
+}
+
+fn record_configuration_packets(
+    recorder: Option<ResMut<Recorder>>,
+    mut events: EventReader<ConfigurationEvent>,
+) {
+    if let Some(mut recorder) = recorder {
+        for event in events.read() {
+            if let Err(error) = recorder.save_packet(&event.packet) {
+                error!("failed to record configuration packet: {error:?}");
+            }
+        }
+    }
+}
+
+fn record_game_packets(recorder: Option<ResMut<Recorder>>, query: Query<&RawConnection>) {
+    if let Some(mut recorder) = recorder
+        && let Ok(raw_conn) = query.get_single()
+    {
+        let queue = raw_conn.incoming_packet_queue();
+        for raw_packet in queue.lock().iter() {
+            if let Err(error) = recorder.save_raw_packet(raw_packet) {
+                error!("failed to record game packet: {error:?}");
+            }
+        }
+    }
+}