Repository: inbucket/inbucket
Pull request #321 (2023-01-24 17:58:45)

lua: Init with config and pool
Signed-off-by: James Hillyerd <james@hillyerd.com>



Patch:
diff --git a/doc/config.md b/doc/config.md
index ae4c10c..6495f95 100644
--- a/doc/config.md
+++ b/doc/config.md
@@ -9,6 +9,7 @@ variables it supports:
 
     KEY                                 DEFAULT             DESCRIPTION
     INBUCKET_LOGLEVEL                   info                debug, info, warn, or error
+    INBUCKET_LUA_SCRIPT                 inbucket.lua        Lua script path
     INBUCKET_MAILBOXNAMING              local               Use local or full addressing
     INBUCKET_SMTP_ADDR                  0.0.0.0:2500        SMTP server IP4 host:port
     INBUCKET_SMTP_DOMAIN                inbucket            HELO domain
@@ -56,6 +57,16 @@ off with `warn` or `error`.
 - Default: `info`
 - Values: one of `debug`, `info`, `warn`, or `error`
 
+### Lua Script
+
+`INBUCKET_LUA_SCRIPT`
+
+This is the path to the (optional) Inbucket Lua script.  If the specified file
+is present, Inbucket will load it during startup.  Ignored if the file is not
+found, or the setting is empty.
+
+- Default: `inbucket.lua`
+
 ### Mailbox Naming
 
 `INBUCKET_MAILBOXNAMING`
diff --git a/go.mod b/go.mod
index e516f87..510970f 100644
--- a/go.mod
+++ b/go.mod
@@ -11,6 +11,7 @@ require (
 	github.com/microcosm-cc/bluemonday v1.0.21
 	github.com/rs/zerolog v1.28.0
 	github.com/stretchr/testify v1.8.1
+	github.com/yuin/gopher-lua v1.0.0
 	golang.org/x/net v0.5.0
 )
 
diff --git a/go.sum b/go.sum
index ead332a..9da5e42 100644
--- a/go.sum
+++ b/go.sum
@@ -77,6 +77,8 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/yuin/gopher-lua v1.0.0 h1:pQCf0LN67Kf7M5u7vRd40A8M1I8IMLrxlqngUJgZ0Ow=
+github.com/yuin/gopher-lua v1.0.0/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw=
 golang.org/x/net v0.0.0-20210501142056-aec3718b3fa0/go.mod h1:OJAsFXCWl8Ukc7SiCT/9KSuxbyM7479/AVlXFRxuMCk=
 golang.org/x/net v0.5.0 h1:GyT4nK/YDHSqa1c4753ouYCDajOYKTja9Xb/OHtgvSw=
 golang.org/x/net v0.5.0/go.mod h1:DivGGAXEgPSlEBzxGzZI+ZLohi+xUj054jfeKui00ws=
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 0747a33..449ba88 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -58,7 +58,8 @@ func (n *mbNaming) Decode(v string) error {
 
 // Root contains global configuration, and structs with for specific sub-systems.
 type Root struct {
-	LogLevel      string   `required:"true" default:"info" desc:"debug, info, warn, or error"`
+	LogLevel      string `required:"true" default:"info" desc:"debug, info, warn, or error"`
+	Lua           Lua
 	MailboxNaming mbNaming `required:"true" default:"local" desc:"Use local, full or domain addressing"`
 	SMTP          SMTP
 	POP3          POP3
@@ -66,6 +67,11 @@ type Root struct {
 	Storage       Storage
 }
 
+// Lua contains the Lua extension host configuration.
+type Lua struct {
+	Path string `required:"false" default:"inbucket.lua" desc:"Lua script path"`
+}
+
 // SMTP contains the SMTP server configuration.
 type SMTP struct {
 	Addr            string        `required:"true" default:"0.0.0.0:2500" desc:"SMTP server IP4 host:port"`
diff --git a/pkg/extension/luahost/lua.go b/pkg/extension/luahost/lua.go
new file mode 100644
index 0000000..346a9c3
--- /dev/null
+++ b/pkg/extension/luahost/lua.go
@@ -0,0 +1,85 @@
+package luahost
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/inbucket/inbucket/pkg/config"
+	"github.com/inbucket/inbucket/pkg/extension"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+	lua "github.com/yuin/gopher-lua"
+	"github.com/yuin/gopher-lua/parse"
+)
+
+// Host of Lua extensions.
+type Host struct {
+	Functions  []string // Functions detected in lua script.
+	extHost    *extension.Host
+	pool       *statePool
+	logContext zerolog.Context
+}
+
+// New constructs a new Lua Host, pre-compiling the source.
+func New(conf config.Lua, extHost *extension.Host) (*Host, error) {
+	scriptPath := conf.Path
+	if scriptPath == "" {
+		return nil, nil
+	}
+
+	logContext := log.With().Str("module", "lua")
+	logger := logContext.Str("phase", "startup").Str("path", scriptPath).Logger()
+
+	// Pre-load, parse, and compile script.
+	if fi, err := os.Stat(scriptPath); err != nil {
+		logger.Info().Msg("Script file not found")
+		return nil, nil
+	} else if fi.IsDir() {
+		return nil, fmt.Errorf("Lua script %v is a directory", scriptPath)
+	}
+
+	logger.Info().Msg("Loading script")
+	file, err := os.Open(scriptPath)
+	defer file.Close()
+	if err != nil {
+		return nil, err
+	}
+
+	return NewFromReader(extHost, bufio.NewReader(file), scriptPath)
+}
+
+// NewFromReader constructs a new Lua Host, loading Lua source from the provided reader.
+// The provided path is used in logging and error messages.
+func NewFromReader(extHost *extension.Host, r io.Reader, path string) (*Host, error) {
+	logContext := log.With().Str("module", "lua")
+
+	// Pre-parse, and compile script.
+	chunk, err := parse.Parse(r, path)
+	if err != nil {
+		return nil, err
+	}
+	proto, err := lua.Compile(chunk, path)
+	if err != nil {
+		return nil, err
+	}
+
+	// Build the pool and confirm LState is retrievable.
+	pool := newStatePool(proto)
+	h := &Host{extHost: extHost, pool: pool, logContext: logContext}
+	if ls, err := pool.getState(); err == nil {
+		// State creation works, put it back.
+		pool.putState(ls)
+	} else {
+		return nil, err
+	}
+
+	return h, nil
+}
+
+// CreateChannel creates a channel and places it into the named global variable
+// in newly created LStates.
+func (h *Host) CreateChannel(name string) chan lua.LValue {
+	return h.pool.createChannel(name)
+}
diff --git a/pkg/extension/luahost/pool.go b/pkg/extension/luahost/pool.go
new file mode 100644
index 0000000..f6f5d8d
--- /dev/null
+++ b/pkg/extension/luahost/pool.go
@@ -0,0 +1,91 @@
+package luahost
+
+import (
+	"sync"
+
+	lua "github.com/yuin/gopher-lua"
+)
+
+type statePool struct {
+	sync.Mutex
+	funcProto *lua.FunctionProto         // Compiled lua.
+	states    []*lua.LState              // Pool of available LStates.
+	channels  map[string]chan lua.LValue // Global interop channels.
+}
+
+func newStatePool(funcProto *lua.FunctionProto) *statePool {
+	return &statePool{
+		funcProto: funcProto,
+		channels:  make(map[string]chan lua.LValue),
+	}
+}
+
+// newState creates a new LState and configures it. Lock must be held.
+func (lp *statePool) newState() (*lua.LState, error) {
+	ls := lua.NewState()
+
+	// Setup channels.
+	for name, ch := range lp.channels {
+		ls.SetGlobal(name, lua.LChannel(ch))
+	}
+
+	// Run compiled script.
+	ls.Push(ls.NewFunctionFromProto(lp.funcProto))
+	if err := ls.PCall(0, lua.MultRet, nil); err != nil {
+		return nil, err
+	}
+
+	return ls, nil
+}
+
+// getState returns a free LState, or creates a new one.
+func (lp *statePool) getState() (*lua.LState, error) {
+	lp.Lock()
+	defer lp.Unlock()
+
+	ln := len(lp.states)
+	if ln == 0 {
+		return lp.newState()
+	}
+
+	state := lp.states[ln-1]
+	lp.states = lp.states[0 : ln-1]
+
+	return state, nil
+}
+
+// putState returns the LState to the pool.
+func (lp *statePool) putState(state *lua.LState) {
+	if state.IsClosed() {
+		return
+	}
+
+	// Clear stack.
+	state.Pop(state.GetTop())
+
+	lp.Lock()
+	defer lp.Unlock()
+
+	lp.states = append(lp.states, state)
+}
+
+// createChannel creates a new channel, which will become a global variable in
+// newly created LStates.  We also destroy any pooled states.
+//
+// Warning: There may still be checked out LStates that will not have the value
+// set, which could be put back into the pool.
+func (lp *statePool) createChannel(name string) chan lua.LValue {
+	lp.Lock()
+	defer lp.Unlock()
+
+	ch := make(chan lua.LValue, 10)
+	lp.channels[name] = ch
+
+	// Flush state pool.
+	for _, s := range lp.states {
+		s.Close()
+	}
+	lp.states = lp.states[:0]
+
+	return ch
+}
diff --git a/pkg/server/lifecycle.go b/pkg/server/lifecycle.go
index d4d4f3d..757b397 100644
--- a/pkg/server/lifecycle.go
+++ b/pkg/server/lifecycle.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/inbucket/inbucket/pkg/config"
 	"github.com/inbucket/inbucket/pkg/extension"
+	"github.com/inbucket/inbucket/pkg/extension/luahost"
 	"github.com/inbucket/inbucket/pkg/message"
 	"github.com/inbucket/inbucket/pkg/msghub"
 	"github.com/inbucket/inbucket/pkg/policy"
@@ -26,6 +27,7 @@ type Services struct {
 	SMTPServer       *smtp.Server
 	WebServer        *web.Server
 	ExtHost          *extension.Host
+	LuaHost          *luahost.Host
 	notify           chan error      // Combined notification for failed services.
 	ready            *sync.WaitGroup // Tracks services that have not reported ready.
 }
@@ -34,6 +36,10 @@ type Services struct {
 func FullAssembly(conf *config.Root) (*Services, error) {
 	// Configure extensions.
 	extHost := extension.NewHost()
+	luaHost, err := luahost.New(conf.Lua, extHost)
+	if err != nil {
+		return nil, err
+	}
 
 	// Configure storage.
 	store, err := storage.FromConfig(conf.Storage)
@@ -65,6 +71,7 @@ func FullAssembly(conf *config.Root) (*Services, error) {
 		SMTPServer:       smtpServer,
 		WebServer:        webServer,
 		ExtHost:          extHost,
+		LuaHost:          luaHost,
 		ready:            &sync.WaitGroup{},
 	}, nil
 }


Test patch:
diff --git a/pkg/extension/luahost/lua_test.go b/pkg/extension/luahost/lua_test.go
new file mode 100644
index 0000000..4cac14f
--- /dev/null
+++ b/pkg/extension/luahost/lua_test.go
@@ -0,0 +1,18 @@
+package luahost_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/inbucket/inbucket/pkg/extension"
+	"github.com/inbucket/inbucket/pkg/extension/luahost"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEmptyScript(t *testing.T) {
+	script := ""
+	extHost := extension.NewHost()
+
+	_, err := luahost.NewFromReader(extHost, strings.NewReader(script), "test.lua")
+	require.NoError(t, err)
+}
diff --git a/pkg/extension/luahost/pool_test.go b/pkg/extension/luahost/pool_test.go
new file mode 100644
index 0000000..b033e52
--- /dev/null
+++ b/pkg/extension/luahost/pool_test.go
@@ -0,0 +1,101 @@
+package luahost
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	lua "github.com/yuin/gopher-lua"
+	"github.com/yuin/gopher-lua/parse"
+)
+
+func makeEmptyPool() *statePool {
+	source := strings.NewReader("-- Empty source")
+
+	chunk, err := parse.Parse(source, "from string")
+	if err != nil {
+		panic(err)
+	}
+
+	proto, err := lua.Compile(chunk, "from string")
+	if err != nil {
+		panic(err)
+	}
+
+	return newStatePool(proto)
+}
+
+func TestPoolGetsDistinct(t *testing.T) {
+	pool := makeEmptyPool()
+
+	a, err := pool.getState()
+	require.NoError(t, err)
+	b, err := pool.getState()
+	require.NoError(t, err)
+
+	if a == b {
+		t.Error("Got pool a == b, expected distinct pools")
+	}
+}
+
+func TestPoolGrowsWithPuts(t *testing.T) {
+	pool := makeEmptyPool()
+
+	a, err := pool.getState()
+	require.NoError(t, err)
+	b, err := pool.getState()
+	require.NoError(t, err)
+	assert.Equal(t, 0, len(pool.states), "Wanted pool to be empty")
+
+	pool.putState(a)
+	pool.putState(b)
+
+	want := 2
+	if got := len(pool.states); got != want {
+		t.Errorf("len pool.states got %v, want %v", got, want)
+	}
+}
+
+// Closed LStates should not be added to the pool.
+func TestPoolPutDiscardsClosed(t *testing.T) {
+	pool := makeEmptyPool()
+
+	a, err := pool.getState()
+	require.NoError(t, err)
+	assert.Equal(t, 0, len(pool.states), "Wanted pool to be empty")
+
+	a.Close()
+	pool.putState(a)
+	assert.Equal(t, 0, len(pool.states), "Wanted pool to remain empty")
+}
+
+func TestPoolPutClearsStack(t *testing.T) {
+	pool := makeEmptyPool()
+
+	ls, err := pool.getState()
+	require.NoError(t, err)
+	assert.Equal(t, 0, len(pool.states), "Wanted pool to be empty")
+
+	// Setup stack.
+	ls.Push(lua.LNumber(4))
+	ls.Push(lua.LString("bacon"))
+	require.Equal(t, 2, ls.GetTop(), "Want stack to have two items")
+
+	// Return and verify stack cleared.
+	pool.putState(ls)
+	assert.Equal(t, 1, len(pool.states), "Wanted pool to have one item")
+	require.Equal(t, 0, ls.GetTop(), "Want stack to be empty")
+}
+
+func TestPoolSetsChannels(t *testing.T) {
+	pool := makeEmptyPool()
+	pool.createChannel("test_chan")
+
+	s, err := pool.getState()
+	require.NoError(t, err)
+
+	got := s.GetGlobal("test_chan")
+	assert.Equal(t, lua.LTChannel, got.Type(),
+		"Got global type %v, wanted LTChannel", got.Type().String())
+}