From 4b552308eee8964e249caaf9b5bfa4c3d87cbec8 Mon Sep 17 00:00:00 2001
From: Michael Fero <6863207+mikefero@users.noreply.github.com>
Date: Fri, 16 Dec 2022 08:40:51 -0500
Subject: [PATCH] feat: add the ability to load a Lua filesystem instead of
 from OS

This feature extends the Options to load a Lua filesystem instead of
loading directly from the OS. This extension uses an interface to allow
for user level overrides by requiring Open and Stat functions utilized
by LState.
---
 _state.go      | 24 ++++++++++++++++++++++++
 auxlib.go      | 12 ++++++++----
 auxlib_test.go | 32 ++++++++++++++++++++++++++++++++
 loadlib.go     |  2 +-
 state.go       | 24 ++++++++++++++++++++++++
 5 files changed, 89 insertions(+), 5 deletions(-)

diff --git a/_state.go b/_state.go
index 645589fe..edb75003 100644
--- a/_state.go
+++ b/_state.go
@@ -105,6 +105,8 @@ type Options struct {
 	// If `MinimizeStackMemory` is set, the call stack will be automatically grown or shrank up to a limit of
 	// `CallStackSize` in order to minimize memory usage. This does incur a slight performance penalty.
 	MinimizeStackMemory bool
+	// Load lua files from LuaFileSystem instead of OS file-system.
+	LuaFileSystem LuaFileSystem
 }
 
 /* }}} */
@@ -533,6 +535,28 @@ func (rg *registry) IsFull() bool {
 
 /* }}} */
 
+/* luaFileSystem {{{ */
+type LuaFileSystem interface {
+	Open(path string) (io.ReadCloser, error)
+	Stat(luapath string) (os.FileInfo, error)
+}
+
+func (ls *LState) Open(path string) (io.ReadCloser, error) {
+	if ls.Options.LuaFileSystem != nil {
+		return ls.Options.LuaFileSystem.Open(path)
+	}
+	return os.Open(path)
+}
+
+func (ls *LState) Stat(luapath string) (os.FileInfo, error) {
+	if ls.Options.LuaFileSystem != nil {
+		return ls.Options.LuaFileSystem.Stat(luapath)
+	}
+	return os.Stat(luapath)
+}
+
+/* }}} */
+
 /* Global {{{ */
 
 func newGlobal() *Global {
diff --git a/auxlib.go b/auxlib.go
index 603fdafe..7be0a142 100644
--- a/auxlib.go
+++ b/auxlib.go
@@ -350,16 +350,20 @@ func (ls *LState) CallMeta(obj LValue, event string) LValue {
 /* load and function call operations {{{ */
 
 func (ls *LState) LoadFile(path string) (*LFunction, error) {
-	var file *os.File
-	var err error
+	var file io.Reader
 	if len(path) == 0 {
 		file = os.Stdin
 	} else {
-		file, err = os.Open(path)
-		defer file.Close()
+		readCloser, err := ls.Open(path)
+		defer func() {
+			if readCloser != nil {
+				readCloser.Close()
+			}
+		}()
 		if err != nil {
 			return nil, newApiErrorE(ApiErrorFile, err)
 		}
+		file = readCloser
 	}
 
 	reader := bufio.NewReader(file)
diff --git a/auxlib_test.go b/auxlib_test.go
index a2de37cc..87f14c30 100644
--- a/auxlib_test.go
+++ b/auxlib_test.go
@@ -1,6 +1,8 @@
 package lua
 
 import (
+	"embed"
+	"io"
 	"io/ioutil"
 	"os"
 	"testing"
@@ -335,3 +337,33 @@ func TestLoadFileForEmptyFile(t *testing.T) {
 	_, err = L.LoadFile(tmpFile.Name())
 	errorIfNotNil(t, err)
 }
+
+//go:embed _lua5.1-tests/all.lua
+var luaTree embed.FS
+
+type luaFileSystem struct {
+	fileSystem embed.FS
+}
+
+func (lfs *luaFileSystem) Open(path string) (io.ReadCloser, error) {
+	return lfs.fileSystem.Open(path)
+}
+
+func (lfs *luaFileSystem) Stat(path string) (os.FileInfo, error) {
+	file, err := lfs.fileSystem.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	return file.Stat()
+}
+
+func TestLoadLuaFileSystemFile(t *testing.T) {
+	L := NewState(Options{LuaFileSystem: &luaFileSystem{fileSystem: luaTree}})
+	defer L.Close()
+
+	_, err := L.LoadFile("_lua5.1-tests/all.lua")
+	errorIfNotNil(t, err)
+
+	_, err = L.LoadFile("invalid.lua")
+	errorIfNil(t, err)
+}
diff --git a/loadlib.go b/loadlib.go
index 40ce122b..b853e038 100644
--- a/loadlib.go
+++ b/loadlib.go
@@ -37,7 +37,7 @@ func loFindFile(L *LState, name, pname string) (string, string) {
 	messages := []string{}
 	for _, pattern := range strings.Split(string(path), ";") {
 		luapath := strings.Replace(pattern, "?", name, -1)
-		if _, err := os.Stat(luapath); err == nil {
+		if _, err := L.Stat(luapath); err == nil {
 			return luapath, ""
 		} else {
 			messages = append(messages, err.Error())
diff --git a/state.go b/state.go
index 50e70ed6..4b6901f3 100644
--- a/state.go
+++ b/state.go
@@ -109,6 +109,8 @@ type Options struct {
 	// If `MinimizeStackMemory` is set, the call stack will be automatically grown or shrank up to a limit of
 	// `CallStackSize` in order to minimize memory usage. This does incur a slight performance penalty.
 	MinimizeStackMemory bool
+	// Load lua files from LuaFileSystem instead of OS file-system.
+	LuaFileSystem LuaFileSystem
 }
 
 /* }}} */
@@ -579,6 +581,28 @@ func (rg *registry) IsFull() bool {
 
 /* }}} */
 
+/* luaFileSystem {{{ */
+type LuaFileSystem interface {
+	Open(path string) (io.ReadCloser, error)
+	Stat(luapath string) (os.FileInfo, error)
+}
+
+func (ls *LState) Open(path string) (io.ReadCloser, error) {
+	if ls.Options.LuaFileSystem != nil {
+		return ls.Options.LuaFileSystem.Open(path)
+	}
+	return os.Open(path)
+}
+
+func (ls *LState) Stat(luapath string) (os.FileInfo, error) {
+	if ls.Options.LuaFileSystem != nil {
+		return ls.Options.LuaFileSystem.Stat(luapath)
+	}
+	return os.Stat(luapath)
+}
+
+/* }}} */
+
 /* Global {{{ */
 
 func newGlobal() *Global {