Add possibly unsafe timeout hook to factoids

3 files changed, 30 insertions, 4 deletions

diff --git a/src/plugins/factoids/mod.rs b/src/plugins/factoids/mod.rs
index 252ba83..08e8f12 100644
--- a/src/plugins/factoids/mod.rs
+++ b/src/plugins/factoids/mod.rs
@@ -15,7 +15,7 @@ pub mod database;
 use self::database::{Database, DbResponse};
 
 mod utils;
-use self::utils::download;
+use self::utils::*;
 
 static LUA_SANDBOX: &'static str = include_str!("sandbox.lua");
 
@@ -247,16 +247,17 @@ impl<T: Database> Factoids<T> {
             .map(ToOwned::to_owned)
             .collect::<Vec<String>>();
 
-        let lua = Lua::new();
+        let lua = unsafe { Lua::new_with_debug() };
         let globals = lua.globals();
 
         globals.set("factoid", code)?;
-        globals.set("download", lua.create_function(download))?;
+        globals.set("download", lua.create_function(download)?)?;
+        globals.set("sleep", lua.create_function(sleep)?)?;
         globals.set("args", args)?;
         globals.set("input", command.tokens.join(" "))?;
         globals.set("user", command.source.clone())?;
         globals.set("channel", command.target.clone())?;
-        globals.set("output", lua.create_table())?;
+        globals.set("output", lua.create_table()?)?;
 
         lua.exec::<()>(LUA_SANDBOX, Some(name))?;
         let output: Vec<String> = globals.get::<_, Vec<String>>("output")?;
diff --git a/src/plugins/factoids/sandbox.lua b/src/plugins/factoids/sandbox.lua
index 8ebf518..3fde65e 100644
--- a/src/plugins/factoids/sandbox.lua
+++ b/src/plugins/factoids/sandbox.lua
@@ -24,9 +24,26 @@ local env = { print = send,
               table = table,
               string = string,
               tostring = tostring,
+              tonumber = tonumber,
               math = math }
 
 local f, e = load(factoid, nil, nil, env)
+
+-- Check if the factoid timed out
+function checktime(event, line)
+    if os.time() - time >= timeout then
+        error("Timed out after " .. timeout .. " seconds", 0)
+    else
+        -- Limit the cpu usage of factoids
+        sleep(1)
+    end
+end
+
+-- Add timeout hook
+time = os.time()
+timeout = 30
+debug.sethook(checktime, "l")
+
 if f then
   f()
 else
diff --git a/src/plugins/factoids/utils.rs b/src/plugins/factoids/utils.rs
index ad25eef..fc86fb3 100644
--- a/src/plugins/factoids/utils.rs
+++ b/src/plugins/factoids/utils.rs
@@ -1,5 +1,8 @@
 extern crate reqwest;
 
+use std::thread;
+use std::time::Duration;
+
 use utils;
 use super::rlua::prelude::*;
 
@@ -11,3 +14,8 @@ pub fn download(_: &Lua, url: String) -> Result<String, LuaError> {
         None => Err(RuntimeError(format!("Failed to download {}", url))),
     }
 }
+
+pub fn sleep(_: &Lua, dur: u64) -> Result<(), LuaError> {
+    thread::sleep(Duration::from_millis(dur));
+    Ok(())
+}