From 26a7b59f6f9b0de0090875c8aa3340fe75e53100 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Oliveirinha?= <joliveirinha@cloudflare.com>
Date: Wed, 18 May 2022 11:11:48 +0100
Subject: [PATCH] TUN-6248: Fix panic in cloudflared during tracing when origin
 doesn't provide header map

---
 proxy/proxy.go          |  5 +++++
 tracing/tracing.go      |  6 ++++++
 tracing/tracing_test.go | 19 +++++++++++++++++++
 3 files changed, 30 insertions(+)

diff --git a/proxy/proxy.go b/proxy/proxy.go
index bcdab8de..32e3d678 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -203,6 +203,11 @@ func (p *Proxy) proxyHTTPRequest(
 	tracing.EndWithStatus(ttfbSpan, codes.Ok, resp.Status)
 	defer resp.Body.Close()
 
+	// resp headers can be nil
+	if resp.Header == nil {
+		resp.Header = make(http.Header)
+	}
+
 	// Add spans to response header (if available)
 	tr.AddSpans(resp.Header, p.log)
 
diff --git a/tracing/tracing.go b/tracing/tracing.go
index 2cab213b..927798b3 100644
--- a/tracing/tracing.go
+++ b/tracing/tracing.go
@@ -103,6 +103,11 @@ func (cft *TracedRequest) Tracer() trace.Tracer {
 
 // Spans returns the spans as base64 encoded protobuf otlp traces.
 func (cft *TracedRequest) AddSpans(headers http.Header, log *zerolog.Logger) {
+	if headers == nil {
+		log.Error().Msgf("provided headers map is nil")
+		return
+	}
+
 	enc, err := cft.exporter.Spans()
 	switch err {
 	case nil:
@@ -121,6 +126,7 @@ func (cft *TracedRequest) AddSpans(headers http.Header, log *zerolog.Logger) {
 		log.Error().Msgf("no traces provided and no error from exporter")
 		return
 	}
+
 	headers[CanonicalCloudflaredTracingHeader] = []string{enc}
 }
 
diff --git a/tracing/tracing_test.go b/tracing/tracing_test.go
index 1bd23e89..68a272fc 100644
--- a/tracing/tracing_test.go
+++ b/tracing/tracing_test.go
@@ -1,13 +1,16 @@
 package tracing
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
+	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/assert"
 	tracesdk "go.opentelemetry.io/otel/sdk/trace"
 	"go.opentelemetry.io/otel/trace"
+	tracepb "go.opentelemetry.io/proto/otlp/trace/v1"
 )
 
 func TestNewCfTracer(t *testing.T) {
@@ -48,3 +51,19 @@ func TestNewCfTracerInvalidHeaders(t *testing.T) {
 		assert.IsType(t, &NoopOtlpClient{}, tr.exporter)
 	}
 }
+
+func TestAddingSpansWithNilMap(t *testing.T) {
+	req := httptest.NewRequest("GET", "http://localhost", nil)
+	req.Header.Add(TracerContextName, "14cb070dde8e51fc5ae8514e69ba42ca:b38f1bf5eae406f3:0:1")
+	tr := NewTracedRequest(req)
+
+	exporter := tr.exporter.(*InMemoryOtlpClient)
+
+	// add fake spans
+	spans := createResourceSpans([]*tracepb.Span{createOtlpSpan(traceId)})
+	err := exporter.UploadTraces(context.Background(), spans)
+	assert.NoError(t, err)
+
+	// a panic shouldn't occur
+	tr.AddSpans(nil, &zerolog.Logger{})
+}