Moonshark/runner/sandbox.go

package runner

import (
	"Moonshark/runner/lualibs"
	"Moonshark/runner/sqlite"
	"fmt"

	"maps"

	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, stateIndex int, 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)
	}

	// Set the state index as a global variable
	state.PushNumber(float64(stateIndex))
	state.SetGlobal("__STATE_INDEX")

	// 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) {
	// Load script and executor
	if err := state.LoadBytecode(bytecode, "script"); err != nil {
		return nil, fmt.Errorf("failed to load bytecode: %w", err)
	}

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

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

	// Prepare response object
	response := map[string]any{
		"status":   200,
		"headers":  make(map[string]string),
		"cookies":  []any{},
		"metadata": make(map[string]any),
		"session":  make(map[string]any),
		"flash":    make(map[string]any),
	}

	// Call __execute(script_func, ctx, response)
	state.PushCopy(-2) // script function
	state.PushValue(ctx.Values)
	state.PushValue(response)

	if err := state.Call(3, 1); err != nil {
		state.Pop(1)
		return nil, fmt.Errorf("script execution failed: %w", err)
	}

	// Extract result
	result, _ := state.ToValue(-1)
	state.Pop(2) // Clean up

	var modifiedResponse map[string]any
	var scriptResult any

	if arr, ok := result.([]any); ok && len(arr) >= 2 {
		scriptResult = arr[0]
		if resp, ok := arr[1].(map[string]any); ok {
			modifiedResponse = resp
		}
	}

	if modifiedResponse == nil {
		scriptResult = result
		modifiedResponse = response
	}

	return s.buildResponse(modifiedResponse, scriptResult), 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
			}
		}
	}

	// 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 - simplified
	if metadata, ok := luaResp["metadata"].(map[string]any); ok {
		maps.Copy(resp.Metadata, metadata)
	}

	// Extract session data - simplified
	if session, ok := luaResp["session"].(map[string]any); ok {
		maps.Copy(resp.SessionData, session)
	}

	// Extract flash data and add to metadata for processing by server
	if flash, ok := luaResp["flash"].(map[string]any); ok && len(flash) > 0 {
		resp.Metadata["flash"] = flash
	}

	return resp
}

// registerCoreFunctions registers all built-in functions in the Lua state
func (s *Sandbox) registerCoreFunctions(state *luajit.State) error {
	if err := lualibs.RegisterCryptoFunctions(state); err != nil {
		return err
	}

	if err := lualibs.RegisterEnvFunctions(state); err != nil {
		return err
	}

	if err := lualibs.RegisterFSFunctions(state); err != nil {
		return err
	}

	if err := lualibs.RegisterHttpFunctions(state); err != nil {
		return err
	}

	if err := lualibs.RegisterPasswordFunctions(state); err != nil {
		return err
	}

	if err := sqlite.RegisterSQLiteFunctions(state); err != nil {
		return err
	}

	if err := lualibs.RegisterUtilFunctions(state); err != nil {
		return err
	}

	return nil
}