Refactor auto extension resolution to not use k6deps and esbuild

internal/cmd/cloud.go

@@ -85,14 +85,6 @@ func (c *cmdCloud) preRun(cmd *cobra.Command, _ []string) error {
 //
 //nolint:funlen,gocognit,cyclop
 func (c *cmdCloud) run(cmd *cobra.Command, args []string) error {
-	printBanner(c.gs)
-
-	progressBar := pb.New(
-		pb.WithConstLeft("Init"),
-		pb.WithConstProgress(0, "Loading test script..."),
-	)
-	printBar(c.gs, progressBar)
-
 	test, err := loadAndConfigureLocalTest(c.gs, cmd, args, getPartialConfig)
 	if err != nil {
 		return err
@@ -111,6 +103,13 @@ func (c *cmdCloud) run(cmd *cobra.Command, args []string) error {
 	// TODO: validate for usage of execution segment
 	// TODO: validate for externally controlled executor (i.e. executors that aren't distributable)
 	// TODO: move those validations to a separate function and reuse validateConfig()?
+	printBanner(c.gs)
+
+	progressBar := pb.New(
+		pb.WithConstLeft("Init"),
+		pb.WithConstProgress(0, "Loading test script..."),
+	)
+	printBar(c.gs, progressBar)
 
 	modifyAndPrintBar(c.gs, progressBar, pb.WithConstProgress(0, "Building the archive..."))
 	arc := testRunState.Runner.MakeArchive()

internal/cmd/launcher.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"io/fs"
 	"maps"
@@ -12,6 +13,8 @@ import (
 	"slices"
 	"strings"
 	"syscall"
+	"unicode"
+	"unicode/utf8"
 
 	"github.com/Masterminds/semver/v3"
 	"github.com/grafana/k6deps"
@@ -20,6 +23,8 @@ import (
 
 	"go.k6.io/k6/cloudapi"
 	"go.k6.io/k6/cmd/state"
+	"go.k6.io/k6/errext"
+	"go.k6.io/k6/errext/exitcodes"
 	"go.k6.io/k6/ext"
 	"go.k6.io/k6/internal/build"
 	"go.k6.io/k6/lib/fsext"
@@ -218,6 +223,10 @@ func (b *customBinary) run(gs *state.GlobalState) error {
 	for {
 		select {
 		case err := <-done:
+			var exitError *exec.ExitError
+			if errors.As(err, &exitError) {
+				return errext.WithExitCodeIfNone(errAlreadyReported, exitcodes.ExitCode(exitError.ExitCode())) //nolint:gosec
+			}
 			return err
 		case sig := <-sigC:
 			gs.Logger.
@@ -227,10 +236,12 @@ func (b *customBinary) run(gs *state.GlobalState) error {
 	}
 }
 
+// used just to signal we shouldn't print error again
+var errAlreadyReported = fmt.Errorf("already reported error")
+
 // isCustomBuildRequired checks if there is at least one dependency that are not satisfied by the binary
 // considering the version of k6 and any built-in extension
-func isCustomBuildRequired(deps map[string]*semver.Constraints, k6Version string, exts []*ext.Extension) bool {
-	// return early if there are no dependencies
+func isCustomBuildRequired(deps dependencies, k6Version string, exts []*ext.Extension) bool {
 	if len(deps) == 0 {
 		return false
 	}
@@ -328,7 +339,7 @@ func formatDependencies(deps map[string]string) string {
 // Presently, only the k6 input script or archive (if any) is passed to k6deps for scanning.
 // TODO: if k6 receives the input from stdin, it is not used for scanning because we don't know
 // if it is a script or an archive
-func analyze(gs *state.GlobalState, args []string) (map[string]*semver.Constraints, error) {
+func analyze(gs *state.GlobalState, args []string) (dependencies, error) {
 	dopts := &k6deps.Options{
 		LookupEnv: func(key string) (string, bool) { v, ok := gs.Env[key]; return v, ok },
 		Manifest:  k6deps.Source{Ignore: true},
@@ -363,7 +374,7 @@ func analyze(gs *state.GlobalState, args []string) (map[string]*semver.Constrain
 	if err != nil {
 		return nil, err
 	}
-	result := make(map[string]*semver.Constraints, len(deps))
+	result := make(dependencies, len(deps))
 	for n, dep := range deps {
 		result[n] = dep.Constraints
 	}
@@ -408,3 +419,70 @@ func extractToken(gs *state.GlobalState) (string, error) {
 
 	return config.Token.String, nil
 }
+
+func processUseDirectives(name string, text []byte, deps dependencies) error {
+	directives := findDirectives(text)
+
+	for _, directive := range directives {
+		// normalize spaces
+		directive = strings.ReplaceAll(directive, "  ", " ")
+		if !strings.HasPrefix(directive, "use k6") {
+			continue
+		}
+		directive = strings.TrimSpace(strings.TrimPrefix(directive, "use k6"))
+		if !strings.HasPrefix(directive, "with k6/x/") {
+			err := deps.update("k6", directive)
+			if err != nil {
+				return fmt.Errorf("error while parsing use directives in %q: %w", name, err)
+			}
+			continue
+		}
+		directive = strings.TrimSpace(strings.TrimPrefix(directive, "with "))
+		dep, constraint, _ := strings.Cut(directive, " ")
+		err := deps.update(dep, constraint)
+		if err != nil {
+			return fmt.Errorf("error while parsing use directives in %q: %w", name, err)
+		}
+	}
+
+	return nil
+}
+
+func findDirectives(text []byte) []string {
+	// parse #! at beginning of file
+	if bytes.HasPrefix(text, []byte("#!")) {
+		_, text, _ = bytes.Cut(text, []byte("\n"))
+	}
+
+	var result []string
+
+	for i := 0; i < len(text); {
+		r, width := utf8.DecodeRune(text[i:])
+		switch {
+		case unicode.IsSpace(r) || r == rune(';'): // skip all spaces and ;
+			i += width
+		case r == '"' || r == '\'': // string literals
+			idx := bytes.IndexRune(text[i+width:], r)
+			if idx < 0 {
+				return result
+			}
+			result = append(result, string(text[i+width:i+width+idx]))
+			i += width + idx + 1
+		case bytes.HasPrefix(text[i:], []byte("//")):
+			idx := bytes.IndexRune(text[i+width:], '\n')
+			if idx < 0 {
+				return result
+			}
+			i += width + idx + 1
+		case bytes.HasPrefix(text[i:], []byte("/*")):
+			idx := bytes.Index(text[i+width:], []byte("*/"))
+			if idx < 0 {
+				return result
+			}
+			i += width + idx + 2
+		default:
+			return result
+		}
+	}
+	return result
+}

internal/cmd/launcher_test.go

@@ -210,6 +210,7 @@ func TestLauncherLaunch(t *testing.T) {
 			t.Parallel()
 
 			ts := tests.NewGlobalTestState(t)
+			ts.Env["K6_OLD_RESOLUTION"] = "true"
 
 			k6Args := append([]string{"k6"}, tc.k6Cmd)
 			k6Args = append(k6Args, tc.k6Args...)
@@ -277,6 +278,7 @@ func TestLauncherViaStdin(t *testing.T) {
 	k6Args := []string{"k6", "archive", "-"}
 
 	ts := tests.NewGlobalTestState(t)
+	ts.Env["K6_OLD_RESOLUTION"] = "true"
 	ts.CmdArgs = k6Args
 
 	// k6deps uses os package to access files. So we need to use it in the global state
@@ -567,3 +569,171 @@ func TestGetProviderConfig(t *testing.T) {
 		})
 	}
 }
+
+func TestProcessUseDirectives(t *testing.T) {
+	t.Parallel()
+	tests := map[string]struct {
+		input          string
+		expectedOutput map[string]string
+		expectedError  string
+	}{
+		"nothing": {
+			input: "export default function() {}",
+		},
+		"nothing really": {
+			input: `"use k6"`,
+			expectedOutput: map[string]string{
+				"k6": "",
+			},
+		},
+		"k6 pinning": {
+			input: `"use k6 > 1.4.0"`,
+			expectedOutput: map[string]string{
+				"k6": "> 1.4.0",
+			},
+		},
+		"a extension": {
+			input: `"use k6 with k6/x/sql"`,
+			expectedOutput: map[string]string{
+				"k6/x/sql": "",
+			},
+		},
+		"an extension with constraint": {
+			input: `"use k6 with k6/x/sql > 1.4.0"`,
+			expectedOutput: map[string]string{
+				"k6/x/sql": "> 1.4.0",
+			},
+		},
+		"complex": {
+			input: `
+				// something here
+				"use k6 with k6/x/A"
+				function a (){
+					"use k6 with k6/x/B"
+					let s = JSON.stringify( "use k6 with k6/x/C")
+					"use k6 with k6/x/D"
+
+					return s
+				}
+
+				export const b = "use k6 with k6/x/E"
+				"use k6 with k6/x/F"
+
+				// Here for esbuild and k6 warnings
+				a()
+				export default function(){}
+				`,
+			expectedOutput: map[string]string{
+				"k6/x/A": "",
+			},
+		},
+
+		"repeat": {
+			input: `
+				"use k6 with k6/x/A"
+				"use k6 with k6/x/A"
+				`,
+			expectedOutput: map[string]string{
+				"k6/x/A": "",
+			},
+		},
+		"repeat with constraint first": {
+			input: `
+				"use k6 with k6/x/A > 1.4.0"
+				"use k6 with k6/x/A"
+				`,
+			expectedOutput: map[string]string{
+				"k6/x/A": "> 1.4.0",
+			},
+		},
+		"constraint difference": {
+			input: `
+				"use k6 > 1.4.0"
+				"use k6 = 1.2.3"
+				`,
+			expectedError: `error while parsing use directives in "name.js": already have constraint for "k6", when parsing "=1.2.3"`,
+		},
+		"constraint difference for extensions": {
+			input: `
+				"use k6 with k6/x/A > 1.4.0"
+				"use k6 with k6/x/A = 1.2.3"
+				`,
+			expectedError: `error while parsing use directives in "name.js": already have constraint for "k6/x/A", when parsing "=1.2.3"`,
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+			deps := make(dependencies)
+			for k, v := range test.expectedOutput {
+				require.NoError(t, deps.update(k, v))
+			}
+			if len(test.expectedError) > 0 {
+				deps = nil
+			}
+
+			m := make(dependencies)
+			err := processUseDirectives("name.js", []byte(test.input), m)
+			if len(test.expectedError) > 0 {
+				require.ErrorContains(t, err, test.expectedError)
+			} else {
+				require.EqualValues(t, deps, m)
+			}
+		})
+	}
+}
+
+func TestFindDirectives(t *testing.T) {
+	t.Parallel()
+	tests := map[string]struct {
+		input          string
+		expectedOutput []string
+	}{
+		"nothing": {
+			input:          "export default function() {}",
+			expectedOutput: nil,
+		},
+		"nothing really": {
+			input:          `"use k6"`,
+			expectedOutput: []string{"use k6"},
+		},
+		"multiline": {
+			input: `
+			"use k6 with k6/x/sql"
+			"something"
+			`,
+			expectedOutput: []string{"use k6 with k6/x/sql", "something"},
+		},
+		"multiline start at beginning": {
+			input: `
+"use k6 with k6/x/sql"
+"something"
+			`,
+			expectedOutput: []string{"use k6 with k6/x/sql", "something"},
+		},
+		"multiline comments": {
+			input: `#!/bin/sh
+			// here comment "hello"
+"use k6 with k6/x/sql";
+			/*
+			"something else here as well"
+			*/
+	;
+"something";
+const l = 5
+"more"
+			`,
+			expectedOutput: []string{"use k6 with k6/x/sql", "something"},
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			m := findDirectives([]byte(test.input))
+			assert.EqualValues(t, test.expectedOutput, m)
+		})
+	}
+}