Moonshark/runner/embed.go

package runner

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

	"Moonshark/utils/logger"

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

//go:embed lua/sandbox.lua
var sandboxLuaCode string

//go:embed lua/json.lua
var jsonLuaCode string

//go:embed lua/sqlite.lua
var sqliteLuaCode string

//go:embed lua/fs.lua
var fsLuaCode string

//go:embed lua/util.lua
var utilLuaCode string

//go:embed lua/string.lua
var stringLuaCode string

//go:embed lua/table.lua
var tableLuaCode string

//go:embed lua/crypto.lua
var cryptoLuaCode string

//go:embed lua/time.lua
var timeLuaCode string

//go:embed lua/math.lua
var mathLuaCode string

//go:embed lua/env.lua
var envLuaCode string

// ModuleInfo holds information about an embeddable Lua module
type ModuleInfo struct {
	Name          string                 // Module name
	Code          string                 // Module source code
	Bytecode      atomic.Pointer[[]byte] // Cached bytecode
	Once          sync.Once              // For one-time compilation
	DefinesGlobal bool                   // Whether module defines globals directly
}

var (
	sandbox = ModuleInfo{Name: "sandbox", Code: sandboxLuaCode, DefinesGlobal: true}
	modules = []ModuleInfo{
		{Name: "json", Code: jsonLuaCode, DefinesGlobal: true},
		{Name: "sqlite", Code: sqliteLuaCode},
		{Name: "fs", Code: fsLuaCode, DefinesGlobal: true},
		{Name: "util", Code: utilLuaCode, DefinesGlobal: true},
		{Name: "string", Code: stringLuaCode},
		{Name: "table", Code: tableLuaCode},
		{Name: "crypto", Code: cryptoLuaCode, DefinesGlobal: true},
		{Name: "time", Code: timeLuaCode},
		{Name: "math", Code: mathLuaCode},
		{Name: "env", Code: envLuaCode, DefinesGlobal: true},
	}
)

// precompileModule compiles a module's code to bytecode once
func precompileModule(m *ModuleInfo) {
	m.Once.Do(func() {
		tempState := luajit.New(true)
		if tempState == nil {
			logger.Fatalf("Failed to create temp Lua state for %s module compilation", m.Name)
			return
		}
		defer tempState.Close()
		defer tempState.Cleanup()

		code, err := tempState.CompileBytecode(m.Code, m.Name+".lua")
		if err != nil {
			logger.Errorf("Failed to compile %s module: %v", m.Name, err)
			return
		}

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

		logger.Debugf("Successfully precompiled %s.lua to bytecode (%d bytes)", m.Name, len(code))
	})
}

// loadModule loads a module into a Lua state
func loadModule(state *luajit.State, m *ModuleInfo, verbose bool) error {
	// Ensure bytecode is compiled
	precompileModule(m)

	// Attempt to load from bytecode
	bytecode := m.Bytecode.Load()
	if bytecode != nil && len(*bytecode) > 0 {
		if verbose {
			logger.Debugf("Loading %s.lua from precompiled bytecode", m.Name)
		}

		if m.DefinesGlobal {
			// Module defines its own globals, just run it
			if err := state.LoadAndRunBytecode(*bytecode, m.Name+".lua"); err != nil {
				return err
			}
		} else {
			// Module returns a table, capture and set as global
			if err := state.LoadAndRunBytecodeWithResults(*bytecode, m.Name+".lua", 1); err != nil {
				return err
			}
			state.SetGlobal(m.Name)
		}
	} else {
		// Fallback to interpreting the source
		if verbose {
			logger.Warnf("Using non-precompiled %s.lua", m.Name)
		}

		if m.DefinesGlobal {
			if err := state.DoString(m.Code); err != nil {
				return err
			}
		} else {
			if err := state.LoadString(m.Code); err != nil {
				return err
			}
			if err := state.Call(0, 1); err != nil {
				return err
			}
			state.SetGlobal(m.Name)
		}
	}

	return nil
}

// loadSandboxIntoState loads all modules and sandbox into a Lua state
func loadSandboxIntoState(state *luajit.State, verbose bool) error {
	// Load all utility modules first
	for i := range modules {
		if err := loadModule(state, &modules[i], verbose); err != nil {
			return err
		}
	}

	// Initialize any module-specific globals (like SQLite tracking)
	if err := state.DoString(`__active_sqlite_connections = {}`); err != nil {
		return err
	}

	// Load the sandbox last - it defines __execute and other core functions
	precompileModule(&sandbox)
	bytecode := sandbox.Bytecode.Load()
	if bytecode != nil && len(*bytecode) > 0 {
		if verbose {
			logger.Debugf("Loading sandbox.lua from precompiled bytecode")
		}
		return state.LoadAndRunBytecode(*bytecode, "sandbox.lua")
	}

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