From 0259d1a135aac788980925b07056a57fd08f023a Mon Sep 17 00:00:00 2001 From: Sky Johnson Date: Sat, 22 Mar 2025 22:27:58 -0500 Subject: [PATCH] add http utils --- core/http/server.go | 44 +++++++++----- core/runner/luarunner.go | 19 ++++-- core/runner/sandbox.go | 127 +++++++++++++++++++++++---------------- 3 files changed, 115 insertions(+), 75 deletions(-) diff --git a/core/http/server.go b/core/http/server.go index 2f43919..a9f9755 100644 --- a/core/http/server.go +++ b/core/http/server.go @@ -170,27 +170,33 @@ func writeResponse(w http.ResponseWriter, result any, log *logger.Logger) { return } - switch res := result.(type) { - case string: - // String result - w.Header().Set("Content-Type", contentTypePlain) - w.Write([]byte(res)) + // Check for HTTPResponse type + if httpResp, ok := result.(*runner.HTTPResponse); ok { + // Set response headers + for name, value := range httpResp.Headers { + w.Header().Set(name, value) + } - case map[string]any, []any: - // Table or array result - convert to JSON - w.Header().Set("Content-Type", contentTypeJSON) - data, err := json.Marshal(res) - if err != nil { - log.Error("Failed to marshal response: %v", err) - http.Error(w, "Internal Server Error", http.StatusInternalServerError) + // Set status code + w.WriteHeader(httpResp.Status) + + // Process the body based on its type + if httpResp.Body == nil { return } - w.Write(data) + result = httpResp.Body // Set result to body for processing below + } + + switch res := result.(type) { + case string: + // String result - plain text + setContentTypeIfMissing(w, contentTypePlain) + w.Write([]byte(res)) default: - // Other result types - convert to JSON - w.Header().Set("Content-Type", contentTypeJSON) - data, err := json.Marshal(result) + // All other types - convert to JSON + setContentTypeIfMissing(w, contentTypeJSON) + data, err := json.Marshal(res) if err != nil { log.Error("Failed to marshal response: %v", err) http.Error(w, "Internal Server Error", http.StatusInternalServerError) @@ -199,3 +205,9 @@ func writeResponse(w http.ResponseWriter, result any, log *logger.Logger) { w.Write(data) } } + +func setContentTypeIfMissing(w http.ResponseWriter, contentType string) { + if w.Header().Get("Content-Type") == "" { + w.Header().Set("Content-Type", contentType) + } +} diff --git a/core/runner/luarunner.go b/core/runner/luarunner.go index 718b78e..c3764bb 100644 --- a/core/runner/luarunner.go +++ b/core/runner/luarunner.go @@ -103,18 +103,25 @@ func NewRunner(options ...RunnerOption) (*LuaRunner, error) { return nil, ErrInitFailed } + // Initialize HTTP module BEFORE sandbox setup + httpInit := HTTPModuleInitFunc() + if err := httpInit(state); err != nil { + state.Close() + return nil, ErrInitFailed + } + + // Set up sandbox AFTER HTTP module is initialized + if err := runner.sandbox.Setup(state); err != nil { + state.Close() + return nil, ErrInitFailed + } + // Preload all modules into package.loaded if err := runner.moduleLoader.PreloadAllModules(state); err != nil { state.Close() return nil, errors.New("failed to preload modules") } - // Set up sandbox - if err := runner.sandbox.Setup(state); err != nil { - state.Close() - return nil, ErrInitFailed - } - // Run init function if provided if runner.initFunc != nil { if err := runner.initFunc(state); err != nil { diff --git a/core/runner/sandbox.go b/core/runner/sandbox.go index a4984d1..ccacc79 100644 --- a/core/runner/sandbox.go +++ b/core/runner/sandbox.go @@ -12,10 +12,12 @@ type Sandbox struct { // NewSandbox creates a new sandbox func NewSandbox() *Sandbox { - return &Sandbox{ + s := &Sandbox{ modules: make(map[string]any), initialized: false, } + + return s } // AddModule adds a module to the sandbox environment @@ -32,63 +34,69 @@ func (s *Sandbox) Setup(state *luajit.State) error { // Create high-performance persistent environment return state.DoString(` - -- Global shared environment (created once) - __env_system = __env_system or { - base_env = nil, -- Template environment - initialized = false, -- Initialization flag - env_pool = {}, -- Pre-allocated environment pool - pool_size = 0, -- Current pool size - max_pool_size = 8 -- Maximum pool size - } + -- Global shared environment (created once) + __env_system = __env_system or { + base_env = nil, -- Template environment + initialized = false, -- Initialization flag + env_pool = {}, -- Pre-allocated environment pool + pool_size = 0, -- Current pool size + max_pool_size = 8 -- Maximum pool size + } - -- Initialize base environment once - if not __env_system.initialized then - -- Create base environment with all standard libraries - local base = {} + -- Initialize base environment once + if not __env_system.initialized then + -- Create base environment with all standard libraries + local base = {} - -- Safe standard libraries - base.string = string - base.table = table - base.math = math - base.os = { - time = os.time, - date = os.date, - difftime = os.difftime, - clock = os.clock - } + -- Safe standard libraries + base.string = string + base.table = table + base.math = math + base.os = { + time = os.time, + date = os.date, + difftime = os.difftime, + clock = os.clock + } - -- Basic functions - base.tonumber = tonumber - base.tostring = tostring - base.type = type - base.pairs = pairs - base.ipairs = ipairs - base.next = next - base.select = select - base.unpack = unpack - base.pcall = pcall - base.xpcall = xpcall - base.error = error - base.assert = assert + -- Basic functions + base.tonumber = tonumber + base.tostring = tostring + base.type = type + base.pairs = pairs + base.ipairs = ipairs + base.next = next + base.select = select + base.unpack = unpack + base.pcall = pcall + base.xpcall = xpcall + base.error = error + base.assert = assert - -- Package system is shared for performance - base.package = { - loaded = package.loaded, - path = package.path, - preload = package.preload - } + -- Package system is shared for performance + base.package = { + loaded = package.loaded, + path = package.path, + preload = package.preload + } - -- Add registered custom modules - if __sandbox_modules then - for name, mod in pairs(__sandbox_modules) do - base[name] = mod - end - end + -- Add HTTP module explicitly to the base environment + base.http = http - -- Store base environment - __env_system.base_env = base - __env_system.initialized = true - end + -- Add registered custom modules + if __sandbox_modules then + for name, mod in pairs(__sandbox_modules) do + base[name] = mod + end + end + + -- Store base environment + __env_system.base_env = base + __env_system.initialized = true + end + + -- Global variable for tracking current environment + __last_env = nil -- Fast environment creation with pre-allocation function __get_sandbox_env(ctx) @@ -101,6 +109,8 @@ func (s *Sandbox) Setup(state *luajit.State) error { -- Clear any previous context env.ctx = ctx or nil + -- Clear any previous response + env._response = nil else -- Create new environment with metatable inheritance env = setmetatable({}, { @@ -118,6 +128,9 @@ func (s *Sandbox) Setup(state *luajit.State) error { end end + -- Store reference to current environment + __last_env = env + return env end @@ -127,6 +140,7 @@ func (s *Sandbox) Setup(state *luajit.State) error { if __env_system.pool_size < __env_system.max_pool_size then -- Clear context reference to avoid memory leaks env.ctx = nil + -- Don't clear response data - we need it for extraction -- Add to pool table.insert(__env_system.env_pool, env) @@ -239,7 +253,14 @@ func (s *Sandbox) Execute(state *luajit.State, bytecode []byte, ctx map[string]a // Get result result, err := state.ToValue(-1) - state.Pop(1) + state.Pop(1) // Pop result + + // Check if HTTP response was set + httpResponse, hasHTTPResponse := GetHTTPResponse(state) + if hasHTTPResponse { + httpResponse.Body = result + return httpResponse, err + } return result, err }