feat: finish replay recording automatically on exit

This partially backs out commit f77aea28d1afbc173a1becaf0ecc7a59568caa4c.

Cargo.lock

@@ -937,8 +937,7 @@ dependencies = [
 [[package]]
 name = "built"
 version = "0.7.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
+source = "git+https://github.com/lukaslueg/built#9468c3a117265cc5976bb82631861aa133f32aec"
-checksum = "56ed6191a7e78c36abdb16ab65341eefd73d64d303fffccdbb00d51e4205967b"
 dependencies = [
  "git2",
 ]
@@ -1406,6 +1405,7 @@ dependencies = [
  "built",
  "clap",
  "console-subscriber",
+ "ctrlc",
  "futures",
  "futures-locks",
  "http-body-util",

Cargo.toml

@@ -16,7 +16,7 @@ lto = true
 strip = true
 
 [build-dependencies]
-built = { version = "0", features = ["git2"] }
+built = { git = "https://github.com/lukaslueg/built", features = ["git2"] }
 
 [dependencies]
 anyhow = "1"
@@ -26,6 +26,7 @@ bevy_ecs = "0"
 bevy_log = "0"
 clap = { version = "4", features = ["derive", "string"] }
 console-subscriber = { version = "0", optional = true }
+ctrlc = { version = "3", features = ["termination"] }
 futures = "0"
 futures-locks = "0"
 http-body-util = "0"
@@ -39,10 +40,7 @@ parking_lot = "0"
 serde_json = "1"
 smallvec = { version = "1", features = ["write"] }
 tokio = { version = "1", features = ["macros"] }
-zip = { version = "2", default-features = false, features = [
+zip = { version = "2", default-features = false, features = ["flate2"] }
-  "deflate-flate2",
-  "flate2",
-] }
 
 [features]
 console-subscriber = ["dep:console-subscriber"]

src/events.rs

@@ -16,7 +16,7 @@ use hyper_util::rt::TokioIo;
 use log::{debug, error, info, trace};
 use mlua::{Error, Function, IntoLuaMulti, Table};
 use ncr::utils::trim_header;
-use std::net::SocketAddr;
+use std::{net::SocketAddr, process::exit};
 use tokio::net::TcpListener;
 
 #[allow(clippy::too_many_lines)]
@@ -177,26 +177,17 @@ pub async fn handle_event(client: Client, event: Event, state: State) -> Result<
         Event::Init => {
             debug!("received initialize event");
 
-            let globals = state.lua.globals();
             let ecs = client.ecs.clone();
-            globals.set(
+            ctrlc::set_handler(move || {
-                "finish_replay_recording",
+                debug!("finishing replay recording");
-                state.lua.create_function_mut(move |_, (): ()| {
                 ecs.lock()
                     .remove_resource::<Recorder>()
-                        .context("recording not active")
+                    .map(Recorder::finish);
-                        .map_err(Error::external)?
+                exit(0);
-                        .finish()
+            })?;
-                        .map_err(Error::external)
+
-                })?,
+            let globals = state.lua.globals();
-            )?;
+            lua_init(client, &state, &globals).await?;
-            globals.set(
-                "client",
-                client::Client {
-                    inner: Some(client),
-                },
-            )?;
-            call_listeners(&state, "init", ()).await;
 
             let Some(address): Option<SocketAddr> = globals
                 .get::<String>("HttpAddress")
@@ -237,6 +228,29 @@ pub async fn handle_event(client: Client, event: Event, state: State) -> Result<
     Ok(())
 }
 
+async fn lua_init(client: Client, state: &State, globals: &Table) -> Result<()> {
+    let ecs = client.ecs.clone();
+    globals.set(
+        "finish_replay_recording",
+        state.lua.create_function_mut(move |_, (): ()| {
+            ecs.lock()
+                .remove_resource::<Recorder>()
+                .context("recording not active")
+                .map_err(Error::external)?
+                .finish()
+                .map_err(Error::external)
+        })?,
+    )?;
+    globals.set(
+        "client",
+        client::Client {
+            inner: Some(client),
+        },
+    )?;
+    call_listeners(state, "init", ()).await;
+    Ok(())
+}
+
 async fn call_listeners<T: Clone + IntoLuaMulti + Send + 'static>(
     state: &State,
     event_type: &'static str,

src/main.rs

@@ -60,7 +60,8 @@ async fn main() -> anyhow::Result<()> {
     lua::register_globals(&lua, &globals, event_listeners.clone())?;
     globals.set("SCRIPT_PATH", &*script_path)?;
     lua.load(
-        read_to_string(&script_path).with_context(|| format!("failed to read {script_path:?}"))?,
+        read_to_string(&script_path)
+            .with_context(|| format!("failed to read {}", script_path.display()))?,
     )
     .exec()?;
     if let Some(code) = args.exec {

src/replay/plugin.rs

@@ -30,7 +30,11 @@ impl Plugin for RecordPlugin {
     }
 }
 
-fn record_login_packets(mut recorder: ResMut<Recorder>, mut events: EventReader<LoginPacketEvent>) {
+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.should_ignore_compression
                 && let ClientboundLoginPacket::LoginCompression(_) = *event.packet
@@ -42,21 +46,26 @@ fn record_login_packets(mut recorder: ResMut<Recorder>, mut events: EventReader<
                 error!("failed to record login packet: {error:?}");
             }
         }
+    }
 }
 
 fn record_configuration_packets(
-    mut recorder: ResMut<Recorder>,
+    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(mut recorder: ResMut<Recorder>, query: Query<&RawConnection>) {
+fn record_game_packets(recorder: Option<ResMut<Recorder>>, query: Query<&RawConnection>) {
-    for raw_conn in query.iter() {
+    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) {

src/replay/recorder.rs

@@ -65,13 +65,9 @@ impl Recorder {
         Ok(())
     }
 
-    fn get_timestamp(&self) -> Result<[u8; 4]> {
-        Ok(TryInto::<u32>::try_into(self.start.elapsed().as_millis())?.to_be_bytes())
-    }
-
     pub fn save_raw_packet(&mut self, raw_packet: &[u8]) -> Result<()> {
         let mut data = Vec::with_capacity(raw_packet.len() + 8);
-        data.extend(self.get_timestamp()?);
+        data.extend(&TryInto::<u32>::try_into(self.start.elapsed().as_millis())?.to_be_bytes());
         data.extend(&TryInto::<u32>::try_into(raw_packet.len())?.to_be_bytes());
         data.extend(raw_packet);
         self.zip_writer.write_all(&data)?;
@@ -79,7 +75,7 @@ impl Recorder {
     }
 
     pub fn save_packet<T: ProtocolPacket>(&mut self, packet: &T) -> Result<()> {
-        let mut raw_packet = SmallVec::<[u8; 256]>::new();
+        let mut raw_packet = SmallVec::<[u8; 64]>::new();
         packet.id().azalea_write_var(&mut raw_packet)?;
         packet.write(&mut raw_packet)?;
         self.save_raw_packet(&raw_packet)