Moonshark/core/runner/SessionModule.go

package runner

import (
	"Moonshark/core/utils/logger"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
)

// LuaSessionModule provides session functionality to Lua scripts
const LuaSessionModule = `
-- Global table to store session data
__session_data = __session_data or {}
__session_id = __session_id or nil
__session_modified = false

-- Session module implementation
local session = {
    -- Get a session value
    get = function(key)
        if type(key) ~= "string" then
            error("session.get: key must be a string", 2)
        end

        if __session_data and __session_data[key] then
            return __session_data[key]
        end

        return nil
    end,

    -- Set a session value
    set = function(key, value)
        if type(key) ~= "string" then
            error("session.set: key must be a string", 2)
        end

        -- Ensure session data table exists
        __session_data = __session_data or {}

        -- Store value
        __session_data[key] = value

        -- Mark session as modified
        __session_modified = true

        return true
    end,

    -- Delete a session value
    delete = function(key)
        if type(key) ~= "string" then
            error("session.delete: key must be a string", 2)
        end

        if __session_data then
            __session_data[key] = nil
            __session_modified = true
        end

        return true
    end,

    -- Clear all session data
    clear = function()
        __session_data = {}
        __session_modified = true
        return true
    end,

    -- Get the session ID
    get_id = function()
        return __session_id or nil
    end,

    -- Get all session data
    get_all = function()
        local result = {}
        for k, v in pairs(__session_data or {}) do
            result[k] = v
        end
        return result
    end,

    -- Check if session has a key
    has = function(key)
        if type(key) ~= "string" then
            error("session.has: key must be a string", 2)
        end

        return __session_data and __session_data[key] ~= nil
    end
}

-- Install session module
_G.session = session

-- Make sure the session module is accessible in sandbox
if __env_system and __env_system.base_env then
    __env_system.base_env.session = session
end

-- Hook into script execution to preserve session state
local old_execute_script = __execute_script
if old_execute_script then
  __execute_script = function(fn, ctx)
    -- Reset modification flag at the start of request
    __session_modified = false

    -- Execute original function
    return old_execute_script(fn, ctx)
  end
end
`

// GetSessionData extracts session data from Lua state
func GetSessionData(state *luajit.State) (string, map[string]any, bool) {
	// Check if session was modified
	state.GetGlobal("__session_modified")
	modified := state.ToBoolean(-1)
	state.Pop(1)

	if !modified {
		return "", nil, false
	}

	// Get session ID
	state.GetGlobal("__session_id")
	sessionID := state.ToString(-1)
	state.Pop(1)

	// Get session data
	state.GetGlobal("__session_data")
	if !state.IsTable(-1) {
		state.Pop(1)
		return sessionID, nil, false
	}

	data, err := state.ToTable(-1)
	state.Pop(1)

	if err != nil {
		logger.Error("Failed to extract session data: %v", err)
		return sessionID, nil, false
	}

	return sessionID, data, true
}

// SetSessionData sets session data in Lua state
func SetSessionData(state *luajit.State, sessionID string, data map[string]any) error {
	// Set session ID
	state.PushString(sessionID)
	state.SetGlobal("__session_id")

	// Set session data
	if data == nil {
		data = make(map[string]any)
	}

	if err := state.PushTable(data); err != nil {
		return err
	}
	state.SetGlobal("__session_data")

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

	return nil
}

// SessionModuleInitFunc returns an initializer for the session module
func SessionModuleInitFunc() StateInitFunc {
	return func(state *luajit.State) error {
		return state.DoString(LuaSessionModule)
	}
}