Tweaked macros to insert runtime checks, rather than compile-time.

2 files changed, 57 insertions, 13 deletions

diff --git a/nomsu.lua b/nomsu.lua
index bd04c46..2520f27 100644
--- a/nomsu.lua
+++ b/nomsu.lua
@@ -259,8 +259,8 @@ do
       if def.is_macro and self.callstack[#self.callstack] ~= "#macro" then
         self:error("Attempt to call macro at runtime: " .. tostring(stub) .. "\nThis can be caused by using a macro in a function that is defined before the macro.")
       end
-      if not (self:check_permission(def)) then
-        self:error("You do not have the authority to call: " .. tostring(stub))
+      if not (def.is_macro) then
+        self:assert_permission(stub)
       end
       local thunk, arg_names
       thunk, arg_names = def.thunk, def.arg_names
@@ -283,10 +283,7 @@ do
       remove(self.callstack)
       return unpack(rets)
     end,
-    run_macro = function(self, tree, kind)
-      if kind == nil then
-        kind = "Expression"
-      end
+    run_macro = function(self, tree)
       local stub, arg_names, args = self:get_stub(tree)
       if self.debug then
         self:write(tostring(colored.bright("RUNNING MACRO")) .. " " .. tostring(colored.underscore(colored.magenta(stub))) .. " ")
@@ -297,6 +294,24 @@ do
       remove(self.callstack)
       return expr, statement
     end,
+    assert_permission = function(self, stub)
+      local fn_def = self.defs[stub]
+      if not (fn_def) then
+        self:error("Undefined function: " .. tostring(fn_name))
+      end
+      local whiteset = fn_def.whiteset
+      if whiteset == nil then
+        return true
+      end
+      local _list_0 = self.callstack
+      for _index_0 = 1, #_list_0 do
+        local caller = _list_0[_index_0]
+        if whiteset[caller] then
+          return true
+        end
+      end
+      return self:error("You do not have the authority to call: " .. tostring(stub))
+    end,
     check_permission = function(self, fn_def)
       if getmetatable(fn_def) ~= functiondef_mt then
         local fn_name = fn_def
@@ -452,8 +467,18 @@ do
                     end)]]):format(concat(lua_bits, "\n"))
       elseif "FunctionCall" == _exp_0 then
         local stub = self:get_stub(tree)
-        if self.defs[stub] and self.defs[stub].is_macro then
-          return self:run_macro(tree, "Expression")
+        local def = self.defs[stub]
+        if def and def.is_macro then
+          local expr, statement = self:run_macro(tree)
+          if def.whiteset then
+            if expr then
+              expr = "(nomsu:assert_permission(" .. tostring(repr(stub)) .. ") and " .. tostring(expr) .. ")"
+            end
+            if statement then
+              statement = "nomsu:assert_permission(" .. tostring(repr(stub)) .. ");\n" .. statement
+            end
+          end
+          return expr, statement
         end
         local args = {
           repr(stub)
diff --git a/nomsu.moon b/nomsu.moon
index 7e45d0d..76818d3 100755
--- a/nomsu.moon
+++ b/nomsu.moon
@@ -228,8 +228,8 @@ class NomsuCompiler
         -- I use a hash sign in "#macro" so it's guaranteed to not be a valid function name
         if def.is_macro and @callstack[#@callstack] != "#macro"
             @error "Attempt to call macro at runtime: #{stub}\nThis can be caused by using a macro in a function that is defined before the macro."
-        unless @check_permission(def)
-            @error "You do not have the authority to call: #{stub}"
+        unless def.is_macro
+            @assert_permission(stub)
         {:thunk, :arg_names} = def
         args = {name, select(i,...) for i,name in ipairs(arg_names)}
         if @debug
@@ -241,7 +241,7 @@ class NomsuCompiler
         remove @callstack
         return unpack(rets)
     
-    run_macro: (tree, kind="Expression")=>
+    run_macro: (tree)=>
         stub,arg_names,args = @get_stub tree
         if @debug
             @write "#{colored.bright "RUNNING MACRO"} #{colored.underscore colored.magenta(stub)} "
@@ -251,6 +251,18 @@ class NomsuCompiler
         remove @callstack
         return expr, statement
 
+    assert_permission: (stub)=>
+        fn_def = @defs[stub]
+        unless fn_def
+            @error "Undefined function: #{fn_name}"
+        whiteset = fn_def.whiteset
+        if whiteset == nil then return true
+        -- TODO: maybe optimize this by making the callstack a Counter and using a 
+        --    move-to-front optimization on the whitelist to check most likely candidates sooner
+        for caller in *@callstack
+            if whiteset[caller] then return true
+        @error "You do not have the authority to call: #{stub}"
+
     check_permission: (fn_def)=>
         if getmetatable(fn_def) != functiondef_mt
             fn_name = fn_def
@@ -371,8 +383,15 @@ class NomsuCompiler
 
             when "FunctionCall"
                 stub = @get_stub(tree)
-                if @defs[stub] and @defs[stub].is_macro
-                    return @run_macro(tree, "Expression")
+                def = @defs[stub]
+                if def and def.is_macro
+                    expr, statement = @run_macro(tree)
+                    if def.whiteset
+                        if expr
+                            expr = "(nomsu:assert_permission(#{repr stub}) and #{expr})"
+                        if statement
+                            statement = "nomsu:assert_permission(#{repr stub});\n"..statement
+                    return expr, statement
                 args = {repr(stub)}
                 for arg in *tree.value
                     if arg.type == 'Word' then continue