Moonshark/core/runner/Sandbox.go

package runner

import (
	"fmt"
	"sync"

	"github.com/valyala/fasthttp"

	"Moonshark/core/utils/logger"

	"maps"

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

// Error represents a simple error string
type Error string

func (e Error) Error() string {
	return string(e)
}

// Error types
var (
	ErrSandboxNotInitialized = Error("sandbox not initialized")
)

// Sandbox provides a secure execution environment for Lua scripts
type Sandbox struct {
	modules map[string]any
	debug   bool
	mu      sync.RWMutex
}

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

// AddModule adds a module to the sandbox environment
func (s *Sandbox) AddModule(name string, module any) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.modules[name] = module
	logger.Debug("Added module: %s", name)
}

// Setup initializes the sandbox in a Lua state
func (s *Sandbox) Setup(state *luajit.State, verbose bool) error {
	if verbose {
		logger.Server("Setting up sandbox...")
	}

	if err := loadSandboxIntoState(state, verbose); err != nil {
		logger.ErrorCont("Failed to load sandbox: %v", err)
		return err
	}

	if err := s.registerCoreFunctions(state); err != nil {
		logger.ErrorCont("Failed to register core functions: %v", err)
		return err
	}

	s.mu.RLock()
	for name, module := range s.modules {
		logger.DebugCont("Registering module: %s", name)
		if err := state.PushValue(module); err != nil {
			s.mu.RUnlock()
			logger.ErrorCont("Failed to register module %s: %v", name, err)
			return err
		}
		state.SetGlobal(name)
	}
	s.mu.RUnlock()

	if verbose {
		logger.ServerCont("Sandbox setup complete")
	}
	return nil
}

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

	if err := state.RegisterGoFunction("__generate_token", generateToken); err != nil {
		return err
	}

	// Additional registrations can be added here

	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) {
	// Get the execution function first
	state.GetGlobal("__execute_script")
	if !state.IsFunction(-1) {
		state.Pop(1)
		return nil, ErrSandboxNotInitialized
	}

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

	// Push context values
	if err := state.PushTable(ctx.Values); err != nil {
		state.Pop(2) // Pop bytecode and __execute_script
		return nil, err
	}

	// Execute with 2 args, 1 result
	if err := state.Call(2, 1); err != nil {
		return nil, fmt.Errorf("script execution failed: %w", err)
	}

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

	response := NewResponse()
	if err == nil {
		response.Body = body
	}

	extractHTTPResponseData(state, response)

	return response, nil
}

// extractResponseData pulls response info from the Lua state
func extractHTTPResponseData(state *luajit.State, response *Response) {
	state.GetGlobal("__http_responses")
	if !state.IsTable(-1) {
		state.Pop(1)
		return
	}

	state.PushNumber(1)
	state.GetTable(-2)
	if !state.IsTable(-1) {
		state.Pop(2)
		return
	}

	// Extract status
	state.GetField(-1, "status")
	if state.IsNumber(-1) {
		response.Status = int(state.ToNumber(-1))
	}
	state.Pop(1)

	// Extract headers
	state.GetField(-1, "headers")
	if state.IsTable(-1) {
		state.PushNil() // Start iteration
		for state.Next(-2) {
			if state.IsString(-2) && state.IsString(-1) {
				key := state.ToString(-2)
				value := state.ToString(-1)
				response.Headers[key] = value
			}
			state.Pop(1)
		}
	}
	state.Pop(1)

	// Extract cookies
	state.GetField(-1, "cookies")
	if state.IsTable(-1) {
		length := state.GetTableLength(-1)
		for i := 1; i <= length; i++ {
			state.PushNumber(float64(i))
			state.GetTable(-2)

			if state.IsTable(-1) {
				extractCookie(state, response)
			}
			state.Pop(1)
		}
	}
	state.Pop(1)

	// Extract metadata
	state.GetField(-1, "metadata")
	if state.IsTable(-1) {
		table, err := state.ToTable(-1)
		if err == nil {
			maps.Copy(response.Metadata, table)
		}
	}
	state.Pop(1)

	// Clean up
	state.Pop(2)
}

// extractCookie pulls cookie data from the current table on the stack
func extractCookie(state *luajit.State, response *Response) {
	cookie := fasthttp.AcquireCookie()

	// Get name (required)
	state.GetField(-1, "name")
	if !state.IsString(-1) {
		state.Pop(1)
		fasthttp.ReleaseCookie(cookie)
		return
	}
	cookie.SetKey(state.ToString(-1))
	state.Pop(1)

	// Get value
	state.GetField(-1, "value")
	if state.IsString(-1) {
		cookie.SetValue(state.ToString(-1))
	}
	state.Pop(1)

	// Get path
	state.GetField(-1, "path")
	if state.IsString(-1) {
		cookie.SetPath(state.ToString(-1))
	} else {
		cookie.SetPath("/") // Default
	}
	state.Pop(1)

	// Get domain
	state.GetField(-1, "domain")
	if state.IsString(-1) {
		cookie.SetDomain(state.ToString(-1))
	}
	state.Pop(1)

	// Get other parameters
	state.GetField(-1, "http_only")
	if state.IsBoolean(-1) {
		cookie.SetHTTPOnly(state.ToBoolean(-1))
	}
	state.Pop(1)

	state.GetField(-1, "secure")
	if state.IsBoolean(-1) {
		cookie.SetSecure(state.ToBoolean(-1))
	}
	state.Pop(1)

	state.GetField(-1, "max_age")
	if state.IsNumber(-1) {
		cookie.SetMaxAge(int(state.ToNumber(-1)))
	}
	state.Pop(1)

	response.Cookies = append(response.Cookies, cookie)
}