Moonshark/core/runner/Sandbox.go

package runner

import (
	"fmt"
	"sync"

	"github.com/valyala/bytebufferpool"
	"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,
	}
}

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

// debugLog logs a message if debug mode is enabled
func (s *Sandbox) debugLog(format string, args ...interface{}) {
	if s.debug {
		logger.Debug("Sandbox "+format, args...)
	}
}

// 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
	s.debugLog("Added module: %s", name)
}

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

	// Load the sandbox code
	if err := loadSandboxIntoState(state); err != nil {
		s.debugLog("Failed to load sandbox: %v", err)
		return err
	}

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

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

	s.debugLog("Sandbox setup complete")
	return nil
}

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

	// Register utility functions
	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) {
	s.debugLog("Executing script...")

	// Create a response object
	response := NewResponse()

	// Get a buffer for string operations
	buf := bytebufferpool.Get()
	defer bytebufferpool.Put(buf)

	// Load bytecode
	if err := state.LoadBytecode(bytecode, "script"); err != nil {
		ReleaseResponse(response)
		s.debugLog("Failed to load bytecode: %v", err)
		return nil, fmt.Errorf("failed to load script: %w", err)
	}

	// Initialize session data in Lua
	if ctx.SessionID != "" {
		// Set session ID
		state.PushString(ctx.SessionID)
		state.SetGlobal("__session_id")

		// Set session data
		if err := state.PushTable(ctx.SessionData); err != nil {
			ReleaseResponse(response)
			s.debugLog("Failed to push session data: %v", err)
			return nil, err
		}
		state.SetGlobal("__session_data")

		// Reset modification flag and tracking
		state.PushBoolean(false)
		state.SetGlobal("__session_modified")

		// Create empty modified keys table
		state.NewTable()
		state.SetGlobal("__session_modified_keys")
	} else {
		// Initialize empty session
		if err := state.DoString("__session_data = {}; __session_modified = false; __session_modified_keys = {}"); err != nil {
			s.debugLog("Failed to initialize empty session data: %v", err)
		}
	}

	// Set up context values for execution
	if err := state.PushTable(ctx.Values); err != nil {
		ReleaseResponse(response)
		s.debugLog("Failed to push context values: %v", err)
		return nil, err
	}

	// Get the execution function
	state.GetGlobal("__execute_script")
	if !state.IsFunction(-1) {
		state.Pop(1) // Pop non-function
		ReleaseResponse(response)
		s.debugLog("__execute_script is not a function")
		return nil, ErrSandboxNotInitialized
	}

	// Push function and context to stack
	state.PushCopy(-2) // bytecode
	state.PushCopy(-2) // context

	// Remove duplicates
	state.Remove(-4)
	state.Remove(-3)

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

	// Set response body from result
	body, err := state.ToValue(-1)
	if err == nil {
		response.Body = body
	}
	state.Pop(1)

	// Extract HTTP response data from Lua state
	s.extractResponseData(state, response)

	return response, nil
}

// extractResponseData pulls response info from the Lua state
func (s *Sandbox) extractResponseData(state *luajit.State, response *Response) {
	// Get HTTP response
	state.GetGlobal("__http_responses")
	if !state.IsNil(-1) && state.IsTable(-1) {
		state.PushNumber(1)
		state.GetTable(-2)

		if !state.IsNil(-1) && state.IsTable(-1) {
			// 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) {
						s.extractCookie(state, response)
					}
					state.Pop(1)
				}
			}
			state.Pop(1)

			// Extract metadata if present
			state.GetField(-1, "metadata")
			if state.IsTable(-1) {
				table, err := state.ToTable(-1)
				if err == nil {
					for k, v := range table {
						response.Metadata[k] = v
					}
				}
			}
			state.Pop(1)
		}
		state.Pop(1)
	}
	state.Pop(1)

	// Extract session data
	state.GetGlobal("__session_modified")
	if state.IsBoolean(-1) && state.ToBoolean(-1) {
		response.SessionModified = true

		// Get session ID
		state.GetGlobal("__session_id")
		if state.IsString(-1) {
			response.SessionID = state.ToString(-1)
		}
		state.Pop(1)

		// Get session data
		state.GetGlobal("__session_data")
		if state.IsTable(-1) {
			sessionData, err := state.ToTable(-1)
			if err == nil {
				maps.Copy(response.SessionData, sessionData)
			}
		}
		state.Pop(1)
	}
	state.Pop(1)
}

// extractCookie pulls cookie data from the current table on the stack
func (s *Sandbox) 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)
}