Allow passing fs.FS when calling functions (#571)

Fixes #563 

Signed-off-by: knqyf263 <[email protected]>
Co-authored-by: Adrian Cole <[email protected]>

config.go

@@ -3,14 +3,14 @@ package wazero
 import (
 	"context"
 	"errors"
-	"fmt"
 	"io"
 	"io/fs"
 	"math"
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/engine/compiler"
 	"github.com/tetratelabs/wazero/internal/engine/interpreter"
+	fs2 "github.com/tetratelabs/wazero/internal/fs"
 	"github.com/tetratelabs/wazero/internal/wasm"
 )
 
@@ -452,21 +452,15 @@ type moduleConfig struct {
 	// environKeys allow overwriting of existing values.
 	environKeys map[string]int
 
-	// preopenFD has the next FD number to use
-	preopenFD uint32
-	// preopens are keyed on file descriptor and only include the Path and FS fields.
-	preopens map[uint32]*wasm.FileEntry
-	// preopenPaths allow overwriting of existing paths.
-	preopenPaths map[string]uint32
+	fs *fs2.FSConfig
 }
 
 func NewModuleConfig() ModuleConfig {
 	return &moduleConfig{
 		startFunctions: []string{"_start"},
 		environKeys:    map[string]int{},
-		preopenFD:      uint32(3), // after stdin/stdout/stderr
-		preopens:       map[uint32]*wasm.FileEntry{},
-		preopenPaths:   map[string]uint32{},
+
+		fs: fs2.NewFSConfig(),
 	}
 }
 
@@ -493,7 +487,7 @@ func (c *moduleConfig) WithEnv(key, value string) ModuleConfig {
 // WithFS implements ModuleConfig.WithFS
 func (c *moduleConfig) WithFS(fs fs.FS) ModuleConfig {
 	ret := *c // copy
-	ret.setFS("/", fs)
+	ret.fs = ret.fs.WithFS(fs)
 	return &ret
 }
 
@@ -542,23 +536,10 @@ func (c *moduleConfig) WithRandSource(source io.Reader) ModuleConfig {
 // WithWorkDirFS implements ModuleConfig.WithWorkDirFS
 func (c *moduleConfig) WithWorkDirFS(fs fs.FS) ModuleConfig {
 	ret := *c // copy
-	ret.setFS(".", fs)
+	ret.fs = ret.fs.WithWorkDirFS(fs)
 	return &ret
 }
 
-// setFS maps a path to a file-system. This is only used for base paths: "/" and ".".
-func (c *moduleConfig) setFS(path string, fs fs.FS) {
-	// Check to see if this key already exists and update it.
-	entry := &wasm.FileEntry{Path: path, FS: fs}
-	if fd, ok := c.preopenPaths[path]; ok {
-		c.preopens[fd] = entry
-	} else {
-		c.preopens[c.preopenFD] = entry
-		c.preopenPaths[path] = c.preopenFD
-		c.preopenFD++
-	}
-}
-
 // toSysContext creates a baseline wasm.SysContext configured by ModuleConfig.
 func (c *moduleConfig) toSysContext() (sys *wasm.SysContext, err error) {
 	var environ []string // Intentionally doesn't pre-allocate to reduce logic to default to nil.
@@ -578,24 +559,9 @@ func (c *moduleConfig) toSysContext() (sys *wasm.SysContext, err error) {
 		environ = append(environ, key+"="+value)
 	}
 
-	// Ensure no-one set a nil FD. We do this here instead of at the call site to allow chaining as nil is unexpected.
-	rootFD := uint32(0) // zero is invalid
-	setWorkDirFS := false
-	preopens := c.preopens
-	for fd, entry := range preopens {
-		if entry.FS == nil {
-			err = fmt.Errorf("FS for %s is nil", entry.Path)
-			return
-		} else if entry.Path == "/" {
-			rootFD = fd
-		} else if entry.Path == "." {
-			setWorkDirFS = true
-		}
-	}
-
-	// Default the working directory to the root FS if it exists.
-	if rootFD != 0 && !setWorkDirFS {
-		preopens[c.preopenFD] = &wasm.FileEntry{Path: ".", FS: preopens[rootFD].FS}
+	preopens, err := c.fs.Preopens()
+	if err != nil {
+		return nil, err
 	}
 
 	return wasm.NewSysContext(math.MaxUint32, c.args, environ, c.stdin, c.stdout, c.stderr, c.randSource, preopens)

config_test.go

@@ -9,6 +9,7 @@ import (
 	"testing/fstest"
 
 	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/internal/fs"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 	"github.com/tetratelabs/wazero/internal/wasm"
 )
@@ -438,7 +439,7 @@ func TestModuleConfig_toSysContext(t *testing.T) {
 				nil,            // stdout
 				nil,            // stderr
 				nil,            // randSource
-				map[uint32]*wasm.FileEntry{ // openedFiles
+				map[uint32]*fs.FileEntry{ // openedFiles
 					3: {Path: "/", FS: testFS},
 					4: {Path: ".", FS: testFS},
 				},
@@ -455,7 +456,7 @@ func TestModuleConfig_toSysContext(t *testing.T) {
 				nil,            // stdout
 				nil,            // stderr
 				nil,            // randSource
-				map[uint32]*wasm.FileEntry{ // openedFiles
+				map[uint32]*fs.FileEntry{ // openedFiles
 					3: {Path: "/", FS: testFS2},
 					4: {Path: ".", FS: testFS2},
 				},
@@ -472,7 +473,7 @@ func TestModuleConfig_toSysContext(t *testing.T) {
 				nil,            // stdout
 				nil,            // stderr
 				nil,            // randSource
-				map[uint32]*wasm.FileEntry{ // openedFiles
+				map[uint32]*fs.FileEntry{ // openedFiles
 					3: {Path: ".", FS: testFS},
 				},
 			),
@@ -488,7 +489,7 @@ func TestModuleConfig_toSysContext(t *testing.T) {
 				nil,            // stdout
 				nil,            // stderr
 				nil,            // randSource
-				map[uint32]*wasm.FileEntry{ // openedFiles
+				map[uint32]*fs.FileEntry{ // openedFiles
 					3: {Path: "/", FS: testFS},
 					4: {Path: ".", FS: testFS2},
 				},
@@ -505,7 +506,7 @@ func TestModuleConfig_toSysContext(t *testing.T) {
 				nil,            // stdout
 				nil,            // stderr
 				nil,            // randSource
-				map[uint32]*wasm.FileEntry{ // openedFiles
+				map[uint32]*fs.FileEntry{ // openedFiles
 					3: {Path: ".", FS: testFS},
 					4: {Path: "/", FS: testFS2},
 				},
@@ -576,7 +577,7 @@ func TestModuleConfig_toSysContext_Errors(t *testing.T) {
 }
 
 // requireSysContext ensures wasm.NewSysContext doesn't return an error, which makes it usable in test matrices.
-func requireSysContext(t *testing.T, max uint32, args, environ []string, stdin io.Reader, stdout, stderr io.Writer, randsource io.Reader, openedFiles map[uint32]*wasm.FileEntry) *wasm.SysContext {
+func requireSysContext(t *testing.T, max uint32, args, environ []string, stdin io.Reader, stdout, stderr io.Writer, randsource io.Reader, openedFiles map[uint32]*fs.FileEntry) *wasm.SysContext {
 	sys, err := wasm.NewSysContext(max, args, environ, stdin, stdout, stderr, randsource, openedFiles)
 	require.NoError(t, err)
 	return sys

examples/wasi/testdata/sub/test.txt

@@ -0,0 +1 @@
+greet sub dir

experimental/fs.go

@@ -0,0 +1,21 @@
+package experimental
+
+import (
+	"context"
+	"io/fs"
+
+	"github.com/tetratelabs/wazero/api"
+	internalfs "github.com/tetratelabs/wazero/internal/fs"
+)
+
+// WithFS overrides fs.FS in the context-based manner. Caller needs to take responsibility for closing the filesystem.
+func WithFS(ctx context.Context, fs fs.FS) (context.Context, api.Closer, error) {
+	fsConfig := internalfs.NewFSConfig().WithFS(fs)
+	preopens, err := fsConfig.Preopens()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	fsCtx := internalfs.NewContext(preopens)
+	return context.WithValue(ctx, internalfs.Key{}, fsCtx), fsCtx, nil
+}

experimental/fs_test.go

@@ -0,0 +1,41 @@
+package experimental_test
+
+import (
+	"context"
+	_ "embed"
+	"log"
+	"testing"
+	"testing/fstest"
+
+	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/internal/fs"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+// This is a very basic integration of fs config. The main goal is to show how it is configured.
+func TestWithFS(t *testing.T) {
+	fileName := "animals.txt"
+	mapfs := fstest.MapFS{fileName: &fstest.MapFile{Data: []byte(`animals`)}}
+
+	// Set context to one that has experimental fs config
+	ctx, closer, err := experimental.WithFS(context.Background(), mapfs)
+	if err != nil {
+		log.Panicln(err)
+	}
+	defer closer.Close(ctx)
+
+	v := ctx.Value(fs.Key{})
+	require.NotNil(t, v)
+	fsCtx, ok := v.(*fs.Context)
+	require.True(t, ok)
+
+	entry, ok := fsCtx.OpenedFile(3)
+	require.True(t, ok)
+	require.Equal(t, "/", entry.Path)
+	require.Equal(t, mapfs, entry.FS)
+
+	entry, ok = fsCtx.OpenedFile(4)
+	require.True(t, ok)
+	require.Equal(t, ".", entry.Path)
+	require.Equal(t, mapfs, entry.FS)
+}