Moonshark/core/runner/embed.go

package runner

import (
	_ "embed"
	"sync"
	"sync/atomic"

	"Moonshark/core/utils/logger"

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

//go:embed sandbox.lua
var sandboxLuaCode string

//go:embed json.lua
var jsonLuaCode string

//go:embed sqlite.lua
var sqliteLuaCode string

// Global bytecode cache to improve performance
var (
	sandboxBytecode    atomic.Pointer[[]byte]
	jsonBytecode       atomic.Pointer[[]byte]
	sqliteBytecode     atomic.Pointer[[]byte]
	bytecodeOnce       sync.Once
	jsonBytecodeOnce   sync.Once
	sqliteBytecodeOnce sync.Once
)

// precompileSandboxCode compiles the sandbox.lua code to bytecode once
func precompileSandboxCode() {
	tempState := luajit.New()
	if tempState == nil {
		logger.Fatal("Failed to create temp Lua state for bytecode compilation")
	}
	defer tempState.Close()
	defer tempState.Cleanup()

	code, err := tempState.CompileBytecode(sandboxLuaCode, "sandbox.lua")
	if err != nil {
		logger.Error("Failed to compile sandbox code: %v", err)
		return
	}

	bytecode := make([]byte, len(code))
	copy(bytecode, code)
	sandboxBytecode.Store(&bytecode)

	logger.Debug("Successfully precompiled sandbox.lua to bytecode (%d bytes)", len(code))
}

// precompileJsonModule compiles the json.lua code to bytecode once
func precompileJsonModule() {
	tempState := luajit.New()
	if tempState == nil {
		logger.Fatal("Failed to create temp Lua state for JSON module compilation")
	}
	defer tempState.Close()
	defer tempState.Cleanup()

	code, err := tempState.CompileBytecode(jsonLuaCode, "json.lua")
	if err != nil {
		logger.Error("Failed to compile JSON module: %v", err)
		return
	}

	bytecode := make([]byte, len(code))
	copy(bytecode, code)
	jsonBytecode.Store(&bytecode)

	logger.Debug("Successfully precompiled json.lua to bytecode (%d bytes)", len(code))
}

// precompileSqliteModule compiles the sqlite.lua code to bytecode once
func precompileSqliteModule() {
	tempState := luajit.New()
	if tempState == nil {
		logger.Fatal("Failed to create temp Lua state for SQLite module compilation")
	}
	defer tempState.Close()
	defer tempState.Cleanup()

	code, err := tempState.CompileBytecode(sqliteLuaCode, "sqlite.lua")
	if err != nil {
		logger.Error("Failed to compile SQLite module: %v", err)
		return
	}

	bytecode := make([]byte, len(code))
	copy(bytecode, code)
	sqliteBytecode.Store(&bytecode)

	logger.Debug("Successfully precompiled sqlite.lua to bytecode (%d bytes)", len(code))
}

// loadSandboxIntoState loads the sandbox code into a Lua state
func loadSandboxIntoState(state *luajit.State, verbose bool) error {
	bytecodeOnce.Do(precompileSandboxCode)
	jsonBytecodeOnce.Do(precompileJsonModule)
	sqliteBytecodeOnce.Do(precompileSqliteModule)

	// First load and execute the JSON module
	jsBytecode := jsonBytecode.Load()
	if jsBytecode != nil && len(*jsBytecode) > 0 {
		if verbose {
			logger.Debug("Loading json.lua from precompiled bytecode")
		}

		if err := state.LoadBytecode(*jsBytecode, "json.lua"); err != nil {
			return err
		}

		if err := state.RunBytecodeWithResults(1); err != nil {
			return err
		}

		state.SetGlobal("json")
	} else {
		if verbose {
			logger.Warning("Using non-precompiled json.lua")
		}

		if err := state.DoString(jsonLuaCode); err != nil {
			return err
		}
	}

	// Initialize active connections tracking
	if err := state.DoString(`__active_sqlite_connections = {}`); err != nil {
		return err
	}

	// Load SQLite module
	sqlBytecode := sqliteBytecode.Load()
	if sqlBytecode != nil && len(*sqlBytecode) > 0 {
		if verbose {
			logger.Debug("Loading sqlite.lua from precompiled bytecode")
		}

		if err := state.LoadBytecode(*sqlBytecode, "sqlite.lua"); err != nil {
			return err
		}

		if err := state.RunBytecodeWithResults(1); err != nil {
			return err
		}

		state.SetGlobal("sqlite")
	} else {
		if verbose {
			logger.Warning("Using non-precompiled sqlite.lua")
		}

		if err := state.DoString(sqliteLuaCode); err != nil {
			return err
		}
	}

	bytecode := sandboxBytecode.Load()
	if bytecode != nil && len(*bytecode) > 0 {
		if verbose {
			logger.Debug("Loading sandbox.lua from precompiled bytecode")
		}
		return state.LoadAndRunBytecode(*bytecode, "sandbox.lua")
	}

	if verbose {
		logger.Warning("Using non-precompiled sandbox.lua")
	}
	return state.DoString(sandboxLuaCode)
}