Support custom HTTP headers when reporting spans over HTTP (#479)

Signed-off-by: Albert Teoh <[email protected]>

config/config.go

@@ -134,6 +134,10 @@ type ReporterConfig struct {
 	// Password instructs reporter to include a password for basic http authentication when sending spans to
 	// jaeger-collector. Can be set by exporting an environment variable named JAEGER_PASSWORD
 	Password string `yaml:"password"`
+
+	// HTTPHeaders instructs the reporter to add these headers to the http request when reporting spans.
+	// This field takes effect only when using HTTPTransport by setting the CollectorEndpoint.
+	HTTPHeaders map[string]string `yaml:"http_headers"`
 }
 
 // BaggageRestrictionsConfig configures the baggage restrictions manager which can be used to whitelist
@@ -397,11 +401,12 @@ func (rc *ReporterConfig) NewReporter(
 
 func (rc *ReporterConfig) newTransport() (jaeger.Transport, error) {
 	switch {
-	case rc.CollectorEndpoint != "" && rc.User != "" && rc.Password != "":
-		return transport.NewHTTPTransport(rc.CollectorEndpoint, transport.HTTPBatchSize(1),
-			transport.HTTPBasicAuth(rc.User, rc.Password)), nil
 	case rc.CollectorEndpoint != "":
-		return transport.NewHTTPTransport(rc.CollectorEndpoint, transport.HTTPBatchSize(1)), nil
+		httpOptions := []transport.HTTPOption{transport.HTTPBatchSize(1), transport.HTTPHeaders(rc.HTTPHeaders)}
+		if rc.User != "" && rc.Password != "" {
+			httpOptions = append(httpOptions, transport.HTTPBasicAuth(rc.User, rc.Password))
+		}
+		return transport.NewHTTPTransport(rc.CollectorEndpoint, httpOptions...), nil
 	default:
 		return jaeger.NewUDPTransport(rc.LocalAgentHostPort, 0)
 	}

config/config_test.go

@@ -497,6 +497,18 @@ func TestHTTPTransportType(t *testing.T) {
 	require.IsType(t, expect, sender)
 }
 
+func TestHTTPTransportTypeWithAuth(t *testing.T) {
+	rc := &ReporterConfig{
+		CollectorEndpoint: "http://1.2.3.4:5678/api/traces",
+		User:              "auth_user",
+		Password:          "auth_pass",
+	}
+	expect := transport.NewHTTPTransport(rc.CollectorEndpoint)
+	sender, err := rc.newTransport()
+	require.NoError(t, err)
+	require.IsType(t, expect, sender)
+}
+
 func TestDefaultConfig(t *testing.T) {
 	cfg := Configuration{}
 	_, _, err := cfg.NewTracer(Metrics(metrics.NullFactory), Logger(log.NullLogger))

transport/http.go

@@ -39,6 +39,7 @@ type HTTPTransport struct {
 	spans           []*j.Span
 	process         *j.Process
 	httpCredentials *HTTPBasicAuthCredentials
+	headers         map[string]string
 }
 
 // HTTPBasicAuthCredentials stores credentials for HTTP basic auth.
@@ -76,6 +77,13 @@ func HTTPRoundTripper(transport http.RoundTripper) HTTPOption {
 	}
 }
 
+// HTTPHeaders defines the HTTP headers that will be attached to the jaeger client's HTTP request
+func HTTPHeaders(headers map[string]string) HTTPOption {
+	return func(c *HTTPTransport) {
+		c.headers = headers
+	}
+}
+
 // NewHTTPTransport returns a new HTTP-backend transport. url should be an http
 // url of the collector to handle POST request, typically something like:
 //     http://hostname:14268/api/traces?format=jaeger.thrift
@@ -136,6 +144,9 @@ func (c *HTTPTransport) send(spans []*j.Span) error {
 		return err
 	}
 	req.Header.Set("Content-Type", "application/x-thrift")
+	for k, v := range c.headers {
+		req.Header.Set(k, v)
+	}
 
 	if c.httpCredentials != nil {
 		req.SetBasicAuth(c.httpCredentials.username, c.httpCredentials.password)

transport/http_test.go

@@ -36,6 +36,7 @@ func TestHTTPTransport(t *testing.T) {
 		"http://localhost:10001/api/v1/spans",
 		HTTPBatchSize(1),
 		HTTPBasicAuth(httpUsername, httpPassword),
+		HTTPHeaders(map[string]string{"my-key": "my-value"}),
 	)
 
 	tracer, closer := jaeger.NewTracer(
@@ -77,6 +78,7 @@ func TestHTTPTransport(t *testing.T) {
 		&HTTPBasicAuthCredentials{username: httpUsername, password: httpPassword},
 		server.authCredentials[0],
 	)
+	assert.Equal(t, "my-value", server.getHeader().Get("my-key"))
 }
 
 func TestHTTPOptions(t *testing.T) {
@@ -88,17 +90,20 @@ func TestHTTPOptions(t *testing.T) {
 		HTTPBatchSize(123),
 		HTTPTimeout(123*time.Millisecond),
 		HTTPRoundTripper(roundTripper),
+		HTTPHeaders(map[string]string{"my-key": "my-value"}),
 	)
 	assert.Equal(t, 123, sender.batchSize)
 	assert.Equal(t, 123*time.Millisecond, sender.client.Timeout)
 	assert.Equal(t, roundTripper, sender.client.Transport)
+	assert.Equal(t, "my-value", sender.headers["my-key"])
 }
 
 type httpServer struct {
 	t               *testing.T
 	batches         []*j.Batch
 	authCredentials []*HTTPBasicAuthCredentials
 	mutex           sync.RWMutex
+	header          http.Header
 }
 
 func (s *httpServer) getBatches() []*j.Batch {
@@ -113,6 +118,12 @@ func (s *httpServer) credentials() []*HTTPBasicAuthCredentials {
 	return s.authCredentials
 }
 
+func (s *httpServer) getHeader() http.Header {
+	s.mutex.RLock()
+	defer s.mutex.RUnlock()
+	return s.header
+}
+
 // TODO this function and zipkin/http_test.go#newHTTPServer look like twins lost at birth
 func newHTTPServer(t *testing.T) *httpServer {
 	server := &httpServer{
@@ -151,6 +162,8 @@ func newHTTPServer(t *testing.T) *httpServer {
 		server.batches = append(server.batches, batch)
 		u, p, _ := r.BasicAuth()
 		server.authCredentials = append(server.authCredentials, &HTTPBasicAuthCredentials{username: u, password: p})
+
+		server.header = r.Header
 	})
 
 	go func() {