feat: add feature flag for protocol secrets support (#1626)

Add SM_AGENT_ENABLE_PROTOCOL_SECRETS environment variable to enable
protocol secrets support. This allows testing the feature before enabling
it by default.

- Add boolFromEnv helper function to parse boolean environment variables
- Add EnableProtocolSecrets config field and environment variable parsing
- Add SupportsProtocolSecrets field to UpdaterOptions and HandlerOpts
- Update RegisterProbe calls to use the feature flag instead of hardcoded false
- Add tests to verify the feature flag works correctly

Author: d0ugal

cmd/synthetic-monitoring-agent/main.go

@@ -12,6 +12,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"path/filepath"
+	"strconv"
 	"syscall"
 	"time"
 
@@ -59,29 +60,30 @@ func run(args []string, stdout io.Writer) error {
 	var (
 		features = feature.NewCollection()
 		config   = struct {
-			DevMode              bool
-			Debug                bool
-			Verbose              bool
-			ReportVersion        bool
-			GrpcApiServerAddr    string
-			GrpcInsecure         bool
-			ApiToken             Secret
-			EnableChangeLogLevel bool
-			EnableDisconnect     bool
-			EnablePProf          bool
-			HttpListenAddr       string
-			K6URI                string
-			K6BlacklistedIP      string
-			SelectedPublisher    string
-			TelemetryTimeSpan    int
-			AutoMemLimit         bool
-			MemLimitRatio        float64
-			DisableK6            bool
-			DisableUsageReports  bool
-			CacheType            cache.Kind
-			CacheLocalCapacity   int
-			CacheLocalTTL        time.Duration
-			MemcachedServers     StringList
+			DevMode               bool
+			Debug                 bool
+			Verbose               bool
+			ReportVersion         bool
+			GrpcApiServerAddr     string
+			GrpcInsecure          bool
+			ApiToken              Secret
+			EnableChangeLogLevel  bool
+			EnableDisconnect      bool
+			EnablePProf           bool
+			HttpListenAddr        string
+			K6URI                 string
+			K6BlacklistedIP       string
+			SelectedPublisher     string
+			TelemetryTimeSpan     int
+			AutoMemLimit          bool
+			MemLimitRatio         float64
+			DisableK6             bool
+			DisableUsageReports   bool
+			CacheType             cache.Kind
+			CacheLocalCapacity    int
+			CacheLocalTTL         time.Duration
+			MemcachedServers      StringList
+			EnableProtocolSecrets bool
 		}{
 			GrpcApiServerAddr:  "localhost:4031",
 			HttpListenAddr:     "localhost:4050",
@@ -172,6 +174,10 @@ func run(args []string, stdout io.Writer) error {
 	// Using API_TOKEN should be deprecated after March 1st, 2023.
 	config.ApiToken = Secret(stringFromEnv("API_TOKEN", stringFromEnv("SM_AGENT_API_TOKEN", string(config.ApiToken))))
 
+	// Enable protocol secrets support via environment variable.
+	// This is a feature flag to allow testing before enabling by default.
+	config.EnableProtocolSecrets = boolFromEnv("SM_AGENT_ENABLE_PROTOCOL_SECRETS", config.EnableProtocolSecrets)
+
 	if config.ApiToken == "" {
 		return fmt.Errorf("invalid API token")
 	}
@@ -348,21 +354,22 @@ func run(args []string, stdout io.Writer) error {
 	)
 
 	checksUpdater, err := checks.NewUpdater(checks.UpdaterOptions{
-		Conn:                  conn,
-		Logger:                zl.With().Str("subsystem", "updater").Logger(),
-		Backoff:               newConnectionBackoff(),
-		Publisher:             publisher,
-		TenantCh:              tenantCh,
-		IsConnected:           readynessHandler.Set,
-		PromRegisterer:        promRegisterer,
-		Features:              features,
-		K6Runner:              k6Runner,
-		ScraperFactory:        scraper.New,
-		TenantLimits:          limits,
-		SecretProvider:        secretProvider,
-		Telemeter:             telemetry,
-		UsageReporter:         usageReporter,
-		CostAttributionLabels: cals,
+		Conn:                    conn,
+		Logger:                  zl.With().Str("subsystem", "updater").Logger(),
+		Backoff:                 newConnectionBackoff(),
+		Publisher:               publisher,
+		TenantCh:                tenantCh,
+		IsConnected:             readynessHandler.Set,
+		PromRegisterer:          promRegisterer,
+		Features:                features,
+		K6Runner:                k6Runner,
+		ScraperFactory:          scraper.New,
+		TenantLimits:            limits,
+		SecretProvider:          secretProvider,
+		Telemeter:               telemetry,
+		UsageReporter:           usageReporter,
+		CostAttributionLabels:   cals,
+		SupportsProtocolSecrets: config.EnableProtocolSecrets,
 	})
 
 	if err != nil {
@@ -374,15 +381,16 @@ func run(args []string, stdout io.Writer) error {
 	})
 
 	adhocHandler, err := adhoc.NewHandler(adhoc.HandlerOpts{
-		Conn:           conn,
-		Logger:         zl.With().Str("subsystem", "adhoc").Logger(),
-		Backoff:        newConnectionBackoff(),
-		Publisher:      publisher,
-		TenantCh:       tenantCh,
-		PromRegisterer: promRegisterer,
-		Features:       features,
-		K6Runner:       k6Runner,
-		SecretProvider: secretProvider,
+		Conn:                    conn,
+		Logger:                  zl.With().Str("subsystem", "adhoc").Logger(),
+		Backoff:                 newConnectionBackoff(),
+		Publisher:               publisher,
+		TenantCh:                tenantCh,
+		PromRegisterer:          promRegisterer,
+		Features:                features,
+		K6Runner:                k6Runner,
+		SecretProvider:          secretProvider,
+		SupportsProtocolSecrets: config.EnableProtocolSecrets,
 	})
 	if err != nil {
 		return fmt.Errorf("cannot create ad-hoc checks handler: %w", err)
@@ -445,6 +453,21 @@ func stringFromEnv(name string, override string) string {
 	return os.Getenv(name)
 }
 
+func boolFromEnv(name string, defaultValue bool) bool {
+	value := os.Getenv(name)
+	if value == "" {
+		return defaultValue
+	}
+
+	parsed, err := strconv.ParseBool(value)
+	if err != nil {
+		// If parsing fails, return the default value
+		return defaultValue
+	}
+
+	return parsed
+}
+
 func validateK6URI(uri string) (string, error) {
 	u, err := url.Parse(uri)
 	if err != nil {

internal/adhoc/adhoc.go

@@ -42,6 +42,7 @@ type Handler struct {
 	runnerFactory                func(context.Context, *sm.AdHocRequest) (*runner, error)
 	grpcAdhocChecksClientFactory func(conn ClientConn) (sm.AdHocChecksClient, error)
 	proberFactory                prober.ProberFactory
+	supportsProtocolSecrets      bool
 }
 
 // Error represents errors returned from this package.
@@ -106,15 +107,16 @@ func (b constantBackoff) Duration() time.Duration { return time.Duration(b) }
 
 // HandlerOpts is used to pass configuration options to the Handler.
 type HandlerOpts struct {
-	Conn           ClientConn
-	Logger         zerolog.Logger
-	Backoff        Backoffer
-	Publisher      pusher.Publisher
-	TenantCh       chan<- sm.Tenant
-	PromRegisterer prometheus.Registerer
-	Features       feature.Collection
-	K6Runner       k6runner.Runner
-	SecretProvider secrets.SecretProvider
+	Conn                    ClientConn
+	Logger                  zerolog.Logger
+	Backoff                 Backoffer
+	Publisher               pusher.Publisher
+	TenantCh                chan<- sm.Tenant
+	PromRegisterer          prometheus.Registerer
+	Features                feature.Collection
+	K6Runner                k6runner.Runner
+	SecretProvider          secrets.SecretProvider
+	SupportsProtocolSecrets bool
 
 	// these two fields exists so that tests can pass alternate
 	// implementations, they are unexported so that clients of this
@@ -148,6 +150,7 @@ func NewHandler(opts HandlerOpts) (*Handler, error) {
 		runnerFactory:                opts.runnerFactory,
 		grpcAdhocChecksClientFactory: opts.grpcAdhocChecksClientFactory,
 		proberFactory:                prober.NewProberFactory(opts.K6Runner, 0, opts.Features, opts.SecretProvider),
+		supportsProtocolSecrets:      opts.SupportsProtocolSecrets,
 		api: apiInfo{
 			conn: opts.Conn,
 		},
@@ -308,11 +311,10 @@ func (h *Handler) loop(ctx context.Context) error {
 	result, err := client.RegisterProbe(
 		ctx,
 		&sm.ProbeInfo{
-			Version:    version.Short(),
-			Commit:     version.Commit(),
-			Buildstamp: version.Buildstamp(),
-			// TODO(d0ugal): We will switch this to true when the implementation is complete in the Agent and API.
-			SupportsProtocolSecrets: false,
+			Version:                 version.Short(),
+			Commit:                  version.Commit(),
+			Buildstamp:              version.Buildstamp(),
+			SupportsProtocolSecrets: h.supportsProtocolSecrets,
 		},
 	)
 	if err != nil {

internal/adhoc/adhoc_test.go

@@ -51,6 +51,47 @@ func TestNewHandler(t *testing.T) {
 	require.NotNil(t, h.grpcAdhocChecksClientFactory)
 	require.Nil(t, h.probe, "probe should not be set at this point")
 	require.NotNil(t, h.metrics.opsCounter)
+	require.False(t, h.supportsProtocolSecrets, "default value should be false")
+}
+
+func TestHandlerSupportsProtocolSecrets(t *testing.T) {
+	features := feature.NewCollection()
+	require.NoError(t, features.Set("adhoc"))
+
+	var capturedProbeInfo *sm.ProbeInfo
+	testClient := &testClient{
+		logger: zerolog.New(io.Discard),
+		registerProbeHook: func(info *sm.ProbeInfo) {
+			capturedProbeInfo = info
+		},
+	}
+
+	opts := HandlerOpts{
+		Conn:                    &grpcTestConn{},
+		Logger:                  zerolog.New(io.Discard),
+		Publisher:               channelPublisher(make(chan pusher.Payload)),
+		TenantCh:                make(chan sm.Tenant),
+		PromRegisterer:          prometheus.NewPedanticRegistry(),
+		Features:                features,
+		SupportsProtocolSecrets: true,
+		grpcAdhocChecksClientFactory: func(conn ClientConn) (sm.AdHocChecksClient, error) {
+			return testClient, nil
+		},
+	}
+
+	h, err := NewHandler(opts)
+	require.NoError(t, err)
+	require.NotNil(t, h)
+	require.True(t, h.supportsProtocolSecrets, "should be set to true")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	// Run will call RegisterProbe, which should capture the ProbeInfo
+	_ = h.Run(ctx)
+
+	require.NotNil(t, capturedProbeInfo, "RegisterProbe should have been called")
+	require.True(t, capturedProbeInfo.SupportsProtocolSecrets, "SupportsProtocolSecrets should be true")
 }
 
 func TestHandlerRun(t *testing.T) {
@@ -258,11 +299,16 @@ type testClient struct {
 	logger              zerolog.Logger
 	registerProbeError  error
 	getAdhocChecksError error
+	registerProbeHook   func(*sm.ProbeInfo)
 }
 
 func (c *testClient) RegisterProbe(ctx context.Context, in *sm.ProbeInfo, opts ...grpc.CallOption) (*sm.RegisterProbeResult, error) {
 	c.logger.Info().Str("func", "RegisterProbe").Interface("in", in).Interface("opts", opts).Send()
 
+	if c.registerProbeHook != nil {
+		c.registerProbeHook(in)
+	}
+
 	if c.registerProbeError != nil {
 		return nil, c.registerProbeError
 	}