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 } }