package runner import ( "github.com/valyala/fasthttp" "Moonshark/core/runner/sandbox" "Moonshark/core/sessions" "Moonshark/core/utils/logger" luajit "git.sharkk.net/Sky/LuaJIT-to-Go" ) // SessionHandler handles session management for Lua scripts type SessionHandler struct { manager *sessions.SessionManager debugLog bool } // NewSessionHandler creates a new session handler func NewSessionHandler(manager *sessions.SessionManager) *SessionHandler { return &SessionHandler{ manager: manager, debugLog: false, } } // EnableDebug enables debug logging func (h *SessionHandler) EnableDebug() { h.debugLog = true } // debug logs a message if debug is enabled func (h *SessionHandler) debug(format string, args ...interface{}) { if h.debugLog { logger.Debug("[SessionHandler] "+format, args...) } } // WithSessionManager creates a RunnerOption to add session support func WithSessionManager(manager *sessions.SessionManager) RunnerOption { return func(r *Runner) { handler := NewSessionHandler(manager) // Add hooks to the runner r.AddInitHook(handler.preRequestHook) r.AddFinalizeHook(handler.postRequestHook) } } // preRequestHook is called before executing a request func (h *SessionHandler) preRequestHook(state *luajit.State, ctx *Context) error { h.debug("Running pre-request session hook") // Check if we have cookie information in context // Instead of raw request, we now look for the cookie map if ctx == nil || ctx.Values["_request_cookies"] == nil { return nil } // Extract cookies from context cookies, ok := ctx.Values["_request_cookies"].(map[string]any) if !ok { return nil } // Get the session ID from cookies cookieName := h.manager.CookieOptions()["name"].(string) var sessionID string // Check if our session cookie exists if cookieValue, exists := cookies[cookieName]; exists { if strValue, ok := cookieValue.(string); ok && strValue != "" { sessionID = strValue h.debug("Found existing session ID: %s", sessionID) } } // If no session ID found, create new session if sessionID == "" { // Create a new session session := h.manager.CreateSession() sessionID = session.ID h.debug("Created new session with ID: %s", sessionID) } // Store the session ID in the context for later use ctx.Set("_session_id", sessionID) // Get the session data session := h.manager.GetSession(sessionID) sessionData := session.GetAll() // Set session data in Lua state if err := SetSessionData(state, sessionID, sessionData); err != nil { h.debug("Failed to set session data: %v", err) return err } h.debug("Session data initialized successfully") return nil } // postRequestHook is called after executing a request func (h *SessionHandler) postRequestHook(state *luajit.State, ctx *Context, result any) error { h.debug("Running post-request session hook") // Check if session was modified modifiedID, modifiedData, modified := GetSessionData(state) if !modified { h.debug("Session not modified, skipping") return nil } // Get the original session ID from context var sessionID string if ctx != nil { if id, ok := ctx.Values["_session_id"].(string); ok { sessionID = id } } // Use the original session ID if the modified one is empty if modifiedID == "" { modifiedID = sessionID } if modifiedID == "" { h.debug("No session ID found, cannot persist session data") return nil } h.debug("Persisting modified session data for ID: %s", modifiedID) // Update session in manager session := h.manager.GetSession(modifiedID) session.Clear() // clear to sync deleted values for k, v := range modifiedData { session.Set(k, v) } h.manager.SaveSession(session) // Add session cookie to result if it's an HTTP response if httpResp, ok := result.(*sandbox.HTTPResponse); ok { h.addSessionCookie(httpResp, modifiedID) } h.debug("Session data persisted successfully") return nil } // addSessionCookie adds a session cookie to an HTTP response func (h *SessionHandler) addSessionCookie(resp *sandbox.HTTPResponse, sessionID string) { // Get cookie options opts := h.manager.CookieOptions() // Check if session cookie is already set cookieName := opts["name"].(string) for _, cookie := range resp.Cookies { if string(cookie.Key()) == cookieName { h.debug("Session cookie already set in response") return } } h.debug("Adding session cookie to response") // Create and add cookie cookie := fasthttp.AcquireCookie() cookie.SetKey(cookieName) cookie.SetValue(sessionID) cookie.SetPath(opts["path"].(string)) cookie.SetHTTPOnly(opts["http_only"].(bool)) cookie.SetMaxAge(opts["max_age"].(int)) // Optional cookie parameters if domain, ok := opts["domain"].(string); ok && domain != "" { cookie.SetDomain(domain) } if secure, ok := opts["secure"].(bool); ok { cookie.SetSecure(secure) } resp.Cookies = append(resp.Cookies, cookie) } // 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 }