feat(lua): add timeout function

src/lua/mod.rs

@@ -7,11 +7,12 @@ pub mod logging;
 pub mod nochatreports;
 pub mod player;
 pub mod system;
+pub mod thread;
 pub mod vec3;
 
 use crate::{ListenerMap, build_info::built};
 use mlua::{Lua, Table};
-use std::{io, time::Duration};
+use std::io;
 
 #[derive(Debug)]
 #[allow(dead_code)]
@@ -33,19 +34,12 @@ pub fn register_globals(
     globals.set("GIT_COMMIT_HASH", built::GIT_COMMIT_HASH)?;
     globals.set("GIT_COMMIT_HASH_SHORT", built::GIT_COMMIT_HASH_SHORT)?;
 
-    globals.set(
-        "sleep",
-        lua.create_async_function(async |_, duration: u64| {
-            tokio::time::sleep(Duration::from_millis(duration)).await;
-            Ok(())
-        })?,
-    )?;
-
     block::register_globals(lua, globals)?;
     events::register_globals(lua, globals, event_listeners)?;
     logging::register_globals(lua, globals)?;
     nochatreports::register_globals(lua, globals)?;
-    system::register_globals(lua, globals)
+    system::register_globals(lua, globals)?;
+    thread::register_globals(lua, globals)
 }
 
 pub fn reload(lua: &Lua, sender: Option<String>) -> Result<(), Error> {

src/lua/thread.rs

@@ -0,0 +1,27 @@
+use mlua::{Error, Function, Lua, Result, Table};
+use std::time::Duration;
+use tokio::time::{sleep, timeout};
+
+pub fn register_globals(lua: &Lua, globals: &Table) -> Result<()> {
+    globals.set(
+        "sleep",
+        lua.create_async_function(async |_, duration: u64| {
+            sleep(Duration::from_millis(duration)).await;
+            Ok(())
+        })?,
+    )?;
+
+    globals.set(
+        "timeout",
+        lua.create_async_function(async |_, (duration, function): (u64, Function)| {
+            timeout(
+                Duration::from_millis(duration),
+                function.call_async::<()>(()),
+            )
+            .await
+            .map_err(Error::external)
+        })?,
+    )?;
+
+    Ok(())
+}

src/main.rs

@@ -14,6 +14,7 @@ use arguments::Arguments;
 use azalea::{
     DefaultBotPlugins, DefaultPlugins, brigadier::prelude::CommandDispatcher, prelude::*,
 };
+use azalea_hax::HaxPlugin;
 use bevy_app::PluginGroup;
 use bevy_log::{
     LogPlugin,
@@ -113,8 +114,14 @@ async fn main() -> anyhow::Result<()> {
             },
         )),
     };
+    let account = if username.contains('@') {
+        Account::microsoft(&username).await?
+    } else {
+        Account::offline(&username)
+    };
     let Err(error) = ClientBuilder::new_without_plugins()
         .add_plugins(DefaultBotPlugins)
+        .add_plugins(HaxPlugin)
         .add_plugins(default_plugins)
         .add_plugins(record_plugin)
         .set_handler(events::handle_event)
@@ -123,14 +130,7 @@ async fn main() -> anyhow::Result<()> {
             event_listeners,
             commands: Arc::new(commands),
         })
-        .start(
-            if username.contains('@') {
-                Account::microsoft(&username).await?
-            } else {
-                Account::offline(&username)
-            },
-            server,
-        )
+        .start(account, server)
         .await;
     eprintln!("{error}");
 

src/replay/recorder.rs

@@ -68,11 +68,12 @@ impl Recorder {
     }
 
     pub fn save_raw_packet(&mut self, raw_packet: &[u8]) -> Result<()> {
-        let mut data = Vec::with_capacity(raw_packet.len() + 8);
-        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)?;
+        self.zip_writer.write_all(
+            &TryInto::<u32>::try_into(self.start.elapsed().as_millis())?.to_be_bytes(),
+        )?;
+        self.zip_writer
+            .write_all(&TryInto::<u32>::try_into(raw_packet.len())?.to_be_bytes())?;
+        self.zip_writer.write_all(raw_packet)?;
         Ok(())
     }