Add RuntimeState and keylogger to it

grafana · Apr 7, 2022 · 5f8266b · 5f8266b
1 parent ece4633
commit 5f8266b
Show file tree

Hide file tree

Showing 18 changed files with 478 additions and 168 deletions.
diff --git a/api/v1/setup_teardown_routes_test.go b/api/v1/setup_teardown_routes_test.go
@@ -146,12 +146,13 @@ func TestSetupData(t *testing.T) {
 			builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 
 			runner, err := js.New(
-				logger,
+				&lib.RuntimeState{
+					Logger:         logger,
+					BuiltinMetrics: builtinMetrics,
+					Registry:       registry,
+				},
 				&loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: testCase.script},
 				nil,
-				lib.RuntimeOptions{},
-				builtinMetrics,
-				registry,
 			)
 			require.NoError(t, err)
 			runner.SetOptions(lib.Options{

diff --git a/cmd/integration_tests/eventloop/eventloop_test.go b/cmd/integration_tests/eventloop/eventloop_test.go
@@ -34,15 +34,16 @@ func eventLoopTest(t *testing.T, script []byte, testHandle func(context.Context,
 	registry := metrics.NewRegistry()
 	builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 	runner, err := js.New(
-		logger,
+		&lib.RuntimeState{
+			Logger:         logger,
+			BuiltinMetrics: builtinMetrics,
+			Registry:       registry,
+		},
 		&loader.SourceData{
 			URL:  &url.URL{Path: "/script.js"},
 			Data: script,
 		},
 		nil,
-		lib.RuntimeOptions{},
-		builtinMetrics,
-		registry,
 	)
 	require.NoError(t, err)
 

diff --git a/cmd/run.go b/cmd/run.go
@@ -268,6 +268,11 @@ func (c *cmdRun) run(cmd *cobra.Command, args []string) error {
 	logger.Debug("Waiting for engine processes to finish...")
 	engineWait()
 	logger.Debug("Everything has finished, exiting k6!")
+	if test.keywriter != nil {
+		if err := test.keywriter.Close(); err != nil {
+			logger.WithError(err).Warn("error while closing the SSLKEYLOGFILE")
+		}
+	}
 	if interrupt != nil {
 		return interrupt
 	}

diff --git a/cmd/runtime_options.go b/cmd/runtime_options.go
@@ -117,6 +117,12 @@ func getRuntimeOptions(flags *pflag.FlagSet, environment map[string]string) (lib
 		}
 	}
 
+	if envVar, ok := environment["SSLKEYLOGFILE"]; ok {
+		if !opts.KeyWriter.Valid {
+			opts.KeyWriter = null.StringFrom(envVar)
+		}
+	}
+
 	if opts.IncludeSystemEnvVars.Bool { // If enabled, gather the actual system environment variables
 		opts.Env = environment
 	}

diff --git a/cmd/test_load.go b/cmd/test_load.go
@@ -4,6 +4,9 @@ import (
 	"archive/tar"
 	"bytes"
 	"fmt"
+	"io"
+	"os"
+	"sync"
 
 	"github.com/spf13/afero"
 	"github.com/spf13/cobra"
@@ -31,6 +34,7 @@ type loadedTest struct {
 	metricsRegistry *metrics.Registry
 	builtInMetrics  *metrics.BuiltinMetrics
 	initRunner      lib.Runner // TODO: rename to something more appropriate
+	keywriter       io.Closer
 
 	// Only set if cliConfigGetter is supplied to loadAndConfigureTest() or if
 	// consolidateDeriveAndValidateConfig() is manually called.
@@ -100,12 +104,25 @@ func (lt *loadedTest) initializeFirstRunner(gs *globalState) error {
 		testType = detectTestType(lt.source.Data)
 	}
 
+	state := &lib.RuntimeState{
+		Logger:         gs.logger,
+		RuntimeOptions: lt.runtimeOptions,
+		BuiltinMetrics: lt.builtInMetrics,
+		Registry:       lt.metricsRegistry,
+	}
+	if lt.runtimeOptions.KeyWriter.Valid {
+		fs := lt.fileSystems["file"]
+		f, err := fs.OpenFile(lt.runtimeOptions.KeyWriter.String, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0o600)
+		if err != nil {
+			return err
+		}
+		lt.keywriter = f
+		state.KeyLogger = &syncWriter{w: f}
+	}
 	switch testType {
 	case testTypeJS:
 		logger.Debug("Trying to load as a JS test...")
-		runner, err := js.New(
-			gs.logger, lt.source, lt.fileSystems, lt.runtimeOptions, lt.builtInMetrics, lt.metricsRegistry,
-		)
+		runner, err := js.New(state, lt.source, lt.fileSystems)
 		// TODO: should we use common.UnwrapGojaInterruptedError() here?
 		if err != nil {
 			return fmt.Errorf("could not load JS test '%s': %w", testPath, err)
@@ -126,7 +143,7 @@ func (lt *loadedTest) initializeFirstRunner(gs *globalState) error {
 		switch arc.Type {
 		case testTypeJS:
 			logger.Debug("Evaluating JS from archive bundle...")
-			lt.initRunner, err = js.NewFromArchive(gs.logger, arc, lt.runtimeOptions, lt.builtInMetrics, lt.metricsRegistry)
+			lt.initRunner, err = js.NewFromArchive(state, arc)
 			if err != nil {
 				return fmt.Errorf("could not load JS from test archive bundle '%s': %w", testPath, err)
 			}
@@ -202,3 +219,14 @@ func (lt *loadedTest) consolidateDeriveAndValidateConfig(
 
 	return nil
 }
+
+type syncWriter struct {
+	w io.Writer
+	m sync.Mutex
+}
+
+func (cw *syncWriter) Write(b []byte) (int, error) {
+	cw.m.Lock()
+	defer cw.m.Unlock()
+	return cw.w.Write(b)
+}
diff --git a/converter/har/converter_test.go b/converter/har/converter_test.go
@@ -61,10 +61,15 @@ func TestBuildK6RequestObject(t *testing.T) {
 	assert.NoError(t, err)
 	registry := metrics.NewRegistry()
 	builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
-	_, err = js.New(testutils.NewLogger(t), &loader.SourceData{
-		URL:  &url.URL{Path: "/script.js"},
-		Data: []byte(fmt.Sprintf("export default function() { res = http.batch([%v]); }", v)),
-	}, nil, lib.RuntimeOptions{}, builtinMetrics, registry)
+	_, err = js.New(
+		&lib.RuntimeState{
+			Logger:         testutils.NewLogger(t),
+			BuiltinMetrics: builtinMetrics,
+			Registry:       registry,
+		}, &loader.SourceData{
+			URL:  &url.URL{Path: "/script.js"},
+			Data: []byte(fmt.Sprintf("export default function() { res = http.batch([%v]); }", v)),
+		}, nil)
 	assert.NoError(t, err)
 }
 

diff --git a/core/engine_test.go b/core/engine_test.go
@@ -552,12 +552,13 @@ func TestSentReceivedMetrics(t *testing.T) {
 		registry := metrics.NewRegistry()
 		builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 		r, err := js.New(
-			testutils.NewLogger(t),
+			&lib.RuntimeState{
+				Logger:         testutils.NewLogger(t),
+				BuiltinMetrics: builtinMetrics,
+				Registry:       registry,
+			},
 			&loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: []byte(ts.Code)},
 			nil,
-			lib.RuntimeOptions{},
-			builtinMetrics,
-			registry,
 		)
 		require.NoError(t, err)
 
@@ -690,12 +691,13 @@ func TestRunTags(t *testing.T) {
 	registry := metrics.NewRegistry()
 	builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 	r, err := js.New(
-		testutils.NewLogger(t),
+		&lib.RuntimeState{
+			Logger:         testutils.NewLogger(t),
+			BuiltinMetrics: builtinMetrics,
+			Registry:       registry,
+		},
 		&loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: script},
 		nil,
-		lib.RuntimeOptions{},
-		builtinMetrics,
-		registry,
 	)
 	require.NoError(t, err)
 
@@ -785,12 +787,13 @@ func TestSetupTeardownThresholds(t *testing.T) {
 	registry := metrics.NewRegistry()
 	builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 	runner, err := js.New(
-		testutils.NewLogger(t),
+		&lib.RuntimeState{
+			Logger:         testutils.NewLogger(t),
+			BuiltinMetrics: builtinMetrics,
+			Registry:       registry,
+		},
 		&loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: script},
 		nil,
-		lib.RuntimeOptions{},
-		builtinMetrics,
-		registry,
 	)
 	require.NoError(t, err)
 
@@ -838,12 +841,13 @@ func TestSetupException(t *testing.T) {
 	registry := metrics.NewRegistry()
 	builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 	runner, err := js.New(
-		testutils.NewLogger(t),
+		&lib.RuntimeState{
+			Logger:         testutils.NewLogger(t),
+			BuiltinMetrics: builtinMetrics,
+			Registry:       registry,
+		},
 		&loader.SourceData{URL: &url.URL{Scheme: "file", Path: "/script.js"}, Data: script},
 		map[string]afero.Fs{"file": memfs},
-		lib.RuntimeOptions{},
-		builtinMetrics,
-		registry,
 	)
 	require.NoError(t, err)
 
@@ -891,11 +895,13 @@ func TestVuInitException(t *testing.T) {
 	registry := metrics.NewRegistry()
 	builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 	runner, err := js.New(
-		logger,
+		&lib.RuntimeState{
+			Logger:         testutils.NewLogger(t),
+			BuiltinMetrics: builtinMetrics,
+			Registry:       registry,
+		},
 		&loader.SourceData{URL: &url.URL{Scheme: "file", Path: "/script.js"}, Data: script},
-		nil, lib.RuntimeOptions{},
-		builtinMetrics,
-		registry,
+		nil,
 	)
 	require.NoError(t, err)
 
@@ -962,12 +968,13 @@ func TestEmittedMetricsWhenScalingDown(t *testing.T) {
 	registry := metrics.NewRegistry()
 	builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 	runner, err := js.New(
-		testutils.NewLogger(t),
+		&lib.RuntimeState{
+			Logger:         testutils.NewLogger(t),
+			BuiltinMetrics: builtinMetrics,
+			Registry:       registry,
+		},
 		&loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: script},
 		nil,
-		lib.RuntimeOptions{},
-		builtinMetrics,
-		registry,
 	)
 	require.NoError(t, err)
 
@@ -1050,7 +1057,11 @@ func TestMetricsEmission(t *testing.T) {
 			registry := metrics.NewRegistry()
 			builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 			runner, err := js.New(
-				testutils.NewLogger(t),
+				&lib.RuntimeState{
+					Logger:         testutils.NewLogger(t),
+					BuiltinMetrics: builtinMetrics,
+					Registry:       registry,
+				},
 				&loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: []byte(fmt.Sprintf(`
 				import { sleep } from "k6";
 				import { Counter } from "k6/metrics";
@@ -1074,9 +1085,6 @@ func TestMetricsEmission(t *testing.T) {
 				}
 				`, tc.minIterDuration, tc.defaultBody))},
 				nil,
-				lib.RuntimeOptions{},
-				builtinMetrics,
-				registry,
 			)
 			require.NoError(t, err)
 
@@ -1162,12 +1170,13 @@ func TestMinIterationDurationInSetupTeardownStage(t *testing.T) {
 			registry := metrics.NewRegistry()
 			builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
 			runner, err := js.New(
-				testutils.NewLogger(t),
+				&lib.RuntimeState{
+					Logger:         testutils.NewLogger(t),
+					BuiltinMetrics: builtinMetrics,
+					Registry:       registry,
+				},
 				&loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: []byte(tc.script)},
 				nil,
-				lib.RuntimeOptions{},
-				builtinMetrics,
-				registry,
 			)
 			require.NoError(t, err)
 
@@ -1278,8 +1287,13 @@ func TestActiveVUsCount(t *testing.T) {
 
 	registry := metrics.NewRegistry()
 	builtinMetrics := metrics.RegisterBuiltinMetrics(registry)
-	runner, err := js.New(logger, &loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: script}, nil, rtOpts,
-		builtinMetrics, registry)
+	runner, err := js.New(
+		&lib.RuntimeState{
+			Logger:         logger,
+			BuiltinMetrics: builtinMetrics,
+			Registry:       registry,
+			RuntimeOptions: rtOpts,
+		}, &loader.SourceData{URL: &url.URL{Path: "/script.js"}, Data: script}, nil)
 	require.NoError(t, err)
 
 	mockOutput := mockoutput.New()