From 727ce89da74454205e0a4d0b4f061de120d9bf2a Mon Sep 17 00:00:00 2001 From: Sky Johnson Date: Wed, 26 Mar 2025 09:39:58 -0500 Subject: [PATCH] Implement cookie utilities --- core/http/server.go | 22 ++++- core/runner/cookies.go | 172 +++++++++++++++++++++++++++++++++++++++ core/runner/http.go | 24 ++++++ core/runner/luarunner.go | 8 +- core/runner/sandbox.go | 72 ++++++++++------ 5 files changed, 263 insertions(+), 35 deletions(-) create mode 100644 core/runner/cookies.go diff --git a/core/http/server.go b/core/http/server.go index 2681a2b..336a28d 100644 --- a/core/http/server.go +++ b/core/http/server.go @@ -156,7 +156,7 @@ func (s *Server) handleLuaRoute(w http.ResponseWriter, r *http.Request, bytecode ctx.Set("path", r.URL.Path) ctx.Set("host", r.Host) - // Inline the header conversion (previously makeHeaderMap) + // Add headers to context headerMap := make(map[string]any, len(r.Header)) for name, values := range r.Header { if len(values) == 1 { @@ -167,11 +167,20 @@ func (s *Server) handleLuaRoute(w http.ResponseWriter, r *http.Request, bytecode } ctx.Set("headers", headerMap) + // Add cookies to context + if cookies := r.Cookies(); len(cookies) > 0 { + cookieMap := make(map[string]any, len(cookies)) + for _, cookie := range cookies { + cookieMap[cookie.Name] = cookie.Value + } + ctx.Set("cookies", cookieMap) + } + // Add URL parameters if params.Count > 0 { paramMap := make(map[string]any, params.Count) - for i := 0; i < params.Count; i++ { - paramMap[params.Keys[i]] = params.Values[i] + for i, key := range params.Keys { + paramMap[key] = params.Values[i] } ctx.Set("params", paramMap) } @@ -230,6 +239,11 @@ func writeResponse(w http.ResponseWriter, result any, log *logger.Logger) { w.Header().Set(name, value) } + // Set cookies + for _, cookie := range httpResp.Cookies { + http.SetCookie(w, cookie) + } + // Set status code w.WriteHeader(httpResp.Status) @@ -243,7 +257,7 @@ func writeResponse(w http.ResponseWriter, result any, log *logger.Logger) { switch res := result.(type) { case string: - // String result - plain text + // String result - default to plain text setContentTypeIfMissing(w, contentTypePlain) w.Write([]byte(res)) default: diff --git a/core/runner/cookies.go b/core/runner/cookies.go new file mode 100644 index 0000000..dbca114 --- /dev/null +++ b/core/runner/cookies.go @@ -0,0 +1,172 @@ +package runner + +import ( + "net/http" + "time" + + luajit "git.sharkk.net/Sky/LuaJIT-to-Go" +) + +// extractCookie grabs cookies from the Lua state +func extractCookie(state *luajit.State) *http.Cookie { + cookie := &http.Cookie{ + Path: "/", // Default path + } + + // Get name + state.GetField(-1, "name") + if !state.IsString(-1) { + state.Pop(1) + return nil // Name is required + } + cookie.Name = state.ToString(-1) + state.Pop(1) + + // Get value + state.GetField(-1, "value") + if state.IsString(-1) { + cookie.Value = state.ToString(-1) + } + state.Pop(1) + + // Get path + state.GetField(-1, "path") + if state.IsString(-1) { + cookie.Path = state.ToString(-1) + } + state.Pop(1) + + // Get domain + state.GetField(-1, "domain") + if state.IsString(-1) { + cookie.Domain = state.ToString(-1) + } + state.Pop(1) + + // Get expires + state.GetField(-1, "expires") + if state.IsNumber(-1) { + expiry := int64(state.ToNumber(-1)) + cookie.Expires = time.Unix(expiry, 0) + } + state.Pop(1) + + // Get max age + state.GetField(-1, "max_age") + if state.IsNumber(-1) { + cookie.MaxAge = int(state.ToNumber(-1)) + } + state.Pop(1) + + // Get secure + state.GetField(-1, "secure") + if state.IsBoolean(-1) { + cookie.Secure = state.ToBoolean(-1) + } + state.Pop(1) + + // Get http only + state.GetField(-1, "http_only") + if state.IsBoolean(-1) { + cookie.HttpOnly = state.ToBoolean(-1) + } + state.Pop(1) + + return cookie +} + +// LuaCookieModule provides cookie functionality to Lua scripts +const LuaCookieModule = ` +-- Cookie module implementation +local cookie = { + -- Set a cookie + set = function(name, value, expires, path, domain, secure, http_only) + if type(name) ~= "string" then + error("cookie.set: name must be a string", 2) + end + + -- Get or create response + local resp = __http_responses[1] or {} + resp.cookies = resp.cookies or {} + __http_responses[1] = resp + + -- Create cookie table + local cookie = { + name = name, + value = value or "", + path = path or "/", + domain = domain + } + + -- Handle expiry + if expires then + if type(expires) == "number" then + if expires > 0 then + -- Add seconds to current time + cookie.max_age = expires + local now = os.time() + cookie.expires = now + expires + elseif expires < 0 then + -- Session cookie (default) + else + -- Expire immediately + cookie.expires = 0 + cookie.max_age = 0 + end + end + end + + -- Set flags + cookie.secure = secure or false + cookie.http_only = http_only or false + + -- Store in cookies table + local n = #resp.cookies + 1 + resp.cookies[n] = cookie + + return true + end, + + -- Get a cookie value + get = function(name) + if type(name) ~= "string" then + error("cookie.get: name must be a string", 2) + end + + -- Access values directly from current environment + local env = getfenv(1) + + -- Check if context exists and has cookies + if env.ctx and env.ctx.cookies and env.ctx.cookies[name] then + return tostring(env.ctx.cookies[name]) + end + + return nil + end, + + -- Remove a cookie + remove = function(name, path, domain) + if type(name) ~= "string" then + error("cookie.remove: name must be a string", 2) + end + + -- Create an expired cookie + return cookie.set(name, "", 0, path or "/", domain, false, false) + end +} + +-- Install cookie module +_G.cookie = cookie + +-- Make sure the cookie module is accessible in sandbox +if __env_system and __env_system.base_env then + __env_system.base_env.cookie = cookie +end +` + +// CookieModuleInitFunc returns an initializer for the cookie module +func CookieModuleInitFunc() StateInitFunc { + return func(state *luajit.State) error { + return state.DoString(LuaCookieModule) + } +} diff --git a/core/runner/http.go b/core/runner/http.go index 4f25242..40cd406 100644 --- a/core/runner/http.go +++ b/core/runner/http.go @@ -1,6 +1,8 @@ package runner import ( + "net/http" + luajit "git.sharkk.net/Sky/LuaJIT-to-Go" ) @@ -9,6 +11,7 @@ type HTTPResponse struct { Status int `json:"status"` Headers map[string]string `json:"headers"` Body any `json:"body"` + Cookies []*http.Cookie `json:"-"` } // NewHTTPResponse creates a default HTTP response @@ -16,6 +19,7 @@ func NewHTTPResponse() *HTTPResponse { return &HTTPResponse{ Status: 200, Headers: make(map[string]string), + Cookies: []*http.Cookie{}, } } @@ -128,6 +132,26 @@ func GetHTTPResponse(state *luajit.State) (*HTTPResponse, bool) { } state.Pop(1) + // Get cookies + state.GetField(-1, "cookies") + if state.IsTable(-1) { + // Iterate through cookies array + length := state.GetTableLength(-1) + for i := 1; i <= length; i++ { + state.PushNumber(float64(i)) + state.GetTable(-2) + + if state.IsTable(-1) { + cookie := extractCookie(state) + if cookie != nil { + response.Cookies = append(response.Cookies, cookie) + } + } + state.Pop(1) + } + } + state.Pop(1) + // Clean up state.Pop(2) // Pop response table and __http_responses diff --git a/core/runner/luarunner.go b/core/runner/luarunner.go index c3764bb..667910f 100644 --- a/core/runner/luarunner.go +++ b/core/runner/luarunner.go @@ -103,14 +103,14 @@ func NewRunner(options ...RunnerOption) (*LuaRunner, error) { return nil, ErrInitFailed } - // Initialize HTTP module BEFORE sandbox setup - httpInit := HTTPModuleInitFunc() - if err := httpInit(state); err != nil { + // Preload core modules + moduleInits := CombineInitFuncs(HTTPModuleInitFunc(), CookieModuleInitFunc()) + if err := moduleInits(state); err != nil { state.Close() return nil, ErrInitFailed } - // Set up sandbox AFTER HTTP module is initialized + // Set up sandbox after core modules are initialized if err := runner.sandbox.Setup(state); err != nil { state.Close() return nil, ErrInitFailed diff --git a/core/runner/sandbox.go b/core/runner/sandbox.go index c13b464..d1b9800 100644 --- a/core/runner/sandbox.go +++ b/core/runner/sandbox.go @@ -81,8 +81,8 @@ func (s *Sandbox) Setup(state *luajit.State) error { preload = package.preload } - -- Add HTTP module explicitly to the base environment base.http = http + base.cookie = cookie -- Add registered custom modules if __sandbox_modules then @@ -101,38 +101,56 @@ func (s *Sandbox) Setup(state *luajit.State) error { -- Fast environment creation with pre-allocation function __get_sandbox_env(ctx) - local env + local env - -- Try to reuse from pool - if __env_system.pool_size > 0 then - env = table.remove(__env_system.env_pool) - __env_system.pool_size = __env_system.pool_size - 1 + -- Try to reuse from pool + if __env_system.pool_size > 0 then + env = table.remove(__env_system.env_pool) + __env_system.pool_size = __env_system.pool_size - 1 - -- 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({}, { - __index = __env_system.base_env - }) + -- 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({}, { + __index = __env_system.base_env + }) - -- Set context if provided - if ctx then - env.ctx = ctx - end + -- Set context if provided + if ctx then + env.ctx = ctx + end - -- Install the fast require implementation - env.require = function(modname) - return __fast_require(env, modname) - end - end + -- Install the fast require implementation + env.require = function(modname) + return __fast_require(env, modname) + end - -- Store reference to current environment - __last_env = env + -- Install cookie module methods directly into environment + env.cookie = { + get = function(name) + if type(name) ~= "string" then + error("cookie.get: name must be a string", 2) + end - return env + if env.ctx and env.ctx.cookies and env.ctx.cookies[name] then + return tostring(env.ctx.cookies[name]) + end + + return nil + end, + + set = cookie.set, + remove = cookie.remove + } + end + + -- Store reference to current environment + __last_env = env + + return env end -- Return environment to pool for reuse