sarin/internal/script/lua.go

package script

import (
	"fmt"

	lua "github.com/yuin/gopher-lua"
	"go.aykhans.me/sarin/internal/types"
)

// LuaEngine implements the Engine interface using gopher-lua.
type LuaEngine struct {
	state     *lua.LState
	transform *lua.LFunction
	http      httpBridge
}

// NewLuaEngine creates a new Lua script engine with the given script content.
// The script must define a global `transform` function that takes a request table
// and returns the modified request table.
//
// Example Lua script:
//
//	function transform(req)
//	    req.headers["X-Custom"] = {"value"}
//	    return req
//	end
//
// It can return the following errors:
//   - types.ErrScriptTransformMissing
//   - types.ScriptExecutionError
func NewLuaEngine(scriptContent string) (*LuaEngine, error) {
	L := lua.NewState()
	engine := &LuaEngine{state: L}

	// Register the globals before running the script so its functions can use them
	engine.registerHTTP()
	registerLuaJSON(L)

	// Execute the script to define the transform function
	if err := L.DoString(scriptContent); err != nil {
		L.Close()
		return nil, types.NewScriptExecutionError("Lua", err)
	}

	// Get the transform function
	transform := L.GetGlobal("transform")
	if transform.Type() != lua.LTFunction {
		L.Close()
		return nil, types.ErrScriptTransformMissing
	}
	engine.transform = transform.(*lua.LFunction)

	return engine, nil
}

// SetHTTPDoer sets the doer that the script's http.* calls are sent through.
func (e *LuaEngine) SetHTTPDoer(doer HTTPDoer) {
	e.http.doer = doer
}

// Transform executes the Lua transform function with the given request data.
// It can return the following errors:
//   - types.ScriptExecutionError
func (e *LuaEngine) Transform(req *RequestData) error {
	e.http.active = true
	defer func() { e.http.active = false }()

	// Convert RequestData to Lua table
	reqTable := e.requestDataToTable(req)

	// Call transform(req)
	e.state.Push(e.transform)
	e.state.Push(reqTable)
	if err := e.state.PCall(1, 1, nil); err != nil {
		return types.NewScriptExecutionError("Lua", err)
	}

	// Get the result
	result := e.state.Get(-1)
	e.state.Pop(1)

	if result.Type() != lua.LTTable {
		return types.NewScriptExecutionError("Lua", fmt.Errorf("transform function must return a table, got %s", result.Type()))
	}

	// Update RequestData from the returned table
	e.tableToRequestData(result.(*lua.LTable), req)

	return nil
}

// Close releases the Lua state resources.
func (e *LuaEngine) Close() {
	if e.state != nil {
		e.state.Close()
	}
}

// requestDataToTable converts RequestData to a Lua table.
func (e *LuaEngine) requestDataToTable(req *RequestData) *lua.LTable {
	L := e.state
	t := L.NewTable()

	t.RawSetString("method", lua.LString(req.Method))
	t.RawSetString("path", lua.LString(req.Path))
	t.RawSetString("body", lua.LString(req.Body))

	t.RawSetString("headers", stringSliceMapToTable(L, req.Headers))
	t.RawSetString("params", stringSliceMapToTable(L, req.Params))
	t.RawSetString("cookies", stringSliceMapToTable(L, req.Cookies))

	return t
}

// stringSliceMapToTable converts a Go map[string][]string to a Lua table of arrays.
func stringSliceMapToTable(state *lua.LState, m map[string][]string) *lua.LTable {
	t := state.NewTable()
	for k, values := range m {
		arr := state.NewTable()
		for _, v := range values {
			arr.Append(lua.LString(v))
		}
		t.RawSetString(k, arr)
	}
	return t
}

// tableToRequestData updates RequestData from a Lua table.
func (e *LuaEngine) tableToRequestData(t *lua.LTable, req *RequestData) {
	// Method
	if v := t.RawGetString("method"); v.Type() == lua.LTString {
		req.Method = string(v.(lua.LString))
	}

	// Path
	if v := t.RawGetString("path"); v.Type() == lua.LTString {
		req.Path = string(v.(lua.LString))
	}

	// Body
	if v := t.RawGetString("body"); v.Type() == lua.LTString {
		req.Body = string(v.(lua.LString))
	}

	// Headers
	if v := t.RawGetString("headers"); v.Type() == lua.LTTable {
		req.Headers = e.tableToStringSliceMap(v.(*lua.LTable))
	}

	// Params
	if v := t.RawGetString("params"); v.Type() == lua.LTTable {
		req.Params = e.tableToStringSliceMap(v.(*lua.LTable))
	}

	// Cookies
	if v := t.RawGetString("cookies"); v.Type() == lua.LTTable {
		req.Cookies = e.tableToStringSliceMap(v.(*lua.LTable))
	}
}

// tableToStringSliceMap converts a Lua table to a Go map[string][]string.
// Supports both single string values and array values.
func (e *LuaEngine) tableToStringSliceMap(t *lua.LTable) map[string][]string {
	result := make(map[string][]string)
	t.ForEach(func(k, v lua.LValue) {
		if k.Type() != lua.LTString {
			return
		}
		key := string(k.(lua.LString))

		if v.Type() == lua.LTTable {
			// Array of values, keys outside 1..n are skipped
			arr := v.(*lua.LTable)
			length := arr.Len()
			values := make([]string, 0, length)
			for i := 1; i <= length; i++ {
				if text, ok := luaScalarString(arr.RawGetInt(i)); ok {
					values = append(values, text)
				}
			}
			result[key] = values
			return
		}

		if text, ok := luaScalarString(v); ok {
			result[key] = []string{text}
		}
	})
	return result
}

// luaScalarString renders strings, numbers and booleans as text.
func luaScalarString(v lua.LValue) (string, bool) {
	switch v.Type() {
	case lua.LTString, lua.LTNumber, lua.LTBool:
		return v.String(), true
	default:
		return "", false
	}
}