From f81b0ee9e85504861ee2b6d43e71f6174e1b6968 Mon Sep 17 00:00:00 2001
From: Devin Carr <dcarr@cloudflare.com>
Date: Mon, 11 Apr 2022 12:57:50 -0700
Subject: [PATCH] TUN-5990: Add otlp span export to response header

---
 connection/http2.go |  6 ++++++
 proxy/proxy.go      | 15 +++++++++------
 tracing/client.go   |  5 +++--
 tracing/tracing.go  | 30 ++++++++++++++++++++++++++----
 4 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/connection/http2.go b/connection/http2.go
index 67c67862..715c827c 100644
--- a/connection/http2.go
+++ b/connection/http2.go
@@ -231,6 +231,12 @@ func (rp *http2RespWriter) WriteRespHeaders(status int, header http.Header) erro
 			dest[name] = values
 		}
 
+		if h2name == tracing.IntCloudflaredTracingHeader {
+			// Add cf-int-cloudflared-tracing header outside of serialized userHeaders
+			rp.w.Header()[tracing.CanonicalCloudflaredTracingHeader] = values
+			continue
+		}
+
 		if !IsControlResponseHeader(h2name) || IsWebsocketClientHeader(h2name) {
 			// User headers, on the other hand, must all be serialized so that
 			// HTTP/2 header validation won't be applied to HTTP/1 header values
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 76e11fe1..f24f0290 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -87,7 +87,7 @@ func (p *Proxy) ProxyHTTP(
 	case ingress.HTTPOriginProxy:
 		if err := p.proxyHTTPRequest(
 			w,
-			req,
+			tr,
 			originProxy,
 			isWebsocket,
 			rule.Config.DisableChunkedEncoding,
@@ -159,15 +159,15 @@ func ruleField(ing ingress.Ingress, ruleNum int) (ruleID string, srv string) {
 // ProxyHTTPRequest proxies requests of underlying type http and websocket to the origin service.
 func (p *Proxy) proxyHTTPRequest(
 	w connection.ResponseWriter,
-	req *http.Request,
+	tr *tracing.TracedRequest,
 	httpService ingress.HTTPOriginProxy,
 	isWebsocket bool,
 	disableChunkedEncoding bool,
 	fields logFields,
 ) error {
-	roundTripReq := req
+	roundTripReq := tr.Request
 	if isWebsocket {
-		roundTripReq = req.Clone(req.Context())
+		roundTripReq = tr.Request.Clone(tr.Request.Context())
 		roundTripReq.Header.Set("Connection", "Upgrade")
 		roundTripReq.Header.Set("Upgrade", "websocket")
 		roundTripReq.Header.Set("Sec-Websocket-Version", "13")
@@ -177,7 +177,7 @@ func (p *Proxy) proxyHTTPRequest(
 		// Support for WSGI Servers by switching transfer encoding from chunked to gzip/deflate
 		if disableChunkedEncoding {
 			roundTripReq.TransferEncoding = []string{"gzip", "deflate"}
-			cLength, err := strconv.Atoi(req.Header.Get("Content-Length"))
+			cLength, err := strconv.Atoi(tr.Request.Header.Get("Content-Length"))
 			if err == nil {
 				roundTripReq.ContentLength = int64(cLength)
 			}
@@ -197,6 +197,9 @@ func (p *Proxy) proxyHTTPRequest(
 	}
 	defer resp.Body.Close()
 
+	// Add spans to response header (if available)
+	tr.AddSpans(resp.Header, p.log)
+
 	err = w.WriteRespHeaders(resp.StatusCode, resp.Header)
 	if err != nil {
 		return errors.Wrap(err, "Error writing response header")
@@ -211,7 +214,7 @@ func (p *Proxy) proxyHTTPRequest(
 
 		eyeballStream := &bidirectionalStream{
 			writer: w,
-			reader: req.Body,
+			reader: tr.Request.Body,
 		}
 
 		websocket.Stream(eyeballStream, rwc, p.log)
diff --git a/tracing/client.go b/tracing/client.go
index 6783912e..5ed4efb3 100644
--- a/tracing/client.go
+++ b/tracing/client.go
@@ -16,7 +16,8 @@ const (
 )
 
 var (
-	errNoTraces = errors.New("no traces recorded to be exported")
+	errNoTraces   = errors.New("no traces recorded to be exported")
+	errNoopTracer = errors.New("noop tracer has no traces")
 )
 
 type InMemoryClient interface {
@@ -86,5 +87,5 @@ func (mc *NoopOtlpClient) UploadTraces(_ context.Context, _ []*tracepb.ResourceS
 
 // Spans always returns no traces error
 func (mc *NoopOtlpClient) Spans() (string, error) {
-	return "", errNoTraces
+	return "", errNoopTracer
 }
diff --git a/tracing/tracing.go b/tracing/tracing.go
index a35d2962..f124e367 100644
--- a/tracing/tracing.go
+++ b/tracing/tracing.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"net/http"
 
+	"github.com/rs/zerolog"
 	otelContrib "go.opentelemetry.io/contrib/propagators/Jaeger"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
@@ -23,11 +24,14 @@ const (
 
 	tracerContextName         = "cf-trace-id"
 	tracerContextNameOverride = "uber-trace-id"
+
+	IntCloudflaredTracingHeader = "cf-int-cloudflared-tracing"
 )
 
 var (
-	Http2TransportAttribute = trace.WithAttributes(TransportAttributeKey.String("http2"))
-	QuicTransportAttribute  = trace.WithAttributes(TransportAttributeKey.String("quic"))
+	CanonicalCloudflaredTracingHeader = http.CanonicalHeaderKey(IntCloudflaredTracingHeader)
+	Http2TransportAttribute           = trace.WithAttributes(TransportAttributeKey.String("http2"))
+	QuicTransportAttribute            = trace.WithAttributes(TransportAttributeKey.String("quic"))
 
 	TransportAttributeKey = attribute.Key("transport")
 	TrafficAttributeKey   = attribute.Key("traffic")
@@ -75,8 +79,26 @@ func (cft *TracedRequest) Tracer() trace.Tracer {
 }
 
 // Spans returns the spans as base64 encoded protobuf otlp traces.
-func (cft *TracedRequest) Spans() (string, error) {
-	return cft.exporter.Spans()
+func (cft *TracedRequest) AddSpans(headers http.Header, log *zerolog.Logger) {
+	enc, err := cft.exporter.Spans()
+	switch err {
+	case nil:
+		break
+	case errNoTraces:
+		log.Error().Err(err).Msgf("expected traces to be available")
+		return
+	case errNoopTracer:
+		return // noop tracer has no traces
+	default:
+		log.Error().Err(err)
+		return
+	}
+	// No need to add header if no traces
+	if enc == "" {
+		log.Error().Msgf("no traces provided and no error from exporter")
+		return
+	}
+	headers[CanonicalCloudflaredTracingHeader] = []string{enc}
 }
 
 // EndWithStatus will set a status for the span and then end it.