contrib/internal/httptrace: add support for inferred proxy spans (#3052)

Co-authored-by: Zarir Hamza <[email protected]>
Co-authored-by: Rodrigo Arguello <[email protected]>

contrib/emicklei/go-restful.v3/restful.go

@@ -46,9 +46,9 @@ func FilterFunc(configOpts ...Option) restful.FilterFunction {
 			spanOpts = append(spanOpts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
 		}
 		spanOpts = append(spanOpts, httptrace.HeaderTagsFromRequest(req.Request, cfg.headerTags))
-		span, ctx := httptrace.StartRequestSpan(req.Request, spanOpts...)
+		_, ctx, finishSpans := httptrace.StartRequestSpan(req.Request, spanOpts...)
 		defer func() {
-			httptrace.FinishRequestSpan(span, resp.StatusCode(), nil, tracer.WithError(resp.Error()))
+			finishSpans(resp.StatusCode(), nil, tracer.WithError(resp.Error()))
 		}()
 
 		// pass the span through the request context

contrib/emicklei/go-restful/restful.go

@@ -46,9 +46,9 @@ func FilterFunc(configOpts ...Option) restful.FilterFunction {
 			spanOpts = append(spanOpts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
 		}
 		spanOpts = append(spanOpts, httptrace.HeaderTagsFromRequest(req.Request, cfg.headerTags))
-		span, ctx := httptrace.StartRequestSpan(req.Request, spanOpts...)
+		_, ctx, finishSpans := httptrace.StartRequestSpan(req.Request, spanOpts...)
 		defer func() {
-			httptrace.FinishRequestSpan(span, resp.StatusCode(), nil, tracer.WithError(resp.Error()))
+			finishSpans(resp.StatusCode(), nil, tracer.WithError(resp.Error()))
 		}()
 
 		// pass the span through the request context
@@ -59,9 +59,9 @@ func FilterFunc(configOpts ...Option) restful.FilterFunction {
 
 // Filter is deprecated. Please use FilterFunc.
 func Filter(req *restful.Request, resp *restful.Response, chain *restful.FilterChain) {
-	span, ctx := httptrace.StartRequestSpan(req.Request, tracer.ResourceName(req.SelectedRoutePath()))
+	_, ctx, finishSpans := httptrace.StartRequestSpan(req.Request, tracer.ResourceName(req.SelectedRoutePath()))
 	defer func() {
-		httptrace.FinishRequestSpan(span, resp.StatusCode(), nil, tracer.WithError(resp.Error()))
+		finishSpans(resp.StatusCode(), nil, tracer.WithError(resp.Error()))
 	}()
 
 	// pass the span through the request context

contrib/gin-gonic/gin/gintrace.go

@@ -52,9 +52,9 @@ func Middleware(service string, opts ...Option) gin.HandlerFunc {
 		}
 		opts = append(opts, tracer.Tag(ext.HTTPRoute, c.FullPath()))
 		opts = append(opts, httptrace.HeaderTagsFromRequest(c.Request, cfg.headerTags))
-		span, ctx := httptrace.StartRequestSpan(c.Request, opts...)
+		span, ctx, finishSpans := httptrace.StartRequestSpan(c.Request, opts...)
 		defer func() {
-			httptrace.FinishRequestSpan(span, c.Writer.Status(), nil)
+			finishSpans(c.Writer.Status(), nil)
 		}()
 
 		// pass the span through the request context

contrib/go-chi/chi.v5/chi.go

@@ -51,11 +51,11 @@ func Middleware(opts ...Option) func(next http.Handler) http.Handler {
 				opts = append(opts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
 			}
 			opts = append(opts, httptrace.HeaderTagsFromRequest(r, cfg.headerTags))
-			span, ctx := httptrace.StartRequestSpan(r, opts...)
+			span, ctx, finishSpans := httptrace.StartRequestSpan(r, opts...)
 			ww := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
 			defer func() {
 				status := ww.Status()
-				httptrace.FinishRequestSpan(span, status, cfg.isStatusError)
+				finishSpans(status, cfg.isStatusError)
 			}()
 
 			// pass the span through the request context

contrib/go-chi/chi/chi.go

@@ -51,11 +51,11 @@ func Middleware(opts ...Option) func(next http.Handler) http.Handler {
 				opts = append(opts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
 			}
 			opts = append(opts, httptrace.HeaderTagsFromRequest(r, cfg.headerTags))
-			span, ctx := httptrace.StartRequestSpan(r, opts...)
+			span, ctx, finishSpans := httptrace.StartRequestSpan(r, opts...)
 			ww := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
 			defer func() {
 				status := ww.Status()
-				httptrace.FinishRequestSpan(span, status, cfg.isStatusError)
+				finishSpans(status, cfg.isStatusError)
 			}()
 
 			// pass the span through the request context

contrib/internal/httptrace/before_handle.go

@@ -57,11 +57,11 @@ func BeforeHandle(cfg *ServeConfig, w http.ResponseWriter, r *http.Request) (htt
 	if cfg.Route != "" {
 		opts = append(opts, tracer.Tag(ext.HTTPRoute, cfg.Route))
 	}
-	span, ctx := StartRequestSpan(r, opts...)
+	span, ctx, finishSpans := StartRequestSpan(r, opts...)
 	rw, ddrw := wrapResponseWriter(w)
 	rt := r.WithContext(ctx)
 	closeSpan := func() {
-		FinishRequestSpan(span, ddrw.status, cfg.IsStatusError, cfg.FinishOpts...)
+		finishSpans(ddrw.status, cfg.IsStatusError, cfg.FinishOpts...)
 	}
 	afterHandle := closeSpan
 	handled := false

contrib/internal/httptrace/config.go

@@ -26,16 +26,19 @@ const (
 	envTraceClientIPEnabled = "DD_TRACE_CLIENT_IP_ENABLED"
 	// envServerErrorStatuses is the name of the env var used to specify error status codes on http server spans
 	envServerErrorStatuses = "DD_TRACE_HTTP_SERVER_ERROR_STATUSES"
+	// envInferredProxyServicesEnabled is the name of the env var used for enabling inferred span tracing
+	envInferredProxyServicesEnabled = "DD_TRACE_INFERRED_PROXY_SERVICES_ENABLED"
 )
 
 // defaultQueryStringRegexp is the regexp used for query string obfuscation if [EnvQueryStringRegexp] is empty.
 var defaultQueryStringRegexp = regexp.MustCompile("(?i)(?:p(?:ass)?w(?:or)?d|pass(?:_?phrase)?|secret|(?:api_?|private_?|public_?|access_?|secret_?)key(?:_?id)?|token|consumer_?(?:id|key|secret)|sign(?:ed|ature)?|auth(?:entication|orization)?)(?:(?:\\s|%20)*(?:=|%3D)[^&]+|(?:\"|%22)(?:\\s|%20)*(?::|%3A)(?:\\s|%20)*(?:\"|%22)(?:%2[^2]|%[^2]|[^\"%])+(?:\"|%22))|bearer(?:\\s|%20)+[a-z0-9\\._\\-]|token(?::|%3A)[a-z0-9]{13}|gh[opsu]_[0-9a-zA-Z]{36}|ey[I-L](?:[\\w=-]|%3D)+\\.ey[I-L](?:[\\w=-]|%3D)+(?:\\.(?:[\\w.+\\/=-]|%3D|%2F|%2B)+)?|[\\-]{5}BEGIN(?:[a-z\\s]|%20)+PRIVATE(?:\\s|%20)KEY[\\-]{5}[^\\-]+[\\-]{5}END(?:[a-z\\s]|%20)+PRIVATE(?:\\s|%20)KEY|ssh-rsa(?:\\s|%20)*(?:[a-z0-9\\/\\.+]|%2F|%5C|%2B){100,}")
 
 type config struct {
-	queryStringRegexp *regexp.Regexp // specifies the regexp to use for query string obfuscation.
-	queryString       bool           // reports whether the query string should be included in the URL span tag.
-	traceClientIP     bool
-	isStatusError     func(statusCode int) bool
+	queryStringRegexp            *regexp.Regexp // specifies the regexp to use for query string obfuscation.
+	queryString                  bool           // reports whether the query string should be included in the URL span tag.
+	traceClientIP                bool
+	isStatusError                func(statusCode int) bool
+	inferredProxyServicesEnabled bool
 }
 
 // ResetCfg sets local variable cfg back to its defaults (mainly useful for testing)
@@ -45,10 +48,11 @@ func ResetCfg() {
 
 func newConfig() config {
 	c := config{
-		queryString:       !internal.BoolEnv(envQueryStringDisabled, false),
-		queryStringRegexp: QueryStringRegexp(),
-		traceClientIP:     internal.BoolEnv(envTraceClientIPEnabled, false),
-		isStatusError:     isServerError,
+		queryString:                  !internal.BoolEnv(envQueryStringDisabled, false),
+		queryStringRegexp:            QueryStringRegexp(),
+		traceClientIP:                internal.BoolEnv(envTraceClientIPEnabled, false),
+		isStatusError:                isServerError,
+		inferredProxyServicesEnabled: internal.BoolEnv(envInferredProxyServicesEnabled, false),
 	}
 	v := os.Getenv(envServerErrorStatuses)
 	if fn := GetErrorCodesFromInput(v); fn != nil {

contrib/internal/httptrace/httptrace.go

@@ -10,6 +10,7 @@ package httptrace
 import (
 	"context"
 	"fmt"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 	"net/http"
 	"strconv"
 	"strings"
@@ -26,17 +27,50 @@ var (
 	cfg = newConfig()
 )
 
+type inferredSpanCreatedCtxKey struct{}
+
+type FinishSpanFunc = func(status int, errorFn func(int) bool, opts ...tracer.FinishOption)
+
 // StartRequestSpan starts an HTTP request span with the standard list of HTTP request span tags (http.method, http.url,
 // http.useragent). Any further span start option can be added with opts.
-func StartRequestSpan(r *http.Request, opts ...ddtrace.StartSpanOption) (tracer.Span, context.Context) {
+func StartRequestSpan(r *http.Request, opts ...ddtrace.StartSpanOption) (tracer.Span, context.Context, FinishSpanFunc) {
 	// Append our span options before the given ones so that the caller can "overwrite" them.
 	// TODO(): rework span start option handling (https://github.com/DataDog/dd-trace-go/issues/1352)
 
 	var ipTags map[string]string
 	if cfg.traceClientIP {
 		ipTags, _ = httpsec.ClientIPTags(r.Header, true, r.RemoteAddr)
 	}
+
 	nopts := make([]ddtrace.StartSpanOption, 0, len(opts)+1+len(ipTags))
+
+	var inferredProxySpan tracer.Span
+
+	if cfg.inferredProxyServicesEnabled {
+		inferredProxySpanCreated := false
+
+		if created, ok := r.Context().Value(inferredSpanCreatedCtxKey{}).(bool); ok {
+			inferredProxySpanCreated = created
+		}
+
+		if !inferredProxySpanCreated {
+			var inferredStartSpanOpts []ddtrace.StartSpanOption
+
+			requestProxyContext, err := extractInferredProxyContext(r.Header)
+			if err != nil {
+				log.Debug(err.Error())
+			} else {
+				spanParentCtx, spanParentErr := tracer.Extract(tracer.HTTPHeadersCarrier(r.Header))
+				if spanParentErr == nil {
+					if spanLinksCtx, spanLinksOk := spanParentCtx.(ddtrace.SpanContextWithLinks); spanLinksOk {
+						inferredStartSpanOpts = append(inferredStartSpanOpts, tracer.WithSpanLinks(spanLinksCtx.SpanLinks()))
+					}
+				}
+				inferredProxySpan = startInferredProxySpan(requestProxyContext, spanParentCtx, inferredStartSpanOpts...)
+			}
+		}
+	}
+
 	nopts = append(nopts,
 		func(ssCfg *ddtrace.StartSpanConfig) {
 			if ssCfg.Tags == nil {
@@ -50,19 +84,38 @@ func StartRequestSpan(r *http.Request, opts ...ddtrace.StartSpanOption) (tracer.
 			if r.Host != "" {
 				ssCfg.Tags["http.host"] = r.Host
 			}
-			if spanctx, err := tracer.Extract(tracer.HTTPHeadersCarrier(r.Header)); err == nil {
-				// If there are span links as a result of context extraction, add them as a StartSpanOption
-				if linksCtx, ok := spanctx.(ddtrace.SpanContextWithLinks); ok && linksCtx.SpanLinks() != nil {
-					tracer.WithSpanLinks(linksCtx.SpanLinks())(ssCfg)
+
+			if inferredProxySpan != nil {
+				tracer.ChildOf(inferredProxySpan.Context())(ssCfg)
+			} else {
+				spanParentCtx, spanParentErr := tracer.Extract(tracer.HTTPHeadersCarrier(r.Header))
+				if spanParentErr == nil {
+					if spanLinksCtx, spanLinksOk := spanParentCtx.(ddtrace.SpanContextWithLinks); spanLinksOk {
+						tracer.WithSpanLinks(spanLinksCtx.SpanLinks())(ssCfg)
+					}
+					tracer.ChildOf(spanParentCtx)(ssCfg)
 				}
-				tracer.ChildOf(spanctx)(ssCfg)
 			}
+
 			for k, v := range ipTags {
 				ssCfg.Tags[k] = v
 			}
 		})
+
 	nopts = append(nopts, opts...)
-	return tracer.StartSpanFromContext(r.Context(), namingschema.OpName(namingschema.HTTPServer), nopts...)
+
+	var requestContext = r.Context()
+	if inferredProxySpan != nil {
+		requestContext = context.WithValue(r.Context(), inferredSpanCreatedCtxKey{}, true)
+	}
+
+	span, ctx := tracer.StartSpanFromContext(requestContext, namingschema.OpName(namingschema.HTTPServer), nopts...)
+	return span, ctx, func(status int, errorFn func(int) bool, opts ...tracer.FinishOption) {
+		FinishRequestSpan(span, status, errorFn, opts...)
+		if inferredProxySpan != nil {
+			FinishRequestSpan(inferredProxySpan, status, errorFn, opts...)
+		}
+	}
 }
 
 // FinishRequestSpan finishes the given HTTP request span and sets the expected response-related tags such as the status