Add header to ease tracking HTTP requests (#624)

* add a header to identify user for probe http check

* blank test id

* don't allow users to override this header.

* lint

* update + include MultiHTTP

* use http.Header

* review

internal/adhoc/adhoc.go

@@ -139,7 +139,7 @@ func NewHandler(opts HandlerOpts) (*Handler, error) {
 		tenantCh:                     opts.TenantCh,
 		runnerFactory:                opts.runnerFactory,
 		grpcAdhocChecksClientFactory: opts.grpcAdhocChecksClientFactory,
-		proberFactory:                prober.NewProberFactory(opts.K6Runner),
+		proberFactory:                prober.NewProberFactory(opts.K6Runner, 0),
 		api: apiInfo{
 			conn: opts.Conn,
 		},

internal/prober/http/http.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"hash/fnv"
+	"net/http"
 	"net/url"
 	"strconv"
 	"strings"
@@ -28,11 +29,15 @@ type Prober struct {
 	cacheBustingQueryParamName string
 }
 
-func NewProber(ctx context.Context, check sm.Check, logger zerolog.Logger) (Prober, error) {
+func NewProber(ctx context.Context, check sm.Check, logger zerolog.Logger, reservedHeaders http.Header) (Prober, error) {
 	if check.Settings.Http == nil {
 		return Prober{}, errUnsupportedCheck
 	}
 
+	if len(reservedHeaders) > 0 {
+		augmentHttpHeaders(&check, reservedHeaders)
+	}
+
 	cfg, err := settingsToModule(ctx, check.Settings.Http, logger)
 	if err != nil {
 		return Prober{}, err
@@ -252,6 +257,35 @@ func convertOAuth2Config(ctx context.Context, cfg *sm.OAuth2Config, logger zerol
 	return r, nil
 }
 
+// Overrides any user-provided headers with our own augmented values
+// for reserved headers.
+func augmentHttpHeaders(check *sm.Check, reservedHeaders http.Header) {
+	headers := []string{}
+	for _, header := range check.Settings.Http.Headers {
+		name, _ := strToHeaderNameValue(header)
+
+		_, present := reservedHeaders[http.CanonicalHeaderKey(name)]
+		if present {
+			continue // users can't override reserved headers with their own values
+		}
+
+		headers = append(headers, header)
+	}
+
+	for key, values := range reservedHeaders {
+		var b strings.Builder
+		for _, value := range values {
+			b.Reset()
+			b.WriteString(key)
+			b.WriteRune(':')
+			b.WriteString(value)
+			headers = append(headers, b.String())
+		}
+	}
+
+	check.Settings.Http.Headers = headers
+}
+
 func buildHttpHeaders(headers []string) map[string]string {
 	userAgentHeader := "user-agent"
 

internal/prober/http/http_test.go

@@ -2,7 +2,9 @@ package http
 
 import (
 	"context"
+	"fmt"
 	"io"
+	"net/http"
 	"net/url"
 	"strings"
 	"testing"
@@ -32,18 +34,26 @@ func TestNewProber(t *testing.T) {
 	}{
 		"default": {
 			input: sm.Check{
+				Id:     3,
 				Target: "www.grafana.com",
 				Settings: sm.CheckSettings{
-					Http: &sm.HttpSettings{},
+					Http: &sm.HttpSettings{
+						Headers: []string{
+							"X-SM-ID: 9880-98",
+						},
+					},
 				},
 			},
 			expected: Prober{
-				config: getDefaultModule().getConfigModule(),
+				config: getDefaultModule().
+					addHttpHeader("X-Sm-Id", "3-3"). // is checkId twice since probeId is unavailable here
+					getConfigModule(),
 			},
 			ExpectError: false,
 		},
 		"no-settings": {
 			input: sm.Check{
+				Id:     1,
 				Target: "www.grafana.com",
 				Settings: sm.CheckSettings{
 					Http: nil,
@@ -54,12 +64,14 @@ func TestNewProber(t *testing.T) {
 		},
 		"headers": {
 			input: sm.Check{
+				Id:     5,
 				Target: "www.grafana.com",
 				Settings: sm.CheckSettings{
 					Http: &sm.HttpSettings{
 						Headers: []string{
 							"uSeR-aGeNt: test-user-agent",
 							"some-header: some-value",
+							"x-SM-iD: 3232-32",
 						},
 					},
 				},
@@ -68,6 +80,7 @@ func TestNewProber(t *testing.T) {
 				config: getDefaultModule().
 					addHttpHeader("uSeR-aGeNt", "test-user-agent").
 					addHttpHeader("some-header", "some-value").
+					addHttpHeader("X-Sm-Id", "5-5").
 					getConfigModule(),
 			},
 			ExpectError: false,
@@ -78,7 +91,12 @@ func TestNewProber(t *testing.T) {
 		ctx := context.Background()
 		logger := zerolog.New(io.Discard)
 		t.Run(name, func(t *testing.T) {
-			actual, err := NewProber(ctx, testcase.input, logger)
+			// origin identifier for http requests is checkId-probeId; testing with checkId twice in the absence of probeId
+			checkId := testcase.input.Id
+			reservedHeaders := http.Header{}
+			reservedHeaders.Add("x-sm-id", fmt.Sprintf("%d-%d", checkId, checkId))
+
+			actual, err := NewProber(ctx, testcase.input, logger, reservedHeaders)
 			require.Equal(t, &testcase.expected, &actual)
 			if testcase.ExpectError {
 				require.Error(t, err, "unsupported check")
@@ -211,7 +229,7 @@ func TestProbe(t *testing.T) {
 			zl := zerolog.Logger{}
 			kl := log.NewLogfmtLogger(io.Discard)
 
-			prober, err := NewProber(ctx, check, zl)
+			prober, err := NewProber(ctx, check, zl, http.Header{})
 			require.NoError(t, err)
 			require.Equal(t, tc.expectFailure, !prober.Probe(ctx, check.Target, registry, kl))
 

internal/prober/multihttp/multihttp.go

@@ -3,6 +3,8 @@ package multihttp
 import (
 	"context"
 	"errors"
+	"net/http"
+	"strings"
 	"time"
 
 	"github.com/grafana/synthetic-monitoring-agent/internal/k6runner"
@@ -27,7 +29,7 @@ type Prober struct {
 	script *k6runner.Script
 }
 
-func NewProber(ctx context.Context, check sm.Check, logger zerolog.Logger, runner k6runner.Runner) (Prober, error) {
+func NewProber(ctx context.Context, check sm.Check, logger zerolog.Logger, runner k6runner.Runner, reservedHeaders http.Header) (Prober, error) {
 	var p Prober
 
 	if check.Settings.Multihttp == nil {
@@ -38,6 +40,10 @@ func NewProber(ctx context.Context, check sm.Check, logger zerolog.Logger, runne
 		return p, err
 	}
 
+	if len(reservedHeaders) > 0 {
+		augmentHttpHeaders(&check, reservedHeaders)
+	}
+
 	p.config = settingsToModule(check.Settings.Multihttp)
 	timeout := time.Duration(check.Timeout) * time.Millisecond
 	p.config.Timeout = timeout
@@ -87,3 +93,27 @@ func settingsToModule(settings *sm.MultiHttpSettings) Module {
 
 	return m
 }
+
+// Overrides any user-provided headers with our own augmented values
+// for 'reserved' headers.
+func augmentHttpHeaders(check *sm.Check, reservedHeaders http.Header) {
+	updatedHeaders := []*sm.HttpHeader{}
+	for key, values := range reservedHeaders {
+		updatedHeaders = append(updatedHeaders, &sm.HttpHeader{Name: key, Value: strings.Join(values, ",")})
+	}
+
+	for _, entry := range check.Settings.Multihttp.Entries {
+		heads := entry.Request.Headers
+		for _, headerPtr := range heads {
+			_, present := reservedHeaders[http.CanonicalHeaderKey(headerPtr.Name)]
+
+			if present {
+				continue // users can't override reserved headers with their own values
+			}
+
+			updatedHeaders = append(updatedHeaders, headerPtr)
+		}
+
+		entry.Request.Headers = updatedHeaders
+	}
+}

internal/prober/multihttp/multihttp_test.go

@@ -2,6 +2,8 @@ package multihttp
 
 import (
 	"context"
+	"fmt"
+	"net/http"
 	"testing"
 	"time"
 
@@ -25,6 +27,7 @@ func TestNewProber(t *testing.T) {
 		"valid": {
 			expectFailure: false,
 			check: sm.Check{
+				Id:        1,
 				Target:    "http://www.example.org",
 				Job:       "test",
 				Frequency: 10 * 1000,
@@ -52,6 +55,7 @@ func TestNewProber(t *testing.T) {
 		"settings must be valid": {
 			expectFailure: true,
 			check: sm.Check{
+				Id:        2,
 				Target:    "http://www.example.org",
 				Job:       "test",
 				Frequency: 10 * 1000,
@@ -70,6 +74,7 @@ func TestNewProber(t *testing.T) {
 		"must contain multihttp settings": {
 			expectFailure: true,
 			check: sm.Check{
+				Id:        3,
 				Target:    "http://www.example.org",
 				Job:       "test",
 				Frequency: 10 * 1000,
@@ -81,17 +86,50 @@ func TestNewProber(t *testing.T) {
 				},
 			},
 		},
+		"header overwrite protection is case-insensitive": {
+			expectFailure: false,
+			check: sm.Check{
+				Id:        4,
+				Target:    "http://www.example.org",
+				Job:       "test",
+				Frequency: 10 * 1000,
+				Timeout:   10 * 1000,
+				Probes:    []int64{1},
+				Settings: sm.CheckSettings{
+					Multihttp: &sm.MultiHttpSettings{
+						Entries: []*sm.MultiHttpEntry{
+							{
+								Request: &sm.MultiHttpEntryRequest{
+									Url:     "http://www.example.org",
+									Headers: []*sm.HttpHeader{{Name: "X-sM-Id", Value: "9880-9880"}},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
 			var runner noopRunner
-			p, err := NewProber(ctx, tc.check, logger, runner)
+			checkId := tc.check.Id
+			reservedHeaders := http.Header{}
+			reservedHeaders.Add("x-sm-id", fmt.Sprintf("%d-%d", checkId, checkId))
+
+			p, err := NewProber(ctx, tc.check, logger, runner, reservedHeaders)
 			if tc.expectFailure {
 				require.Error(t, err)
 				return
 			}
 
+			requestHeaders := tc.check.Settings.Multihttp.Entries[0].Request.Headers
+			require.Equal(t, len(requestHeaders), 1) // reserved header is present
+
+			require.Equal(t, requestHeaders[0].Name, "X-Sm-Id")
+			require.Equal(t, requestHeaders[0].Value, fmt.Sprintf("%d-%d", checkId, checkId))
+
 			require.NoError(t, err)
 			require.Equal(t, proberName, p.config.Prober)
 			require.Equal(t, 10*time.Second, p.config.Timeout)

internal/prober/multihttp/script_test.go

@@ -3,6 +3,7 @@ package multihttp
 import (
 	"bytes"
 	"context"
+	"net/http"
 	"net/http/httptest"
 	"path/filepath"
 	"strings"
@@ -812,7 +813,7 @@ func TestSettingsToScript(t *testing.T) {
 	k6path := filepath.Join(testhelper.ModuleDir(t), "dist", "k6")
 	runner := k6runner.New(k6path)
 
-	prober, err := NewProber(ctx, check, logger, runner)
+	prober, err := NewProber(ctx, check, logger, runner, http.Header{})
 	require.NoError(t, err)
 
 	reg := prometheus.NewPedanticRegistry()

internal/prober/prober.go

@@ -3,12 +3,13 @@ package prober
 import (
 	"context"
 	"fmt"
+	"net/http"
 
 	"github.com/grafana/synthetic-monitoring-agent/internal/k6runner"
 	"github.com/grafana/synthetic-monitoring-agent/internal/model"
 	"github.com/grafana/synthetic-monitoring-agent/internal/prober/dns"
 	"github.com/grafana/synthetic-monitoring-agent/internal/prober/grpc"
-	"github.com/grafana/synthetic-monitoring-agent/internal/prober/http"
+	httpProber "github.com/grafana/synthetic-monitoring-agent/internal/prober/http"
 	"github.com/grafana/synthetic-monitoring-agent/internal/prober/icmp"
 	"github.com/grafana/synthetic-monitoring-agent/internal/prober/logger"
 	"github.com/grafana/synthetic-monitoring-agent/internal/prober/multihttp"
@@ -34,12 +35,14 @@ type ProberFactory interface {
 }
 
 type proberFactory struct {
-	runner k6runner.Runner
+	runner  k6runner.Runner
+	probeId int64
 }
 
-func NewProberFactory(runner k6runner.Runner) ProberFactory {
+func NewProberFactory(runner k6runner.Runner, probeId int64) ProberFactory {
 	return proberFactory{
-		runner: runner,
+		runner:  runner,
+		probeId: probeId,
 	}
 }
 
@@ -56,7 +59,8 @@ func (f proberFactory) New(ctx context.Context, logger zerolog.Logger, check mod
 		target = check.Target
 
 	case sm.CheckTypeHttp:
-		p, err = http.NewProber(ctx, check.Check, logger)
+		reservedHeaders := f.getReservedHeaders(&check)
+		p, err = httpProber.NewProber(ctx, check.Check, logger, reservedHeaders)
 		target = check.Target
 
 	case sm.CheckTypeDns:
@@ -81,7 +85,8 @@ func (f proberFactory) New(ctx context.Context, logger zerolog.Logger, check mod
 
 	case sm.CheckTypeMultiHttp:
 		if f.runner != nil {
-			p, err = multihttp.NewProber(ctx, check.Check, logger, f.runner)
+			reservedHeaders := f.getReservedHeaders(&check)
+			p, err = multihttp.NewProber(ctx, check.Check, logger, f.runner, reservedHeaders)
 			target = check.Target
 		} else {
 			err = fmt.Errorf("k6 checks are not enabled")
@@ -97,3 +102,14 @@ func (f proberFactory) New(ctx context.Context, logger zerolog.Logger, check mod
 
 	return p, target, err
 }
+
+// Build reserved HTTP request headers for applicable checks.
+func (f proberFactory) getReservedHeaders(check *model.Check) http.Header {
+	reservedHeaders := http.Header{}
+	if f.probeId != 0 {
+		checkProbeIdentifier := fmt.Sprintf("%d-%d", check.GlobalID(), f.probeId)
+		reservedHeaders.Add("x-sm-id", checkProbeIdentifier)
+	}
+
+	return reservedHeaders
+}

internal/scraper/scraper.go

@@ -120,7 +120,7 @@ func New(ctx context.Context, check model.Check, publisher pusher.Publisher, pro
 		Logger:        logger,
 		ScrapeCounter: scrapeCounter,
 		ErrorCounter:  errorCounter,
-		ProbeFactory:  prober.NewProberFactory(k6runner),
+		ProbeFactory:  prober.NewProberFactory(k6runner, probe.Id),
 		LabelsLimiter: labelsLimiter,
 	})
 }