package runner import ( "fmt" "sync" "github.com/valyala/bytebufferpool" "github.com/valyala/fasthttp" "Moonshark/core/utils/logger" 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 state.PushBoolean(false) state.SetGlobal("__session_modified") } else { // Initialize empty session if err := state.DoString("__session_data = {}; __session_modified = false"); 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 { for k, v := range sessionData { response.SessionData[k] = v } } } 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) }