Signed bytecode loading

Question

So I'm looking into securing Lua's load function for bytecode loading. Currently I have this:

local nativeload = load
load = function(chunk, chunkname, mode, ...)
  if mode == nil then
    mode = "bt"
  elseif not (mode == "b" or mode == "t" or mode == "bt") then
    error("Invalid mode")
  end
  local targetenv = [[snip]]
  if select('#', ...) > 0 then targetenv = ... end
  if mode == "t" then
    return nativeload(chunk, chunkname, mode, targetenv)
  elseif type(chunk) == "string" then
    if chunk:sub(1,4) == "\27Lua" then
      local code = chunk:sub(1,-33)
      if HMAC_SHA256(code, getkey()) == chunk:sub(-32) then
        return nativeload(code, chunkname, mode, targetenv)
      else
        error("Invalid signature")
      end
    else
      return nativeload(chunk, chunkname, mode, targetenv)
    end
  elseif type(chunk) == "function" then
    -- How do I do this?!
  end
end

And, while text mode and string chunks are fairly trivial to handle, I have no idea how to handle function chunks.

Do I just collect everything into a string somehow, then do the HMAC stuff and call nativeload with this string? But then I lose being able to load() large files (e.g. 2GB files) without the program crashing (load with function compiles the file in 8kbyte increments, and when the file is mostly empty lines, that means it only needs a few kbytes during compilation - while loading the whole file into a string would clearly use 2GB of RAM).

How would I do this?

Answer 1

local function extract_code(data)
   local code = data:sub(1,-33)
   assert(HMAC_SHA256(code, getkey()) == data:sub(-32), "Invalid signature")
   return code
end

local nativeload = load
load = function(chunk, chunkname, mode, ...)
   local targetenv = [[snip]]
   if select('#', ...) ~= 0 then targetenv = ... end
   local new_chunk
   if type(chunk) == "string" then
      new_chunk = chunk:match"^\27Lua" and extract_code(chunk) or chunk
   elseif type(chunk) == "function" then
      local buffer = ""
      repeat
         local next_part = chunk() or ""
         buffer = buffer..next_part
      until next_part == "" or #buffer >= 4
      if buffer:match"^\27Lua" then -- Bytecode can't be very large, collect it in a string
         local t = {buffer}
         while t[#t] ~= "" do t[#t+1] = chunk() or "" end
         new_chunk = extract_code(table.concat(t))
      else                          -- Source code can be very large, use a function
         local function f()
            f = chunk
            return buffer
         end
         function new_chunk()
            return f()
         end
      end
   end
   return nativeload(new_chunk, chunkname, mode, targetenv)
end