From bfb010e6fce9a3333a6ed31ff3fe6827485c2b22 Mon Sep 17 00:00:00 2001
From: Lukasz Mierzwa <lukasz@cloudflare.com>
Date: Tue, 27 Aug 2024 08:03:49 +0100
Subject: [PATCH] Fix benchmarks

---
 internal/config/config.go | 9 ++++++++-
 internal/config/match.go  | 2 +-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/internal/config/config.go b/internal/config/config.go
index abe59a2a..83cff768 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -83,7 +83,7 @@ func (cfg Config) String() string {
 func (cfg *Config) GetChecksForEntry(ctx context.Context, gen *PrometheusGenerator, entry discovery.Entry) []checks.RuleChecker {
 	enabled := []checks.RuleChecker{}
 
-	defaultMatch := []Match{{State: defaultMatchStates(ctx.Value(CommandKey).(ContextCommandVal))}}
+	defaultMatch := []Match{{State: defaultMatchStates(commandFromContext(ctx))}}
 	proms := gen.ServersForPath(entry.Path.Name)
 
 	if entry.PathError != nil || entry.Rule.Error.Err != nil {
@@ -229,3 +229,10 @@ func parseDuration(d string) (time.Duration, error) {
 	}
 	return time.Duration(mdur), nil
 }
+
+func commandFromContext(ctx context.Context) (cmd ContextCommandVal) {
+	if val := ctx.Value(CommandKey); val != nil {
+		cmd = val.(ContextCommandVal)
+	}
+	return cmd
+}
diff --git a/internal/config/match.go b/internal/config/match.go
index c7ac8b12..4a20a48b 100644
--- a/internal/config/match.go
+++ b/internal/config/match.go
@@ -105,7 +105,7 @@ func (m Match) validate(allowEmpty bool) error {
 }
 
 func (m Match) IsMatch(ctx context.Context, path string, e discovery.Entry) bool {
-	cmd := ctx.Value(CommandKey).(ContextCommandVal)
+	cmd := commandFromContext(ctx)
 
 	if m.Command != nil {
 		if cmd != *m.Command {