287 lines
7.1 KiB
Go
287 lines
7.1 KiB
Go
package runner
|
|
|
|
import (
|
|
luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
|
|
)
|
|
|
|
// Sandbox manages a sandboxed Lua environment
|
|
type Sandbox struct {
|
|
modules map[string]any // Custom modules for environment
|
|
initialized bool // Whether base environment is initialized
|
|
}
|
|
|
|
// NewSandbox creates a new sandbox
|
|
func NewSandbox() *Sandbox {
|
|
s := &Sandbox{
|
|
modules: make(map[string]any),
|
|
initialized: false,
|
|
}
|
|
|
|
return s
|
|
}
|
|
|
|
// AddModule adds a module to the sandbox environment
|
|
func (s *Sandbox) AddModule(name string, module any) {
|
|
s.modules[name] = module
|
|
}
|
|
|
|
// Setup initializes the sandbox in a Lua state
|
|
func (s *Sandbox) Setup(state *luajit.State) error {
|
|
// Register modules
|
|
if err := s.registerModules(state); err != nil {
|
|
return err
|
|
}
|
|
|
|
// Create high-performance persistent environment
|
|
return state.DoString(`
|
|
-- Global shared environment (created once)
|
|
__env_system = __env_system or {
|
|
base_env = nil, -- Template environment
|
|
initialized = false, -- Initialization flag
|
|
env_pool = {}, -- Pre-allocated environment pool
|
|
pool_size = 0, -- Current pool size
|
|
max_pool_size = 8 -- Maximum pool size
|
|
}
|
|
|
|
-- Initialize base environment once
|
|
if not __env_system.initialized then
|
|
-- Create base environment with all standard libraries
|
|
local base = {}
|
|
|
|
-- Safe standard libraries
|
|
base.string = string
|
|
base.table = table
|
|
base.math = math
|
|
base.os = {
|
|
time = os.time,
|
|
date = os.date,
|
|
difftime = os.difftime,
|
|
clock = os.clock
|
|
}
|
|
|
|
-- Basic functions
|
|
base.print = print
|
|
base.tonumber = tonumber
|
|
base.tostring = tostring
|
|
base.type = type
|
|
base.pairs = pairs
|
|
base.ipairs = ipairs
|
|
base.next = next
|
|
base.select = select
|
|
base.unpack = unpack
|
|
base.pcall = pcall
|
|
base.xpcall = xpcall
|
|
base.error = error
|
|
base.assert = assert
|
|
|
|
-- Package system is shared for performance
|
|
base.package = {
|
|
loaded = package.loaded,
|
|
path = package.path,
|
|
preload = package.preload
|
|
}
|
|
|
|
base.http = http
|
|
base.cookie = cookie
|
|
-- http_client module is now part of http.client
|
|
|
|
-- Add registered custom modules
|
|
if __sandbox_modules then
|
|
for name, mod in pairs(__sandbox_modules) do
|
|
base[name] = mod
|
|
end
|
|
end
|
|
|
|
-- Store base environment
|
|
__env_system.base_env = base
|
|
__env_system.initialized = true
|
|
end
|
|
|
|
-- Global variable for tracking current environment
|
|
__last_env = nil
|
|
|
|
-- Fast environment creation with pre-allocation
|
|
function __get_sandbox_env(ctx)
|
|
local env
|
|
|
|
-- Try to reuse from pool
|
|
if __env_system.pool_size > 0 then
|
|
env = table.remove(__env_system.env_pool)
|
|
__env_system.pool_size = __env_system.pool_size - 1
|
|
|
|
-- Clear any previous context
|
|
env.ctx = ctx or nil
|
|
-- Clear any previous response
|
|
env._response = nil
|
|
else
|
|
-- Create new environment with metatable inheritance
|
|
env = setmetatable({}, {
|
|
__index = __env_system.base_env
|
|
})
|
|
|
|
-- Set context if provided
|
|
if ctx then
|
|
env.ctx = ctx
|
|
end
|
|
|
|
-- Install the fast require implementation
|
|
env.require = function(modname)
|
|
return __fast_require(env, modname)
|
|
end
|
|
|
|
-- Install cookie module methods directly into environment
|
|
env.cookie = {
|
|
get = function(name)
|
|
if type(name) ~= "string" then
|
|
error("cookie.get: name must be a string", 2)
|
|
end
|
|
|
|
if env.ctx and env.ctx.cookies and env.ctx.cookies[name] then
|
|
return tostring(env.ctx.cookies[name])
|
|
end
|
|
|
|
return nil
|
|
end,
|
|
|
|
set = cookie.set,
|
|
remove = cookie.remove
|
|
}
|
|
end
|
|
|
|
-- Store reference to current environment
|
|
__last_env = env
|
|
|
|
return env
|
|
end
|
|
|
|
-- Return environment to pool for reuse
|
|
function __recycle_env(env)
|
|
-- Only recycle if pool isn't full
|
|
if __env_system.pool_size < __env_system.max_pool_size then
|
|
-- Clear context reference to avoid memory leaks
|
|
env.ctx = nil
|
|
-- Don't clear response data - we need it for extraction
|
|
|
|
-- Add to pool
|
|
table.insert(__env_system.env_pool, env)
|
|
__env_system.pool_size = __env_system.pool_size + 1
|
|
end
|
|
end
|
|
|
|
-- Hyper-optimized sandbox executor
|
|
function __execute_sandbox(bytecode, ctx)
|
|
-- Get environment (from pool if available)
|
|
local env = __get_sandbox_env(ctx)
|
|
|
|
-- Set environment for bytecode
|
|
setfenv(bytecode, env)
|
|
|
|
-- Execute with protected call
|
|
local success, result = pcall(bytecode)
|
|
|
|
-- Recycle environment for future use
|
|
__recycle_env(env)
|
|
|
|
-- Process result
|
|
if not success then
|
|
error(result, 0)
|
|
end
|
|
|
|
return result
|
|
end
|
|
|
|
-- Run minimal GC for overall health
|
|
collectgarbage("step", 10)
|
|
`)
|
|
}
|
|
|
|
// registerModules registers custom modules in the Lua state
|
|
func (s *Sandbox) registerModules(state *luajit.State) error {
|
|
// Create or get module registry table
|
|
state.GetGlobal("__sandbox_modules")
|
|
if state.IsNil(-1) {
|
|
// Table doesn't exist, create it
|
|
state.Pop(1)
|
|
state.NewTable()
|
|
state.SetGlobal("__sandbox_modules")
|
|
state.GetGlobal("__sandbox_modules")
|
|
}
|
|
|
|
// Add modules to registry
|
|
for name, module := range s.modules {
|
|
state.PushString(name)
|
|
if err := state.PushValue(module); err != nil {
|
|
state.Pop(2)
|
|
return err
|
|
}
|
|
state.SetTable(-3)
|
|
}
|
|
|
|
// Pop module table
|
|
state.Pop(1)
|
|
return nil
|
|
}
|
|
|
|
// Execute runs bytecode in the sandbox
|
|
func (s *Sandbox) Execute(state *luajit.State, bytecode []byte, ctx map[string]any) (any, error) {
|
|
// Update custom modules if needed
|
|
if !s.initialized {
|
|
if err := s.registerModules(state); err != nil {
|
|
return nil, err
|
|
}
|
|
s.initialized = true
|
|
}
|
|
|
|
// Load bytecode
|
|
if err := state.LoadBytecode(bytecode, "script"); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Create context table if provided
|
|
if len(ctx) > 0 {
|
|
// Preallocate table with appropriate size
|
|
state.CreateTable(0, len(ctx))
|
|
|
|
// Add context entries
|
|
for k, v := range ctx {
|
|
state.PushString(k)
|
|
if err := state.PushValue(v); err != nil {
|
|
state.Pop(2)
|
|
return nil, err
|
|
}
|
|
state.SetTable(-3)
|
|
}
|
|
} else {
|
|
state.PushNil() // No context
|
|
}
|
|
|
|
// Get optimized sandbox executor
|
|
state.GetGlobal("__execute_sandbox")
|
|
|
|
// Setup call with correct argument order
|
|
state.PushCopy(-3) // Copy bytecode function
|
|
state.PushCopy(-3) // Copy context
|
|
|
|
// Clean up stack
|
|
state.Remove(-5) // Remove original bytecode
|
|
state.Remove(-4) // Remove original context
|
|
|
|
// Call optimized sandbox executor
|
|
if err := state.Call(2, 1); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Get result
|
|
result, err := state.ToValue(-1)
|
|
state.Pop(1) // Pop result
|
|
|
|
// Check if HTTP response was set
|
|
httpResponse, hasHTTPResponse := GetHTTPResponse(state)
|
|
if hasHTTPResponse {
|
|
httpResponse.Body = result
|
|
return httpResponse, err
|
|
}
|
|
|
|
return result, err
|
|
}
|