309 lines
6.8 KiB
Go
309 lines
6.8 KiB
Go
package runner
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
|
|
)
|
|
|
|
// Common errors
|
|
var (
|
|
ErrRunnerClosed = errors.New("lua runner is closed")
|
|
ErrInitFailed = errors.New("initialization failed")
|
|
)
|
|
|
|
// StateInitFunc is a function that initializes a Lua state
|
|
type StateInitFunc func(*luajit.State) error
|
|
|
|
// LuaRunner runs Lua scripts using a single Lua state
|
|
type LuaRunner struct {
|
|
state *luajit.State // The Lua state
|
|
jobQueue chan job // Channel for incoming jobs
|
|
isRunning atomic.Bool // Flag indicating if the runner is active
|
|
mu sync.RWMutex // Mutex for thread safety
|
|
wg sync.WaitGroup // WaitGroup for clean shutdown
|
|
initFunc StateInitFunc // Optional function to initialize Lua state
|
|
bufferSize int // Size of the job queue buffer
|
|
}
|
|
|
|
// NewRunner creates a new LuaRunner
|
|
func NewRunner(options ...RunnerOption) (*LuaRunner, error) {
|
|
// Default configuration
|
|
runner := &LuaRunner{
|
|
bufferSize: 10, // Default buffer size
|
|
}
|
|
|
|
// Apply options
|
|
for _, opt := range options {
|
|
opt(runner)
|
|
}
|
|
|
|
// Initialize Lua state
|
|
state := luajit.New()
|
|
if state == nil {
|
|
return nil, errors.New("failed to create Lua state")
|
|
}
|
|
runner.state = state
|
|
|
|
// Create job queue
|
|
runner.jobQueue = make(chan job, runner.bufferSize)
|
|
runner.isRunning.Store(true)
|
|
|
|
// Set up sandbox
|
|
if err := runner.setupSandbox(); err != nil {
|
|
state.Close()
|
|
return nil, ErrInitFailed
|
|
}
|
|
|
|
// Run init function if provided
|
|
if runner.initFunc != nil {
|
|
if err := runner.initFunc(state); err != nil {
|
|
state.Close()
|
|
return nil, ErrInitFailed
|
|
}
|
|
}
|
|
|
|
// Start the event loop
|
|
runner.wg.Add(1)
|
|
go runner.eventLoop()
|
|
|
|
return runner, nil
|
|
}
|
|
|
|
// RunnerOption defines a functional option for configuring the LuaRunner
|
|
type RunnerOption func(*LuaRunner)
|
|
|
|
// WithBufferSize sets the job queue buffer size
|
|
func WithBufferSize(size int) RunnerOption {
|
|
return func(r *LuaRunner) {
|
|
if size > 0 {
|
|
r.bufferSize = size
|
|
}
|
|
}
|
|
}
|
|
|
|
// WithInitFunc sets the init function for the Lua state
|
|
func WithInitFunc(initFunc StateInitFunc) RunnerOption {
|
|
return func(r *LuaRunner) {
|
|
r.initFunc = initFunc
|
|
}
|
|
}
|
|
|
|
// setupSandbox initializes the sandbox environment
|
|
func (r *LuaRunner) setupSandbox() error {
|
|
// This is the Lua script that creates our sandbox function
|
|
setupScript := `
|
|
-- Create a function to run code in a sandbox environment
|
|
function __create_sandbox()
|
|
-- Create new environment table
|
|
local env = {}
|
|
|
|
-- Add standard library modules (can be restricted as needed)
|
|
env.string = string
|
|
env.table = table
|
|
env.math = math
|
|
env.os = {
|
|
time = os.time,
|
|
date = os.date,
|
|
difftime = os.difftime,
|
|
clock = os.clock
|
|
}
|
|
env.tonumber = tonumber
|
|
env.tostring = tostring
|
|
env.type = type
|
|
env.pairs = pairs
|
|
env.ipairs = ipairs
|
|
env.next = next
|
|
env.select = select
|
|
env.unpack = unpack
|
|
env.pcall = pcall
|
|
env.xpcall = xpcall
|
|
env.error = error
|
|
env.assert = assert
|
|
|
|
-- Allow access to package.loaded for modules
|
|
env.require = function(name)
|
|
return package.loaded[name]
|
|
end
|
|
|
|
-- Create metatable to restrict access to _G
|
|
local mt = {
|
|
__index = function(t, k)
|
|
-- First check in env table
|
|
local v = rawget(env, k)
|
|
if v ~= nil then return v end
|
|
|
|
-- If not found, check for registered modules/functions
|
|
local moduleValue = _G[k]
|
|
if type(moduleValue) == "table" or
|
|
type(moduleValue) == "function" then
|
|
return moduleValue
|
|
end
|
|
|
|
return nil
|
|
end,
|
|
__newindex = function(t, k, v)
|
|
rawset(env, k, v)
|
|
end
|
|
}
|
|
|
|
setmetatable(env, mt)
|
|
return env
|
|
end
|
|
|
|
-- Create function to execute code with a sandbox
|
|
function __run_sandboxed(f, ctx)
|
|
local env = __create_sandbox()
|
|
|
|
-- Add context to the environment if provided
|
|
if ctx then
|
|
env.ctx = ctx
|
|
end
|
|
|
|
-- Set the environment and run the function
|
|
setfenv(f, env)
|
|
return f()
|
|
end
|
|
`
|
|
|
|
return r.state.DoString(setupScript)
|
|
}
|
|
|
|
// eventLoop processes jobs from the queue
|
|
func (r *LuaRunner) eventLoop() {
|
|
defer r.wg.Done()
|
|
defer r.state.Close()
|
|
|
|
// Process jobs until closure
|
|
for job := range r.jobQueue {
|
|
// Execute the job and send result
|
|
result := r.executeJob(job)
|
|
select {
|
|
case job.Result <- result:
|
|
// Result sent successfully
|
|
default:
|
|
// Result channel closed or full, discard the result
|
|
}
|
|
}
|
|
}
|
|
|
|
// executeJob runs a script in the sandbox environment
|
|
func (r *LuaRunner) executeJob(j job) JobResult {
|
|
// Re-run init function if needed
|
|
if r.initFunc != nil {
|
|
if err := r.initFunc(r.state); err != nil {
|
|
return JobResult{nil, err}
|
|
}
|
|
}
|
|
|
|
// Set up context if provided
|
|
if j.Context != nil {
|
|
// Push context table
|
|
r.state.NewTable()
|
|
|
|
// Add values to context table
|
|
for key, value := range j.Context.Values {
|
|
// Push key
|
|
r.state.PushString(key)
|
|
|
|
// Push value
|
|
if err := r.state.PushValue(value); err != nil {
|
|
return JobResult{nil, err}
|
|
}
|
|
|
|
// Set table[key] = value
|
|
r.state.SetTable(-3)
|
|
}
|
|
} else {
|
|
// Push nil if no context
|
|
r.state.PushNil()
|
|
}
|
|
|
|
// Load bytecode
|
|
if err := r.state.LoadBytecode(j.Bytecode, "script"); err != nil {
|
|
r.state.Pop(1) // Pop context
|
|
return JobResult{nil, err}
|
|
}
|
|
|
|
// Get the sandbox runner function
|
|
r.state.GetGlobal("__run_sandboxed")
|
|
|
|
// Push loaded function and context as arguments
|
|
r.state.PushCopy(-2) // Copy the loaded function
|
|
r.state.PushCopy(-4) // Copy the context table or nil
|
|
|
|
// Remove the original function and context
|
|
r.state.Remove(-5) // Remove original context
|
|
r.state.Remove(-4) // Remove original function
|
|
|
|
// Call the sandbox runner with 2 args (function and context), expecting 1 result
|
|
if err := r.state.Call(2, 1); err != nil {
|
|
return JobResult{nil, err}
|
|
}
|
|
|
|
// Get result
|
|
value, err := r.state.ToValue(-1)
|
|
r.state.Pop(1) // Pop result
|
|
|
|
return JobResult{value, err}
|
|
}
|
|
|
|
// RunWithContext executes a Lua script with context and timeout
|
|
func (r *LuaRunner) RunWithContext(ctx context.Context, bytecode []byte, execCtx *Context) (any, error) {
|
|
r.mu.RLock()
|
|
if !r.isRunning.Load() {
|
|
r.mu.RUnlock()
|
|
return nil, ErrRunnerClosed
|
|
}
|
|
r.mu.RUnlock()
|
|
|
|
resultChan := make(chan JobResult, 1)
|
|
j := job{
|
|
Bytecode: bytecode,
|
|
Context: execCtx,
|
|
Result: resultChan,
|
|
}
|
|
|
|
// Submit job with context
|
|
select {
|
|
case r.jobQueue <- j:
|
|
// Job submitted
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
}
|
|
|
|
// Wait for result with context
|
|
select {
|
|
case result := <-resultChan:
|
|
return result.Value, result.Error
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
}
|
|
}
|
|
|
|
// Run executes a Lua script
|
|
func (r *LuaRunner) Run(bytecode []byte, execCtx *Context) (any, error) {
|
|
return r.RunWithContext(context.Background(), bytecode, execCtx)
|
|
}
|
|
|
|
// Close gracefully shuts down the LuaRunner
|
|
func (r *LuaRunner) Close() error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
if !r.isRunning.Load() {
|
|
return ErrRunnerClosed
|
|
}
|
|
|
|
r.isRunning.Store(false)
|
|
close(r.jobQueue)
|
|
|
|
// Wait for event loop to finish
|
|
r.wg.Wait()
|
|
|
|
return nil
|
|
}
|