fix: Change how the secret store is handled (#1281)

* fix: move handling of secrets store around

In the current version of the code, the secret store is populated when
the prober is created. The problem with that is that the prober is
created *once* in the lifetime of the check, but the secret store needs
to be updated more or less on a continuous basis (because it contains
the token used to access the secrets).

I'm sure there's a better way to handle this, but it took me a while to
figure out the issue.

Signed-off-by: Marcelo E. Magallon <[email protected]>
Co-authored-by: Dougal Matthews <[email protected]>

Author: mem

cmd/synthetic-monitoring-agent/main.go

@@ -43,6 +43,19 @@ import (
 const (
 	exitFail             = 1
 	defTelemetryTimeSpan = 5 // min
+
+	// The expiration time for access tokens is 10 minutes. Until we get
+	// expiration information from the API, set this to a low number.
+	// Normally it should be 15 minutes.
+	//
+	// This specific number is accounting for ~ 3 minutes of total
+	// execution time for scripts, plus 30 seconds of buffer for retrieving
+	// secrets, sending script to runner, etc. We expect secrets to be
+	// loaded near the start of the script, but given it's a program we
+	// cannot enforce that. k6 would have enforced that by making secrets
+	// part of the options block, but we do not want to go anywhere near
+	// that.
+	secretsTimeout = 450 * time.Second // 7.5 minutes
 )
 
 // run is the main entry point for the program.
@@ -277,7 +290,13 @@ func run(args []string, stdout io.Writer) error {
 		})
 	}
 
-	tm := tenants.NewManager(ctx, synthetic_monitoring.NewTenantsClient(conn), tenantCh, 15*time.Minute)
+	tm := tenants.NewManager(
+		ctx,
+		synthetic_monitoring.NewTenantsClient(conn),
+		tenantCh,
+		secretsTimeout,
+		zl.With().Str("subsystem", "tenant_manager").Logger(),
+	)
 
 	pusherRegistry := pusher.NewRegistry[pusher.Factory]()
 	pusherRegistry.MustRegister(pusherV1.Name, pusherV1.NewPublisher)

internal/adhoc/adhoc_test.go

@@ -600,6 +600,6 @@ func (r *testK6Runner) WithLogger(logger *zerolog.Logger) k6runner.Runner {
 	return r
 }
 
-func (r *testK6Runner) Run(ctx context.Context, script k6runner.Script) (*k6runner.RunResponse, error) {
+func (r *testK6Runner) Run(ctx context.Context, script k6runner.Script, secretStore k6runner.SecretStore) (*k6runner.RunResponse, error) {
 	return &k6runner.RunResponse{}, nil
 }

internal/checks/checks_test.go

@@ -504,7 +504,7 @@ func (noopRunner) WithLogger(logger *zerolog.Logger) k6runner.Runner {
 	return r
 }
 
-func (noopRunner) Run(ctx context.Context, script k6runner.Script) (*k6runner.RunResponse, error) {
+func (noopRunner) Run(ctx context.Context, script k6runner.Script, secretStore k6runner.SecretStore) (*k6runner.RunResponse, error) {
 	return &k6runner.RunResponse{}, nil
 }
 

internal/k6runner/http.go

@@ -44,8 +44,9 @@ func (r requestError) Error() string {
 
 // HTTPRunRequest
 type HTTPRunRequest struct {
-	Script   `json:",inline"`
-	NotAfter time.Time `json:"notAfter"`
+	Script      `json:",inline"`
+	SecretStore SecretStore `json:",inline"`
+	NotAfter    time.Time   `json:"notAfter"`
 }
 
 type RunResponse struct {
@@ -63,7 +64,7 @@ func (r HttpRunner) WithLogger(logger *zerolog.Logger) Runner {
 
 var ErrUnexpectedStatus = errors.New("unexpected status code")
 
-func (r HttpRunner) Run(ctx context.Context, script Script) (*RunResponse, error) {
+func (r HttpRunner) Run(ctx context.Context, script Script, secretStore SecretStore) (*RunResponse, error) {
 	if r.backoff == 0 {
 		panic("zero backoff, runner is misconfigured, refusing to DoS")
 	}
@@ -92,7 +93,7 @@ func (r HttpRunner) Run(ctx context.Context, script Script) (*RunResponse, error
 		start := time.Now()
 
 		var err error
-		response, err = r.request(ctx, script)
+		response, err = r.request(ctx, script, secretStore)
 		if err == nil {
 			r.logger.Debug().Bytes("metrics", response.Metrics).Bytes("logs", response.Logs).Msg("script result")
 			r.metrics.Requests.With(map[string]string{metricLabelSuccess: "1", metricLabelRetriable: ""}).Inc()
@@ -136,7 +137,7 @@ func (r HttpRunner) Run(ctx context.Context, script Script) (*RunResponse, error
 // errRetryable indicates that an error is retryable. It is typically joined with another error.
 var errRetryable = errors.New("retryable")
 
-func (r HttpRunner) request(ctx context.Context, script Script) (*RunResponse, error) {
+func (r HttpRunner) request(ctx context.Context, script Script, secretStore SecretStore) (*RunResponse, error) {
 	checkTimeout := time.Duration(script.Settings.Timeout) * time.Millisecond
 	if checkTimeout == 0 {
 		return nil, ErrNoTimeout
@@ -161,8 +162,9 @@ func (r HttpRunner) request(ctx context.Context, script Script) (*RunResponse, e
 	// This allows runners to not waste time running scripts which will not complete before the client gives up on the
 	// request.
 	runRequest := HTTPRunRequest{
-		Script:   script,
-		NotAfter: notAfter,
+		Script:      script,
+		SecretStore: secretStore,
+		NotAfter:    notAfter,
 	}
 
 	reqBody, err := json.Marshal(runRequest)

internal/k6runner/http_test.go

@@ -80,7 +80,7 @@ func TestHttpRunnerRun(t *testing.T) {
 	ctx, cancel := testhelper.Context(ctx, t)
 	t.Cleanup(cancel)
 
-	_, err := runner.Run(ctx, script)
+	_, err := runner.Run(ctx, script, SecretStore{})
 	require.NoError(t, err)
 }
 
@@ -129,7 +129,7 @@ func TestHttpRunnerRunError(t *testing.T) {
 	ctx, cancel = testhelper.Context(ctx, t)
 	t.Cleanup(cancel)
 
-	_, err := runner.Run(ctx, script)
+	_, err := runner.Run(ctx, script, SecretStore{})
 	require.ErrorIs(t, err, ErrUnexpectedStatus)
 }
 
@@ -322,7 +322,7 @@ func TestScriptHTTPRun(t *testing.T) {
 				zlogger  = zerolog.Nop()
 			)
 
-			success, err := script.Run(ctx, registry, logger, zlogger)
+			success, err := script.Run(ctx, registry, logger, zlogger, SecretStore{})
 			require.Equal(t, tc.expectSuccess, success)
 			require.Equal(t, tc.expectLogs, logger.buf.String())
 			if tc.expectErrorAs == nil {
@@ -457,7 +457,7 @@ func TestHTTPProcessorRetries(t *testing.T) {
 					logger   testLogger
 					zlogger  = zerolog.New(io.Discard)
 				)
-				success, err := processor.Run(ctx, registry, &logger, zlogger)
+				success, err := processor.Run(ctx, registry, &logger, zlogger, SecretStore{})
 				require.ErrorIs(t, err, tc.expectError)
 				require.Equal(t, tc.expectError == nil, success)
 				require.Equal(t, tc.expectRequests, requests.Load())
@@ -503,7 +503,7 @@ func TestHTTPProcessorRetries(t *testing.T) {
 			logger   testLogger
 			zlogger  = zerolog.New(io.Discard)
 		)
-		success, err := processor.Run(ctx, registry, &logger, zlogger)
+		success, err := processor.Run(ctx, registry, &logger, zlogger, SecretStore{})
 		require.NoError(t, err)
 		require.True(t, success)
 	})

internal/k6runner/k6runner.go

@@ -29,8 +29,6 @@ type Script struct {
 	// CheckInfo holds information about the SM check that triggered this script.
 	CheckInfo CheckInfo `json:"check"`
 	// SecretStore holds the location and token for accessing secrets
-	SecretStore SecretStore `json:"secretStore"`
-	// TODO: Add features.
 }
 
 type SecretStore struct {
@@ -92,7 +90,7 @@ var ErrNoTimeout = errors.New("check has no timeout")
 
 type Runner interface {
 	WithLogger(logger *zerolog.Logger) Runner
-	Run(ctx context.Context, script Script) (*RunResponse, error)
+	Run(ctx context.Context, script Script, secretStore SecretStore) (*RunResponse, error)
 }
 
 type RunnerOpts struct {
@@ -149,11 +147,11 @@ var (
 	ErrFromRunner  = errors.New("runner reported an error")
 )
 
-func (r Processor) Run(ctx context.Context, registry *prometheus.Registry, logger logger.Logger, internalLogger zerolog.Logger) (bool, error) {
+func (r Processor) Run(ctx context.Context, registry *prometheus.Registry, logger logger.Logger, internalLogger zerolog.Logger, secretStore SecretStore) (bool, error) {
 	k6runner := r.runner.WithLogger(&internalLogger)
 
 	// TODO: This error message is okay to be Debug for local k6 execution, but should be Error for remote runners.
-	result, err := k6runner.Run(ctx, r.script)
+	result, err := k6runner.Run(ctx, r.script, secretStore)
 	if err != nil {
 		internalLogger.Debug().
 			Err(err).

internal/k6runner/k6runner_test.go

@@ -90,7 +90,7 @@ func TestScriptRun(t *testing.T) {
 	// We already know tha parsing the metrics and the logs is working, so
 	// we are only interested in verifying that the script runs without
 	// errors.
-	success, err := processor.Run(ctx, registry, &logger, zlogger)
+	success, err := processor.Run(ctx, registry, &logger, zlogger, SecretStore{})
 	require.NoError(t, err)
 	require.True(t, success)
 }
@@ -130,7 +130,7 @@ type testRunner struct {
 
 var _ Runner = &testRunner{}
 
-func (r *testRunner) Run(ctx context.Context, script Script) (*RunResponse, error) {
+func (r *testRunner) Run(ctx context.Context, script Script, secretStore SecretStore) (*RunResponse, error) {
 	return &RunResponse{
 		Metrics: r.metrics,
 		Logs:    r.logs,

internal/k6runner/local.go

@@ -35,7 +35,7 @@ func (r Local) WithLogger(logger *zerolog.Logger) Runner {
 	return r
 }
 
-func (r Local) Run(ctx context.Context, script Script) (*RunResponse, error) {
+func (r Local) Run(ctx context.Context, script Script, secretStore SecretStore) (*RunResponse, error) {
 	logger := r.logger.With().Object("checkInfo", &script.CheckInfo).Logger()
 
 	afs := afero.Afero{Fs: r.fs}
@@ -78,7 +78,7 @@ func (r Local) Run(ctx context.Context, script Script) (*RunResponse, error) {
 	executable := r.k6path
 	// TODO(d0ugal): This is a short-term hack to use a different k6 binary when
 	// secrets are configured. This should be removed when the default binary has been updated
-	if script.SecretStore.IsConfigured() {
+	if secretStore.IsConfigured() {
 		executable = r.k6path + "-gsm"
 		logger.Info().
 			Str("k6_path", r.k6path).
@@ -96,13 +96,18 @@ func (r Local) Run(ctx context.Context, script Script) (*RunResponse, error) {
 	defer cancel()
 
 	var configFile string
-	if script.SecretStore.IsConfigured() {
+	if secretStore.IsConfigured() {
 		var cleanup func()
-		configFile, cleanup, err = createSecretConfigFile(script.SecretStore.Url, script.SecretStore.Token)
+		configFile, cleanup, err = createSecretConfigFile(secretStore.Url, secretStore.Token)
 		if err != nil {
 			return nil, fmt.Errorf("cannot create secret config file: %w", err)
 		}
 		defer cleanup()
+
+		logger.Debug().
+			Str("secret_config_file", configFile).
+			Str("secrets_url", secretStore.Url).
+			Msg("Using secret config file")
 	}
 
 	args, err := r.buildK6Args(script, metricsFn, logsFn, scriptFn, configFile)
@@ -213,7 +218,7 @@ func (r Local) buildK6Args(script Script, metricsFn, logsFn, scriptFn, configFil
 	}
 
 	// Add secretStore configuration if available
-	if script.SecretStore.IsConfigured() {
+	if configFile != "" {
 		args = append(args, "--secret-source", "grafanasecrets=config="+configFile)
 	}
 

internal/k6runner/local_test.go

@@ -110,12 +110,7 @@ func TestBuildK6Args(t *testing.T) {
 			},
 		},
 		"script with secrets": {
-			script: Script{
-				SecretStore: SecretStore{
-					Url:   secretUrl,
-					Token: "secret-token",
-				},
-			},
+			script:        Script{},
 			metricsFn:     "/tmp/metrics.json",
 			logsFn:        "/tmp/logs.log",
 			scriptFn:      "/tmp/script.js",

internal/prober/browser/browser.go

@@ -24,9 +24,10 @@ type Module struct {
 }
 
 type Prober struct {
-	logger    zerolog.Logger
-	module    Module
-	processor *k6runner.Processor
+	logger           zerolog.Logger
+	module           Module
+	processor        *k6runner.Processor
+	secretsRetriever func(context.Context) (k6runner.SecretStore, error)
 }
 
 func NewProber(ctx context.Context, check model.Check, logger zerolog.Logger, runner k6runner.Runner, store secrets.SecretProvider) (Prober, error) {
@@ -36,11 +37,6 @@ func NewProber(ctx context.Context, check model.Check, logger zerolog.Logger, ru
 		return p, errUnsupportedCheck
 	}
 
-	secretStore, err := store.GetSecretCredentials(ctx, check.GlobalTenantID())
-	if err != nil {
-		return p, err
-	}
-
 	p.module = Module{
 		Prober: sm.CheckTypeBrowser.String(),
 		Script: k6runner.Script{
@@ -52,20 +48,14 @@ func NewProber(ctx context.Context, check model.Check, logger zerolog.Logger, ru
 		},
 	}
 
-	if secretStore != nil {
-		p.module.Script.SecretStore = k6runner.SecretStore{
-			Url:   secretStore.Url,
-			Token: secretStore.Token,
-		}
-	}
-
 	processor, err := k6runner.NewProcessor(p.module.Script, runner)
 	if err != nil {
 		return p, err
 	}
 
 	p.processor = processor
 	p.logger = logger
+	p.secretsRetriever = newCredentialsRetriever(store, check.GlobalTenantID())
 
 	return p, nil
 }
@@ -75,7 +65,14 @@ func (p Prober) Name() string {
 }
 
 func (p Prober) Probe(ctx context.Context, target string, registry *prometheus.Registry, logger logger.Logger) (bool, float64) {
-	success, err := p.processor.Run(ctx, registry, logger, p.logger)
+	secretStore, err := p.secretsRetriever(ctx)
+
+	if err != nil {
+		p.logger.Error().Err(err).Msg("running probe")
+		return false, 0
+	}
+
+	success, err := p.processor.Run(ctx, registry, logger, p.logger, secretStore)
 	if err != nil {
 		p.logger.Error().Err(err).Msg("running probe")
 		return false, 0
@@ -84,3 +81,24 @@ func (p Prober) Probe(ctx context.Context, target string, registry *prometheus.R
 	// TODO(mem): implement custom duration extraction.
 	return success, 0
 }
+
+// TODO(mem): This should probably be in the k6runner package.
+func newCredentialsRetriever(provider secrets.SecretProvider, tenantID model.GlobalID) func(context.Context) (k6runner.SecretStore, error) {
+	return func(ctx context.Context) (k6runner.SecretStore, error) {
+		var store k6runner.SecretStore
+
+		credentials, err := provider.GetSecretCredentials(ctx, tenantID)
+		if err != nil {
+			return store, err
+		}
+
+		if credentials != nil {
+			store = k6runner.SecretStore{
+				Url:   credentials.Url,
+				Token: credentials.Token,
+			}
+		}
+
+		return store, nil
+	}
+}