From 0a477fd288d064d265395c88dbd0f261b723804e Mon Sep 17 00:00:00 2001 From: spiperac Date: Mon, 15 Jun 2026 06:42:13 +0200 Subject: Initial release --- internal/plugin/engine.go | 347 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 347 insertions(+) create mode 100644 internal/plugin/engine.go (limited to 'internal/plugin/engine.go') diff --git a/internal/plugin/engine.go b/internal/plugin/engine.go new file mode 100644 index 0000000..0623d36 --- /dev/null +++ b/internal/plugin/engine.go @@ -0,0 +1,347 @@ +package plugin + +import ( + "fmt" + "io" + "net" + "net/http" + "os" + "path/filepath" + "strings" + + lua "github.com/yuin/gopher-lua" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" + "github.com/shiv/internal/store" +) + +// Engine loads all Lua plugins from a directory and implements +// events.RequestMiddleware, events.ResponseObserver, +// events.WebSocketFrameObserver, and events.PluginEnabledObserver +// so it can be registered directly on the bus. +type Engine struct { + plugins []*Plugin + bus *events.Bus + dir string + st *store.Store +} + +// NewEngine loads all *.lua files from dir, registers each in the store, +// restores persisted enabled state, and returns a ready Engine. +// Plugins that fail to load are skipped. +func NewEngine(dir string, st *store.Store, bus *events.Bus) (*Engine, error) { + plugins, err := loadDir(dir, st) + if err != nil { + return nil, err + } + + e := &Engine{plugins: plugins, bus: bus, dir: dir, st: st} + + for _, p := range plugins { + // Persist plugin record; preserves existing enabled state on conflict. + if err := st.RegisterPlugin(p.name, p.path); err != nil { + logger.Error("plugin: register %s: %v", p.name, err) + } + // Restore enabled state from previous session. + if enabled, err := st.PluginEnabled(p.name); err == nil { + p.enabled.Store(enabled) + } + // Drain any log() calls made during on_load so they reach the UI. + e.emitLogs(p) + } + + active := make([]string, 0, len(plugins)) + for _, p := range plugins { + active = append(active, p.name) + } + if err := st.ScanPlugins(active); err != nil { + logger.Error("plugin: scan: %v", err) + } + + return e, nil +} + +// emitLogs drains a plugin's log buffer and emits each line onto the bus. +// Called after every hook invocation, including on error, so no log line +// is ever lost. +func (e *Engine) emitLogs(p *Plugin) { + for _, line := range p.drainLogs() { + e.bus.EmitPluginLog(events.PluginLogEvent{Name: p.name, Message: line}) + } +} + +// HandleRequest implements events.RequestMiddleware. +// Each plugin's on_request(req) hook is called in load order. Table +// construction and the Lua call happen in a single locked section via +// callWithBuilder — no VM access escapes the lock at any point. +// A plugin may modify method, url, headers, body, or set drop=true. +// The first drop short-circuits all remaining plugins. +func (e *Engine) HandleRequest(ev events.RequestEvent) events.RequestResult { + req := ev.Request + body := ev.Body + + if req == nil || req.URL == nil { + return events.RequestResult{Drop: false, Request: req, Body: body} + } + + for _, p := range e.plugins { + if !p.enabled.Load() { + continue + } + if !p.has("on_request") { + continue + } + + // Snapshot Go values before entering the lock so the builder closure + // captures plain Go strings, not live request fields that could race. + method := req.Method + rawURL := req.URL.String() + bodySnap := string(body) + headers := req.Header + + ret, err := p.callWithBuilder("on_request", func(L *lua.LState) *lua.LTable { + tbl := L.NewTable() + L.SetField(tbl, "method", lua.LString(method)) + L.SetField(tbl, "url", lua.LString(rawURL)) + L.SetField(tbl, "body", lua.LString(bodySnap)) + L.SetField(tbl, "drop", lua.LFalse) + + hTbl := L.NewTable() + for k, vals := range headers { + L.SetField(hTbl, k, lua.LString(strings.Join(vals, ", "))) + } + L.SetField(tbl, "headers", hTbl) + return tbl + }) + e.emitLogs(p) + if err != nil { + logger.Error("plugin %s: on_request: %v", p.name, err) + continue + } + if ret == nil { + continue + } + + // Read back modifications. ret is a Lua table owned by p.L — + // mu must be held while reading it. + p.mu.Lock() + drop := false + if d, ok := ret.RawGetString("drop").(lua.LBool); ok { + drop = bool(d) + } + if !drop { + if m, ok := ret.RawGetString("method").(lua.LString); ok && string(m) != "" { + req.Method = string(m) + } + if u, ok := ret.RawGetString("url").(lua.LString); ok && string(u) != "" { + if parsed, err := req.URL.Parse(string(u)); err == nil { + req.URL = parsed + } + } + if b, ok := ret.RawGetString("body").(lua.LString); ok { + body = []byte(b) + } + if hTbl, ok := ret.RawGetString("headers").(*lua.LTable); ok { + newHeaders := make(http.Header) + hTbl.ForEach(func(k, v lua.LValue) { + newHeaders.Set(k.String(), v.String()) + }) + req.Header = newHeaders + } + } + p.mu.Unlock() + + if drop { + return events.RequestResult{Drop: true, Request: req, Body: body} + } + } + + return events.RequestResult{Drop: false, Request: req, Body: body} +} + +// ObserveResponse implements events.ResponseObserver. +// Each plugin's on_response(resp) hook is called in load order. Table +// construction and the Lua call happen in a single locked section via +// callWithBuilder — no VM access escapes the lock at any point. +// The hook is read-only — return value is ignored. +func (e *Engine) ObserveResponse(ev events.ResponseEvent) { + // Snapshot all event fields once before the plugin loop. Each plugin + // call captures these plain Go values in its builder closure — no live + // event fields are accessed inside the lock. + host := ev.Host + // ev.Host arrives as host:port from the proxy; expose bare hostname to plugins. + if h, _, err := net.SplitHostPort(ev.Host); err == nil { + host = h + } + method := ev.Method + url := ev.URL + proto := ev.Proto + status := ev.StatusCode + durationMs := ev.DurationMs + tls := ev.TLS + reqBody := string(ev.ReqBody) + respBody := string(ev.RespBody) + reqHeaders := ev.ReqHeaders + respHeaders := ev.RespHeaders + + for _, p := range e.plugins { + if !p.enabled.Load() { + continue + } + if !p.has("on_response") { + continue + } + + _, err := p.callWithBuilder("on_response", func(L *lua.LState) *lua.LTable { + tbl := L.NewTable() + L.SetField(tbl, "host", lua.LString(host)) + L.SetField(tbl, "method", lua.LString(method)) + L.SetField(tbl, "url", lua.LString(url)) + L.SetField(tbl, "proto", lua.LString(proto)) + L.SetField(tbl, "status", lua.LNumber(status)) + L.SetField(tbl, "duration_ms", lua.LNumber(durationMs)) + L.SetField(tbl, "tls", lua.LBool(tls)) + L.SetField(tbl, "req_body", lua.LString(reqBody)) + L.SetField(tbl, "resp_body", lua.LString(respBody)) + + rqH := L.NewTable() + for k, vals := range reqHeaders { + L.SetField(rqH, k, lua.LString(strings.Join(vals, ", "))) + } + L.SetField(tbl, "req_headers", rqH) + + rpH := L.NewTable() + for k, vals := range respHeaders { + L.SetField(rpH, k, lua.LString(strings.Join(vals, ", "))) + } + L.SetField(tbl, "resp_headers", rpH) + return tbl + }) + e.emitLogs(p) + if err != nil { + logger.Error("plugin %s: on_response: %v", p.name, err) + } + } +} + +// ObserveWebSocketFrame implements events.WebSocketFrameObserver. +// Each plugin's on_websocket_frame(frame) hook is called in load order. +// The hook may modify frame.payload — the last modification wins. +// Plugins cannot drop frames. +func (e *Engine) ObserveWebSocketFrame(ev events.WebSocketFrameEvent) events.WebSocketFrameResult { + payload := ev.Payload + + // Snapshot Go values before the plugin loop. + connID := ev.ConnectionID + direction := int(ev.Direction) + opcode := int(ev.Opcode) + payloadSnap := string(ev.Payload) + + for _, p := range e.plugins { + if !p.enabled.Load() { + continue + } + if !p.has("on_websocket_frame") { + continue + } + + ret, err := p.callWithBuilder("on_websocket_frame", func(L *lua.LState) *lua.LTable { + tbl := L.NewTable() + L.SetField(tbl, "connection_id", lua.LNumber(connID)) + L.SetField(tbl, "direction", lua.LNumber(direction)) + L.SetField(tbl, "opcode", lua.LNumber(opcode)) + L.SetField(tbl, "payload", lua.LString(payloadSnap)) + return tbl + }) + e.emitLogs(p) + if err != nil { + logger.Error("plugin %s: on_websocket_frame: %v", p.name, err) + continue + } + if ret == nil { + continue + } + + p.mu.Lock() + if b, ok := ret.RawGetString("payload").(lua.LString); ok { + payload = []byte(b) + payloadSnap = string(payload) + } + p.mu.Unlock() + } + + return events.WebSocketFrameResult{Payload: payload} +} + +// ObservePluginEnabled implements events.PluginEnabledObserver. +// Toggles the in-memory enabled flag for the named plugin so the change +// takes effect immediately for all subsequent hook calls. +func (e *Engine) ObservePluginEnabled(ev events.SetPluginEnabledEvent) { + for _, p := range e.plugins { + if p.name == ev.Name { + p.enabled.Store(ev.Enabled) + return + } + } +} + +// ObserveLoadPlugin implements events.LoadPluginObserver. +// Copies the selected file into the plugins directory, loads it, registers +// it in the store, and appends it to the engine's plugin slice so it +// participates in all subsequent hook calls. +// Safe without a mutex — the bus is synchronous so this and all hook +// iteration methods are always called sequentially. +func (e *Engine) ObserveLoadPlugin(ev events.LoadPluginEvent) { + filename := filepath.Base(ev.SourcePath) + destPath := filepath.Join(e.dir, filename) + + if err := copyFile(ev.SourcePath, destPath); err != nil { + logger.Error("plugin: import %s: %v", filename, err) + return + } + + p, err := load(destPath, e.st) + if err != nil { + logger.Error("plugin: load %s: %v", filename, err) + return + } + + if err := e.st.RegisterPlugin(p.name, p.path); err != nil { + logger.Error("plugin: register %s: %v", p.name, err) + } + + e.emitLogs(p) + e.plugins = append(e.plugins, p) +} + +// copyFile copies src to dst, creating the destination directory if needed. +func copyFile(src, dst string) error { + in, err := os.Open(src) + if err != nil { + return fmt.Errorf("open: %w", err) + } + defer in.Close() + + if err := os.MkdirAll(filepath.Dir(dst), 0755); err != nil { + return fmt.Errorf("mkdir: %w", err) + } + + out, err := os.Create(dst) + if err != nil { + return fmt.Errorf("create: %w", err) + } + defer out.Close() + + if _, err := io.Copy(out, in); err != nil { + return fmt.Errorf("copy: %w", err) + } + return out.Sync() +} + +// Close shuts down all plugin VMs. +func (e *Engine) Close() { + for _, p := range e.plugins { + p.close() + } +} -- cgit v1.2.3