From b1d69c89d77fb0b07fc40e4b99ba398cba512ecb Mon Sep 17 00:00:00 2001
From: kikito <kikito@gmail.com>
Date: Tue, 3 Sep 2013 16:41:46 +0200
Subject: [PATCH] cleanup

---
 sandbox.lua           | 30 +++++++++++++++++++++---------
 spec/sandbox_spec.lua | 15 +++++++++++++--
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/sandbox.lua b/sandbox.lua
index 66ff7c4..034f4eb 100644
--- a/sandbox.lua
+++ b/sandbox.lua
@@ -47,15 +47,26 @@ end)
 
 local string_rep = string.rep
 
-local function copy(other)
+local copy -- defined below
+
+local function merge(destination, other)
   if type(other) ~= 'table' then return other end
-  local c = {}
   for k,v in pairs(other) do
-    c[copy(k)] = copy(v)
+    destination[copy(k)] = copy(v)
   end
+  return destination
+end
+
+-- declared above
+copy = function(other)
+  if type(other) ~= 'table' then return other end
+  local c = merge({}, other)
+  local mt = getmetatable(other)
+  if mt then setmetatable(c, copy(mt)) end
   return c
 end
 
+
 local function cleanup()
   debug.sethook()
   string.rep = string_rep
@@ -66,11 +77,12 @@ local function run(f, options)
 
   options = options or {}
 
-  local limit = options.limit or 500000
+  local quota = options.quota or 500000
+  local env   = options.env   or {}
 
-  local env = copy(BASE_ENV)
+  local sandboxed_env = merge(copy(BASE_ENV), env)
 
-  setfenv(f, env)
+  setfenv(f, sandboxed_env)
 
   -- I would love to be able to make step greater than 1
   -- (say, 500000) but any value > 1 seems to choke with a simple while true do end
@@ -80,9 +92,9 @@ local function run(f, options)
   local instructions_count = 0
   local timeout = function(str)
     instructions_count = instructions_count + 1
-    if instructions_count >= limit then
+    if instructions_count >= quota then
       cleanup()
-      error('Quota exceeded: ' .. tostring(instructions_count) .. '/' .. tostring(limit) .. ' instructions')
+      error('Quota exceeded: ' .. tostring(instructions_count) .. '/' .. tostring(quota) .. ' instructions')
     end
   end
   debug.sethook(timeout, "", step)
@@ -91,8 +103,8 @@ local function run(f, options)
   local ok, result = pcall(f)
 
   cleanup()
-  if not ok then error(result) end
 
+  if not ok then error(result) end
   return result
 end
 
diff --git a/spec/sandbox_spec.lua b/spec/sandbox_spec.lua
index 7aa0406..da406d5 100644
--- a/spec/sandbox_spec.lua
+++ b/spec/sandbox_spec.lua
@@ -22,6 +22,7 @@ describe('sandbox', function()
   it('does not allow access to not-safe stuff', function()
     assert.has_error(function() sandbox('return setmetatable({}, {})') end)
     assert.has_error(function() sandbox('return string.rep("hello", 5)') end)
+    assert.has_error(function() sandbox('return _G.string.upper("hello")') end)
   end)
 
   it('does not allow pesky string:rep', function()
@@ -55,13 +56,23 @@ describe('sandbox', function()
       assert.equal('hellohello', string.rep('hello', 2))
     end)
 
-    it('#focus accepts a limit param', function()
+    it('accepts a quota param', function()
       assert.no_has_error(function() sandbox("for i=1,100 do end") end)
-      assert.has_error(function() sandbox("for i=1,100 do end", {limit = 50}) end)
+      assert.has_error(function() sandbox("for i=1,100 do end", {quota = 50}) end)
     end)
 
   end)
 
+  describe('when given an env', function()
+    it('is available on the sandboxed env', function()
+      assert.equal(1, sandbox("return foo", {env = {foo = 1}}))
+    end)
+
+    it('does not hide previous env', function()
+      assert.equal('HELLO', sandbox("return string.upper(foo)", {env = {foo = 'hello'}}))
+    end)
+  end)
+
 
 
 end)