Lock module resolution after initialization (#3680)

Co-authored-by: Ivan <2103732+codebien@users.noreply.github.com>

js/bundle.go

@@ -112,6 +112,7 @@ func newBundle(
 	if err != nil {
 		return nil, err
 	}
+	bundle.ModuleResolver.Lock()
 
 	err = bundle.populateExports(updateOptions, exports)
 	if err != nil {

js/modules/resolution.go

@@ -10,6 +10,8 @@ import (
 	"go.k6.io/k6/loader"
 )
 
+const notPreviouslyResolvedModule = "the module %q was not previously resolved during initialization (__VU==0)"
+
 // FileLoader is a type alias for a function that returns the contents of the referenced file.
 type FileLoader func(specifier *url.URL, name string) ([]byte, error)
 
@@ -32,6 +34,7 @@ type ModuleResolver struct {
 	goModules map[string]interface{}
 	loadCJS   FileLoader
 	compiler  *compiler.Compiler
+	locked    bool
 }
 
 // NewModuleResolver returns a new module resolution instance that will resolve.
@@ -55,6 +58,9 @@ func (mr *ModuleResolver) resolveSpecifier(basePWD *url.URL, arg string) (*url.U
 }
 
 func (mr *ModuleResolver) requireModule(name string) (module, error) {
+	if mr.locked {
+		return nil, fmt.Errorf(notPreviouslyResolvedModule, name)
+	}
 	mod, ok := mr.goModules[name]
 	if !ok {
 		return nil, fmt.Errorf("unknown module: %s", name)
@@ -81,6 +87,14 @@ func (mr *ModuleResolver) resolveLoaded(basePWD *url.URL, arg string, data []byt
 	return mod, err
 }
 
+// Lock locks the module's resolution from any further new resolving operation.
+// It means that it relays only its internal cache and on the fact that it has already
+// seen previously the module during the initialization.
+// It is the same approach used for opening file operations.
+func (mr *ModuleResolver) Lock() {
+	mr.locked = true
+}
+
 func (mr *ModuleResolver) resolve(basePWD *url.URL, arg string) (module, error) {
 	if cached, ok := mr.cache[arg]; ok {
 		return cached.mod, cached.err
@@ -102,6 +116,9 @@ func (mr *ModuleResolver) resolve(basePWD *url.URL, arg string) (module, error)
 			return cached.mod, cached.err
 		}
 
+		if mr.locked {
+			return nil, fmt.Errorf(notPreviouslyResolvedModule, arg)
+		}
 		// Fall back to loading
 		data, err := mr.loadCJS(specifier, arg)
 		if err != nil {

js/runner_test.go

@@ -1569,6 +1569,101 @@ func TestVUDoesNonExistingPathnUnderConditions(t *testing.T) {
 	assert.Contains(t, err.Error(), "open() can't be used with files that weren't previously opened during initialization (__VU==0)")
 }
 
+func TestVUDoesRequireUnderV0Condition(t *testing.T) {
+	t.Parallel()
+
+	baseFS := fsext.NewMemMapFs()
+	data := `
+			if (__VU == 0) {
+				let data = require("/home/somebody/test.js");
+			}
+			exports.default = function() {
+				console.log("hey")
+			}
+		`
+	require.NoError(t, fsext.WriteFile(baseFS, "/home/somebody/test.js", []byte(`exports=42`), fs.ModePerm))
+	require.NoError(t, fsext.WriteFile(baseFS, "/script.js", []byte(data), fs.ModePerm))
+
+	fs := fsext.NewCacheOnReadFs(baseFS, fsext.NewMemMapFs(), 0)
+
+	r, err := getSimpleRunner(t, "/script.js", data, fs)
+	require.NoError(t, err)
+
+	_, err = r.NewVU(context.Background(), 1, 1, make(chan metrics.SampleContainer, 100))
+	require.NoError(t, err)
+}
+
+func TestVUDoesNotRequireUnderConditions(t *testing.T) {
+	t.Parallel()
+
+	baseFS := fsext.NewMemMapFs()
+	data := `
+			if (__VU > 0) {
+				let data = require("/home/somebody/test.js");
+			}
+			exports.default = function() {
+				console.log("hey")
+			}
+		`
+	require.NoError(t, fsext.WriteFile(baseFS, "/home/somebody/test.js", []byte(`exports=42`), fs.ModePerm))
+	require.NoError(t, fsext.WriteFile(baseFS, "/script.js", []byte(data), fs.ModePerm))
+
+	fs := fsext.NewCacheOnReadFs(baseFS, fsext.NewMemMapFs(), 0)
+
+	r, err := getSimpleRunner(t, "/script.js", data, fs)
+	require.NoError(t, err)
+
+	_, err = r.NewVU(context.Background(), 1, 1, make(chan metrics.SampleContainer, 100))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), " was not previously resolved during initialization (__VU==0)")
+}
+
+func TestVUDoesRequireUnderConditions(t *testing.T) {
+	t.Parallel()
+
+	baseFS := fsext.NewMemMapFs()
+	data := `
+			if (__VU == 0) {
+				require("/home/somebody/test.js");
+				require("/home/somebody/test2.js");
+			}
+
+			if (__VU % 2 == 1) {
+				require("/home/somebody/test.js");
+			}
+
+			if (__VU % 2 == 0) {
+				require("/home/somebody/test2.js");
+			}
+
+			exports.default = function() {
+				console.log("hey")
+			}
+		`
+	require.NoError(t, fsext.WriteFile(baseFS, "/home/somebody/test.js", []byte(`console.log("test.js", __VU)`), fs.ModePerm))
+	require.NoError(t, fsext.WriteFile(baseFS, "/home/somebody/test2.js", []byte(`console.log("test2.js", __VU)`), fs.ModePerm))
+	require.NoError(t, fsext.WriteFile(baseFS, "/script.js", []byte(data), fs.ModePerm))
+
+	fs := fsext.NewCacheOnReadFs(baseFS, fsext.NewMemMapFs(), 0)
+
+	logger, hook := testutils.NewLoggerWithHook(t, logrus.InfoLevel)
+	r, err := getSimpleRunner(t, "/script.js", data, fs, logger)
+	require.NoError(t, err)
+	logs := hook.Drain()
+	require.Len(t, logs, 2)
+
+	_, err = r.NewVU(context.Background(), 1, 1, make(chan metrics.SampleContainer, 100))
+	require.NoError(t, err)
+	logs = hook.Drain()
+	require.Len(t, logs, 1)
+	require.Contains(t, logs[0].Message, "test.js 1")
+	_, err = r.NewVU(context.Background(), 2, 2, make(chan metrics.SampleContainer, 100))
+	require.NoError(t, err)
+	logs = hook.Drain()
+	require.Len(t, logs, 1)
+	require.Contains(t, logs[0].Message, "test2.js 2")
+}
+
 func TestVUIntegrationCookiesReset(t *testing.T) {
 	t.Parallel()
 	tb := httpmultibin.NewHTTPMultiBin(t)