Add source map support

This includes both support for any sourcemap found through the source
file and generating sourcemaps when going through babel.

There are additional fixes for trying to fix off by 1 line errors in
imports, but those may need further work.

On not being able to load a sourcemap a warning is emitted but the file
is still parsed and compiled just without sourcemaps

fixes #1789, #1804

core/engine_test.go

@@ -783,8 +783,8 @@ func TestSetupException(t *testing.T) {
 		require.Error(t, err)
 		var exception errext.Exception
 		require.ErrorAs(t, err, &exception)
-		require.Equal(t, "Error: baz\n\tat baz (file:///bar.js:7:8(3))\n"+
-			"\tat file:///bar.js:4:5(3)\n\tat setup (file:///script.js:7:204(4))\n\tat native\n",
+		require.Equal(t, "Error: baz\n\tat baz (file:///bar.js:6:16(3))\n"+
+			"\tat file:///bar.js:3:8(3)\n\tat setup (file:///script.js:4:2(4))\n\tat native\n",
 			err.Error())
 	}
 }
@@ -835,7 +835,7 @@ func TestVuInitException(t *testing.T) {
 
 	var exception errext.Exception
 	require.ErrorAs(t, err, &exception)
-	assert.Equal(t, "Error: oops in 2\n\tat file:///script.js:10:8(31)\n", err.Error())
+	assert.Equal(t, "Error: oops in 2\n\tat file:///script.js:10:9(31)\n", err.Error())
 
 	var errWithHint errext.HasHint
 	require.ErrorAs(t, err, &errWithHint)

js/bundle.go

@@ -30,7 +30,6 @@ import (
 	"runtime"
 
 	"github.com/dop251/goja"
-	"github.com/dop251/goja/parser"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/afero"
 	"gopkg.in/guregu/null.v3"
@@ -84,7 +83,13 @@ func NewBundle(
 	// Compile sources, both ES5 and ES6 are supported.
 	code := string(src.Data)
 	c := compiler.New(logger)
-	pgm, _, err := c.Compile(code, src.URL.String(), "", "", true, compatMode)
+	c.COpts = compiler.Options{
+		CompatibilityMode: compatMode,
+		Strict:            true,
+		SourceMapEnabled:  true,
+		SourceMapLoader:   generateSourceMapLoader(logger, filesystems),
+	}
+	pgm, _, err := c.Compile(code, src.URL.String(), true, c.COpts)
 	if err != nil {
 		return nil, err
 	}
@@ -132,7 +137,13 @@ func NewBundleFromArchive(
 	}
 
 	c := compiler.New(logger)
-	pgm, _, err := c.Compile(string(arc.Data), arc.FilenameURL.String(), "", "", true, compatMode)
+	c.COpts = compiler.Options{
+		Strict:            true,
+		CompatibilityMode: compatMode,
+		SourceMapEnabled:  true,
+		SourceMapLoader:   generateSourceMapLoader(logger, arc.Filesystems),
+	}
+	pgm, _, err := c.Compile(string(arc.Data), arc.FilenameURL.String(), true, c.COpts)
 	if err != nil {
 		return nil, err
 	}
@@ -291,7 +302,6 @@ func (b *Bundle) Instantiate(logger logrus.FieldLogger, vuID uint64) (bi *Bundle
 // Instantiates the bundle into an existing runtime. Not public because it also messes with a bunch
 // of other things, will potentially thrash data and makes a mess in it if the operation fails.
 func (b *Bundle) instantiate(logger logrus.FieldLogger, rt *goja.Runtime, init *InitContext, vuID uint64) error {
-	rt.SetParserOptions(parser.WithDisableSourceMaps)
 	rt.SetFieldNameMapper(common.FieldNameMapper{})
 	rt.SetRandSource(common.NewRandSource())
 
@@ -338,3 +348,18 @@ func (b *Bundle) instantiate(logger logrus.FieldLogger, rt *goja.Runtime, init *
 
 	return nil
 }
+
+func generateSourceMapLoader(logger logrus.FieldLogger, filesystems map[string]afero.Fs,
+) func(path string) ([]byte, error) {
+	return func(path string) ([]byte, error) {
+		u, err := url.Parse(path)
+		if err != nil {
+			return nil, err
+		}
+		data, err := loader.Load(logger, filesystems, u, path)
+		if err != nil {
+			return nil, err
+		}
+		return data.Data, nil
+	}
+}

js/compiler/compiler.go

@@ -22,6 +22,9 @@ package compiler
 
 import (
 	_ "embed" // we need this for embedding Babel
+	"encoding/json"
+	"errors"
+	"strings"
 	"sync"
 	"time"
 
@@ -86,10 +89,13 @@ var (
 	globalBabel        *babel        // nolint:gochecknoglobals
 )
 
+const sourceMapURLFromBabel = "k6://internal-should-not-leak/file.map"
+
 // A Compiler compiles JavaScript source code (ES5.1 or ES6) into a goja.Program
 type Compiler struct {
 	logger logrus.FieldLogger
 	babel  *babel
+	COpts  Options
 }
 
 // New returns a new Compiler
@@ -108,7 +114,7 @@ func (c *Compiler) initializeBabel() error {
 }
 
 // Transform the given code into ES5
-func (c *Compiler) Transform(src, filename string) (code string, srcmap []byte, err error) {
+func (c *Compiler) Transform(src, filename string, inputSrcMap []byte) (code string, srcMap []byte, err error) {
 	if c.babel == nil {
 		onceBabel.Do(func() {
 			globalBabel, err = newBabel()
@@ -119,48 +125,96 @@ func (c *Compiler) Transform(src, filename string) (code string, srcmap []byte,
 		return
 	}
 
-	code, srcmap, err = c.babel.Transform(c.logger, src, filename)
+	code, srcMap, err = c.babel.transformImpl(c.logger, src, filename, c.COpts.SourceMapEnabled, inputSrcMap)
 	return
 }
 
+// Options are options to the compiler
+type Options struct {
+	CompatibilityMode lib.CompatibilityMode
+	SourceMapEnabled  bool
+	SourceMapLoader   func(string) ([]byte, error)
+	Strict            bool
+}
+
+// compilationState is helper struct to keep the state of a compilation
+type compilationState struct {
+	// set when we couldn't load external source map so we can try parsing without loading it
+	couldntLoadSourceMap bool
+	// srcMap is the current full sourceMap that has been generated read so far
+	srcMap []byte
+	main   bool
+
+	compiler *Compiler
+}
+
 // Compile the program in the given CompatibilityMode, wrapping it between pre and post code
-func (c *Compiler) Compile(src, filename, pre, post string,
-	strict bool, compatMode lib.CompatibilityMode) (*goja.Program, string, error) {
-	code := pre + src + post
-	ast, err := parser.ParseFile(nil, filename, code, 0, parser.WithDisableSourceMaps)
-	if err != nil {
-		if compatMode == lib.CompatibilityModeExtended {
-			code, _, err = c.Transform(src, filename)
-			if err != nil {
-				return nil, code, err
-			}
-			// the compatibility mode "decreases" here as we shouldn't transform twice
-			return c.Compile(code, filename, pre, post, strict, lib.CompatibilityModeBase)
+func (c *Compiler) Compile(src, filename string, main bool, cOpts Options) (*goja.Program, string, error) {
+	return c.compileImpl(src, filename, main, cOpts, nil)
+}
+
+// here we take the srcMap as pointer so we can change in case that we have an original srcMap that will be loaded, but
+// the code still needs to go through babel. In this way we can give the correct inputSourceMap to babel which hopefully
+// will generate a sourcemap which goes from the final source file lines to the original ones that were transpiled
+// outside of k6
+// TODO have some kind of "compilationContext" to save this instead and modify and move that
+// This won't be needed if we weren't going through babel so removing babel would kind of remove the need for that
+func (c *compilationState) sourceMapLoader(path string) ([]byte, error) {
+	if path == sourceMapURLFromBabel {
+		if !c.main {
+			return increaseMappingsByOne(c.srcMap)
 		}
-		return nil, code, err
+		return c.srcMap, nil
+	}
+	var err error
+	c.srcMap, err = c.compiler.COpts.SourceMapLoader(path)
+	if err != nil {
+		c.couldntLoadSourceMap = true
+		return nil, err
+	}
+	if !c.main {
+		return increaseMappingsByOne(c.srcMap)
+	}
+	return c.srcMap, err
+}
+
+func (c *Compiler) compileImpl(
+	src, filename string, main bool, cOpts Options, srcMap []byte,
+) (*goja.Program, string, error) {
+	code := src
+	state := compilationState{srcMap: srcMap, compiler: c, main: main}
+	if !main {
+		code = "(function(module, exports){\n" + code + "\n})\n"
+	}
+	if len(srcMap) > 0 {
+		code += "\n//# sourceMappingURL=" + sourceMapURLFromBabel
+	}
+	opts := parser.WithDisableSourceMaps
+	if cOpts.SourceMapEnabled {
+		opts = parser.WithSourceMapLoader(state.sourceMapLoader)
+	}
+	ast, err := parser.ParseFile(nil, filename, code, 0, opts)
+
+	if state.couldntLoadSourceMap {
+		state.couldntLoadSourceMap = false // reset
+		// we probably don't want to abort scripts which have source maps but they can't be found,
+		// this also will be a breaking change, so if we couldn't we retry with it disabled
+		c.logger.WithError(err).Warnf("Couldn't load source map for %s", filename)
+		ast, err = parser.ParseFile(nil, filename, code, 0, parser.WithDisableSourceMaps)
 	}
-	pgm, err := goja.CompileAST(ast, strict)
-	// Parsing only checks the syntax, not whether what the syntax expresses
-	// is actually supported (sometimes).
-	//
-	// For example, destructuring looks a lot like an object with shorthand
-	// properties, but this is only noticeable once the code is compiled, not
-	// while parsing. Even now code such as `let [x] = [2]` doesn't return an
-	// error on the parsing stage but instead in the compilation in base mode.
-	//
-	// So, because of this, if there is an error during compilation, it still might
-	// be worth it to transform the code and try again.
 	if err != nil {
-		if compatMode == lib.CompatibilityModeExtended {
-			code, _, err = c.Transform(src, filename)
+		if cOpts.CompatibilityMode == lib.CompatibilityModeExtended {
+			code, state.srcMap, err = c.Transform(src, filename, state.srcMap)
 			if err != nil {
 				return nil, code, err
 			}
 			// the compatibility mode "decreases" here as we shouldn't transform twice
-			return c.Compile(code, filename, pre, post, strict, lib.CompatibilityModeBase)
+			cOpts.CompatibilityMode = lib.CompatibilityModeBase
+			return c.compileImpl(code, filename, main, cOpts, state.srcMap)
 		}
 		return nil, code, err
 	}
+	pgm, err := goja.CompileAST(ast, cOpts.Strict)
 	return pgm, code, err
 }
 
@@ -194,16 +248,53 @@ func newBabel() (*babel, error) {
 	return result, err
 }
 
-// Transform the given code into ES5, while synchronizing to ensure only a single
+func increaseMappingsByOne(sourceMap []byte) ([]byte, error) {
+	var err error
+	m := make(map[string]interface{})
+	if err = json.Unmarshal(sourceMap, &m); err != nil {
+		return nil, err
+	}
+	mappings, ok := m["mappings"]
+	if !ok {
+		// no mappings, no idea what this will do, but just return it as technically we can have sourcemap with sections
+		// TODO implement incrementing of `offset` in the sections?
+		// TODO (kind of alternatively) drop the newline in the "commonjs" wrapping and have only the first line wrong
+		// except everything being off by one
+		return sourceMap, nil
+	}
+	if str, ok := mappings.(string); ok {
+		// ';' is the separator between lines so just adding 1 will make all mappings be for the line after which they were
+		// originally
+		m["mappings"] = ";" + str
+	} else {
+		// we have mappings but it's not a string - this is some kind of erro
+		return nil, errors.New(`missing "mappings" in sourcemap`)
+	}
+
+	return json.Marshal(m)
+}
+
+// transformImpl the given code into ES5, while synchronizing to ensure only a single
 // bundle instance / Goja VM is in use at a time.
-// TODO the []byte is there to be used as the returned sourcemap and will be done in PR #2082
-func (b *babel) Transform(logger logrus.FieldLogger, src, filename string) (string, []byte, error) {
+func (b *babel) transformImpl(
+	logger logrus.FieldLogger, src, filename string, sourceMapsEnabled bool, inputSrcMap []byte,
+) (string, []byte, error) {
 	b.m.Lock()
 	defer b.m.Unlock()
 	opts := make(map[string]interface{})
 	for k, v := range DefaultOpts {
 		opts[k] = v
 	}
+	if sourceMapsEnabled {
+		opts["sourceMaps"] = true
+		if inputSrcMap != nil {
+			srcMap := new(map[string]interface{})
+			if err := json.Unmarshal(inputSrcMap, &srcMap); err != nil {
+				return "", nil, err
+			}
+			opts["inputSourceMap"] = srcMap
+		}
+	}
 	opts["filename"] = filename
 
 	startTime := time.Now()
@@ -218,7 +309,27 @@ func (b *babel) Transform(logger logrus.FieldLogger, src, filename string) (stri
 	if err = b.vm.ExportTo(vO.Get("code"), &code); err != nil {
 		return code, nil, err
 	}
-	return code, nil, err
+	if !sourceMapsEnabled {
+		return code, nil, nil
+	}
+
+	// this is to make goja try to load a sourcemap.
+	// it is a special url as it should never leak outside of this code
+	// additionally the alternative support from babel is to embed *the whole* sourcemap at the end
+	if i := strings.Index(code, "//# sourceMappingURL="); i > 0 {
+		code = code[:i]
+	}
+	code += "\n//# sourceMappingURL=" + sourceMapURLFromBabel
+	stringify, err := b.vm.RunString("(function(m) { return JSON.stringify(m)})")
+	if err != nil {
+		return code, nil, err
+	}
+	c, _ := goja.AssertFunction(stringify)
+	mapAsJSON, err := c(goja.Undefined(), vO.Get("map"))
+	if err != nil {
+		return code, nil, err
+	}
+	return code, []byte(mapAsJSON.String()), nil
 }
 
 // Pool is a pool of compilers so it can be used easier in parallel tests as they have their own babel.