Moonshark/core/runner/Sandbox.go

package runner

import (
	"fmt"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
	"git.sharkk.net/Sky/Moonshark/core/logger"
)

// Sandbox manages a simplified Lua environment
type Sandbox struct {
	modules map[string]any // Custom modules for environment
	debug   bool           // Enable debug output
}

// NewSandbox creates a new sandbox
func NewSandbox() *Sandbox {
	return &Sandbox{
		modules: make(map[string]any),
		debug:   false,
	}
}

// EnableDebug turns on debug output
func (s *Sandbox) EnableDebug() {
	s.debug = true
}

// AddModule adds a module to the sandbox environment
func (s *Sandbox) AddModule(name string, module any) {
	s.modules[name] = module
}

// debugLog prints debug messages if debug is enabled
func (s *Sandbox) debugLog(format string, args ...interface{}) {
	if s.debug {
		logger.Debug("[Sandbox Debug] "+format, args...)
	}
}

// Setup initializes the sandbox in a Lua state
func (s *Sandbox) Setup(state *luajit.State) error {
	s.debugLog("Setting up sandbox environment")

	// Register modules in the global environment
	for name, module := range s.modules {
		s.debugLog("Registering module: %s", name)
		if err := state.PushValue(module); err != nil {
			s.debugLog("Failed to register module %s: %v", name, err)
			return err
		}
		state.SetGlobal(name)
	}

	// Initialize simple environment setup
	err := state.DoString(`
		-- Global tables for response handling
		__http_responses = __http_responses or {}

		-- Simple environment creation
		function __create_env(ctx)
			-- Create environment inheriting from _G
			local env = setmetatable({}, {__index = _G})

			-- Add context if provided
			if ctx then
				env.ctx = ctx
			end

			return env
		end

		-- Execute script with clean environment
		function __execute_script(fn, ctx)
			-- Clear previous responses
			__http_responses[1] = nil

			-- Create environment
			local env = __create_env(ctx)

			-- Set environment for function
			setfenv(fn, env)

			-- Execute with protected call
			local ok, result = pcall(fn)
			if not ok then
				error(result, 0)
			end

			return result
		end
	`)

	if err != nil {
		s.debugLog("Failed to set up sandbox: %v", err)
		return err
	}

	s.debugLog("Sandbox setup complete")

	// Verify HTTP module is accessible
	httpResult, _ := state.ExecuteWithResult(`
		if type(http) == "table" and
		   type(http.client) == "table" and
		   type(http.client.get) == "function" then
			return "HTTP module verified OK"
		else
			local status = {
				http = type(http),
				client = type(http) == "table" and type(http.client) or "N/A",
				get = type(http) == "table" and type(http.client) == "table" and type(http.client.get) or "N/A"
			}
			return status
		end
	`)

	s.debugLog("HTTP verification result: %v", httpResult)
	return nil
}

// Execute runs bytecode in the sandbox
func (s *Sandbox) Execute(state *luajit.State, bytecode []byte, ctx map[string]any) (any, error) {
	// Load bytecode
	if err := state.LoadBytecode(bytecode, "script"); err != nil {
		s.debugLog("Failed to load bytecode: %v", err)
		return nil, err
	}

	// Prepare context
	if ctx != nil {
		state.CreateTable(0, len(ctx))
		for k, v := range ctx {
			state.PushString(k)
			if err := state.PushValue(v); err != nil {
				state.Pop(2)
				s.debugLog("Failed to push context value %s: %v", k, err)
				return nil, err
			}
			state.SetTable(-3)
		}
	} else {
		state.PushNil() // No context
	}

	// Get execution function
	state.GetGlobal("__execute_script")
	if !state.IsFunction(-1) {
		state.Pop(2) // Pop nil and non-function
		s.debugLog("__execute_script is not a function")
		return nil, fmt.Errorf("sandbox execution function not found")
	}

	// Push arguments
	state.PushCopy(-3) // bytecode function
	state.PushCopy(-3) // context

	// Clean up stack
	state.Remove(-5) // original bytecode
	state.Remove(-4) // original context

	// Call with 2 args, 1 result
	if err := state.Call(2, 1); err != nil {
		s.debugLog("Execution failed: %v", err)
		return nil, err
	}

	// Get result
	result, err := state.ToValue(-1)
	state.Pop(1)

	// Check for HTTP response
	httpResponse, hasResponse := GetHTTPResponse(state)
	if hasResponse {
		httpResponse.Body = result
		return httpResponse, nil
	}

	return result, err
}