3 files changed, 62 insertions, 40 deletions

diff --git a/src/plugins/factoid/mod.rs b/src/plugins/factoid/mod.rs
index 5a484e6..08b96ec 100644
--- a/src/plugins/factoid/mod.rs
+++ b/src/plugins/factoid/mod.rs
@@ -1,11 +1,16 @@
 extern crate rlua;
 
-use self::rlua::prelude::*;
-use antidote::RwLock;
-use irc::client::prelude::*;
 use std::fmt;
 use std::marker::PhantomData;
 use std::str::FromStr;
+use std::sync::Arc;
+use std::thread;
+use std::time::{Duration, Instant};
+
+use self::rlua::prelude::*;
+use self::rlua::HookTriggers;
+use antidote::RwLock;
+use irc::client::prelude::*;
 
 use chrono::NaiveDateTime;
 use time;
@@ -188,7 +193,15 @@ impl<T: Database, C: Client> Factoid<T, C> {
                     match self.run_lua(&name, &content, &command) {
                         Ok(v) => v,
                         Err(e) => match e {
-                            LuaError::CallbackError { cause, .. } => cause.to_string(),
+                            LuaError::CallbackError { cause, .. } => match *cause {
+                                LuaError::MemoryError(_) => {
+                                    String::from("memory error: Factoid used over 1 MiB of ram")
+                                }
+                                _ => cause.to_string(),
+                            },
+                            LuaError::MemoryError(_) => {
+                                String::from("memory error: Factoid used over 1 MiB of ram")
+                            }
                             _ => e.to_string(),
                         },
                     }
@@ -209,13 +222,43 @@ impl<T: Database, C: Client> Factoid<T, C> {
             .map(ToOwned::to_owned)
             .collect::<Vec<String>>();
 
-        let lua = unsafe { Lua::new_with_debug() };
+        let lua = Lua::new();
+        // TODO Is this actually 1 Mib?
+        lua.set_memory_limit(Some(1024 * 1024));
+
+        let start = Instant::now();
+        // Check if the factoid timed out
+        lua.set_hook(
+            HookTriggers {
+                every_line: true,
+                ..Default::default()
+            },
+            move |_, _| {
+                if Instant::now() - start > Duration::from_secs(30) {
+                    return Err(LuaError::ExternalError(Arc::new(
+                        format_err!("Factoid timed out after 30 seconds").compat(),
+                    )));
+                }
+
+                // Limit the cpu usage of factoids
+                thread::sleep(Duration::from_millis(1));
+
+                Ok(())
+            },
+        );
+
         let output = lua.context(|ctx| {
             let globals = ctx.globals();
 
             globals.set("factoid", code)?;
-            globals.set("download", ctx.create_function(|ctx, url| download(&ctx, url))?)?;
-            globals.set("json_decode", ctx.create_function(|ctx, json| json_decode(&ctx, json))?)?;
+            globals.set(
+                "download",
+                ctx.create_function(|ctx, url| download(&ctx, url))?,
+            )?;
+            globals.set(
+                "json_decode",
+                ctx.create_function(|ctx, json| json_decode(&ctx, json))?,
+            )?;
             globals.set("sleep", ctx.create_function(|ctx, ms| sleep(&ctx, ms))?)?;
             globals.set("args", args)?;
             globals.set("input", command.tokens.join(" "))?;
@@ -241,11 +284,13 @@ impl<T: Database, C: FrippyClient> Plugin for Factoid<T, C> {
     type Client = C;
     fn execute(&self, _: &Self::Client, message: &Message) -> ExecutionStatus {
         match message.command {
-            Command::PRIVMSG(_, ref content) => if content.starts_with('!') {
-                ExecutionStatus::RequiresThread
-            } else {
-                ExecutionStatus::Done
-            },
+            Command::PRIVMSG(_, ref content) => {
+                if content.starts_with('!') {
+                    ExecutionStatus::RequiresThread
+                } else {
+                    ExecutionStatus::Done
+                }
+            }
             _ => ExecutionStatus::Done,
         }
     }
diff --git a/src/plugins/factoid/sandbox.lua b/src/plugins/factoid/sandbox.lua
index a927535..108a30d 100644
--- a/src/plugins/factoid/sandbox.lua
+++ b/src/plugins/factoid/sandbox.lua
@@ -90,35 +90,8 @@ end
 sandbox_env.eval = eval
 sandbox_env.sleep = safesleep
 
--- Check if the factoid timed out
-function checktime()
-  if os.time() - time >= timeout then
-    error("Timed out after " .. timeout .. " seconds", 0)
-  else
-    -- Limit the cpu usage of factoids
-    sleep(1)
-  end
-end
-
--- Check if the factoid uses too much memory
-function checkmem()
-  if collectgarbage("count") > maxmem then
-    error("Factoid used over " .. maxmem .. " kbyte of ram")
-  end
-end
-
 local f, e = load(factoid, nil, nil, sandbox_env)
 
--- Add timeout hook
-time = os.time()
--- The timeout is defined in seconds
-timeout = 30
-debug.sethook(checktime, "l")
--- Add memory check hook
--- The max memory is defined in kilobytes
-maxmem = 1000
-debug.sethook(checkmem, "l")
-
 if f then
   f()
 else
diff --git a/src/plugins/factoid/utils.rs b/src/plugins/factoid/utils.rs
index bd1bf43..b764f03 100644
--- a/src/plugins/factoid/utils.rs
+++ b/src/plugins/factoid/utils.rs
@@ -36,7 +36,11 @@ pub fn download(_: &Context, url: String) -> Result<String, LuaError> {
     }
 }
 
-fn convert_value<'l>(lua: &Context<'l>, sval: SerdeValue, max_recurs: usize) -> Result<LuaValue<'l>, LuaError> {
+fn convert_value<'l>(
+    lua: &Context<'l>,
+    sval: SerdeValue,
+    max_recurs: usize,
+) -> Result<LuaValue<'l>, LuaError> {
     if max_recurs == 0 {
         return Err(RuntimeError(String::from(
             "Reached max recursion level - json is nested too deep",