Moonshark/runner/sandbox.go

package runner

import (
	"fmt"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
	"github.com/valyala/fasthttp"
)

// Sandbox provides a secure execution environment for Lua scripts
type Sandbox struct {
	executorBytecode []byte
}

// NewSandbox creates a new sandbox environment
func NewSandbox() *Sandbox {
	return &Sandbox{}
}

// Setup initializes the sandbox in a Lua state
func (s *Sandbox) Setup(state *luajit.State, verbose bool) error {
	// Load all embedded modules and sandbox
	if err := loadSandboxIntoState(state, verbose); err != nil {
		return fmt.Errorf("failed to load sandbox: %w", err)
	}

	// Pre-compile the executor function for reuse
	executorCode := `return __execute`
	bytecode, err := state.CompileBytecode(executorCode, "executor")
	if err != nil {
		return fmt.Errorf("failed to compile executor: %w", err)
	}
	s.executorBytecode = bytecode

	// Register native functions
	if err := s.registerCoreFunctions(state); err != nil {
		return err
	}

	return nil
}

// Execute runs a Lua script in the sandbox with the given context
func (s *Sandbox) Execute(state *luajit.State, bytecode []byte, ctx *Context) (*Response, error) {
	// Create response object in Lua
	response := map[string]any{
		"status":   200,
		"headers":  make(map[string]string),
		"cookies":  []any{},
		"metadata": make(map[string]any),
		"session":  make(map[string]any),
	}

	// Load script bytecode (pushes function)
	if err := state.LoadBytecode(bytecode, "script"); err != nil {
		return nil, fmt.Errorf("failed to load bytecode: %w", err)
	}

	// Load executor (pushes __execute function)
	if err := state.LoadBytecode(s.executorBytecode, "executor"); err != nil {
		state.Pop(1) // Remove script function
		return nil, fmt.Errorf("failed to load executor: %w", err)
	}

	// Call the loaded executor to get __execute
	if err := state.Call(0, 1); err != nil {
		state.Pop(1) // Remove script function
		return nil, fmt.Errorf("failed to get executor: %w", err)
	}

	// Stack: [script_func, __execute]
	state.PushCopy(-2) // Copy script function
	// Stack: [script_func, __execute, script_func]

	// Push context
	if err := state.PushValue(ctx.Values); err != nil {
		state.Pop(3)
		return nil, fmt.Errorf("failed to push context: %w", err)
	}

	// Push response object
	if err := state.PushValue(response); err != nil {
		state.Pop(4)
		return nil, fmt.Errorf("failed to push response: %w", err)
	}

	// Stack: [script_func, __execute, script_func, ctx, response]
	// Call __execute(script_func, ctx, response)
	if err := state.Call(3, 1); err != nil {
		state.Pop(1) // Clean up
		return nil, fmt.Errorf("script execution failed: %w", err)
	}

	// Get the result
	result, _ := state.ToValue(-1)
	state.Pop(2) // Remove result and original script function

	// Extract response data directly from the response object we passed
	return s.buildResponse(response, result), nil
}

// buildResponse converts the Lua response object to a Go Response
func (s *Sandbox) buildResponse(luaResp map[string]any, body any) *Response {
	resp := NewResponse()
	resp.Body = body

	// Extract status
	if status, ok := luaResp["status"].(float64); ok {
		resp.Status = int(status)
	} else if status, ok := luaResp["status"].(int); ok {
		resp.Status = status
	}

	// Extract headers
	if headers, ok := luaResp["headers"].(map[string]any); ok {
		for k, v := range headers {
			if str, ok := v.(string); ok {
				resp.Headers[k] = str
			}
		}
	} else if headers, ok := luaResp["headers"].(map[string]string); ok {
		for k, v := range headers {
			resp.Headers[k] = v
		}
	}

	// Extract cookies
	if cookies, ok := luaResp["cookies"].([]any); ok {
		for _, cookieData := range cookies {
			if cookieMap, ok := cookieData.(map[string]any); ok {
				cookie := fasthttp.AcquireCookie()

				if name, ok := cookieMap["name"].(string); ok && name != "" {
					cookie.SetKey(name)
					if value, ok := cookieMap["value"].(string); ok {
						cookie.SetValue(value)
					}
					if path, ok := cookieMap["path"].(string); ok {
						cookie.SetPath(path)
					}
					if domain, ok := cookieMap["domain"].(string); ok {
						cookie.SetDomain(domain)
					}
					if httpOnly, ok := cookieMap["http_only"].(bool); ok {
						cookie.SetHTTPOnly(httpOnly)
					}
					if secure, ok := cookieMap["secure"].(bool); ok {
						cookie.SetSecure(secure)
					}
					if maxAge, ok := cookieMap["max_age"].(float64); ok {
						cookie.SetMaxAge(int(maxAge))
					} else if maxAge, ok := cookieMap["max_age"].(int); ok {
						cookie.SetMaxAge(maxAge)
					}

					resp.Cookies = append(resp.Cookies, cookie)
				} else {
					fasthttp.ReleaseCookie(cookie)
				}
			}
		}
	}

	// Extract metadata
	if metadata, ok := luaResp["metadata"].(map[string]any); ok {
		for k, v := range metadata {
			resp.Metadata[k] = v
		}
	}

	// Extract session data
	if session, ok := luaResp["session"].(map[string]any); ok {
		for k, v := range session {
			resp.SessionData[k] = v
		}
	}

	return resp
}

// registerCoreFunctions registers all built-in functions in the Lua state
func (s *Sandbox) registerCoreFunctions(state *luajit.State) error {
	// Register your native functions here
	// This stays the same as your current implementation
	return nil
}