Moonshark/core/runner/CookieModule.go

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, options, ...)
        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

        -- Handle options as table or legacy params
        local opts = {}
        if type(options) == "table" then
            opts = options
        elseif options ~= nil then
            -- Legacy support: options is actually 'expires'
            opts.expires = options
            -- Check for other legacy params (4th-7th args)
            local args = {...}
            if args[1] then opts.path = args[1] end
            if args[2] then opts.domain = args[2] end
            if args[3] then opts.secure = args[3] end
            if args[4] ~= nil then opts.http_only = args[4] end
        end

        -- Create cookie table
        local cookie = {
            name = name,
            value = value or "",
            path = opts.path or "/",
            domain = opts.domain
        }

        -- Handle expiry
        if opts.expires then
            if type(opts.expires) == "number" then
                if opts.expires > 0 then
                    cookie.max_age = opts.expires
                    local now = os.time()
                    cookie.expires = now + opts.expires
                elseif opts.expires < 0 then
                    cookie.expires = 1
                    cookie.max_age = 0
                else
                    -- opts.expires == 0: Session cookie
                    -- Do nothing (omitting both expires and max-age creates a session cookie)
                end
            end
        end

        -- Security flags
        cookie.secure = (opts.secure ~= false)
        cookie.http_only = (opts.http_only ~= 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(2)

        -- 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, "", {expires = 0, path = path or "/", domain = domain})
    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)
	}
}