contrib/net/http: refactor tracing (#2921)

Co-authored-by: Dario Castañé <[email protected]>

contrib/internal/httptrace/before_handle.go

@@ -0,0 +1,78 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package httptrace
+
+import (
+	"net/http"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/contrib/internal/options"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/emitter/httpsec"
+)
+
+// ServeConfig specifies the tracing configuration when using TraceAndServe.
+type ServeConfig struct {
+	// Service specifies the service name to use. If left blank, the global service name
+	// will be inherited.
+	Service string
+	// Resource optionally specifies the resource name for this request.
+	Resource string
+	// QueryParams should be true in order to append the URL query values to the  "http.url" tag.
+	QueryParams bool
+	// Route is the request matched route if any, or is empty otherwise
+	Route string
+	// RouteParams specifies framework-specific route parameters (e.g. for route /user/:id coming
+	// in as /user/123 we'll have {"id": "123"}). This field is optional and is used for monitoring
+	// by AppSec. It is only taken into account when AppSec is enabled.
+	RouteParams map[string]string
+	// FinishOpts specifies any options to be used when finishing the request span.
+	FinishOpts []ddtrace.FinishOption
+	// SpanOpts specifies any options to be applied to the request starting span.
+	SpanOpts []ddtrace.StartSpanOption
+}
+
+// BeforeHandle contains functionality that should be executed before a http.Handler runs.
+// It returns the "traced" http.ResponseWriter and http.Request, an additional afterHandle function
+// that should be executed after the Handler runs, and a handled bool that instructs if the request has been handled
+// or not - in case it was handled, the original handler should not run.
+func BeforeHandle(cfg *ServeConfig, w http.ResponseWriter, r *http.Request) (http.ResponseWriter, *http.Request, func(), bool) {
+	if cfg == nil {
+		cfg = new(ServeConfig)
+	}
+	opts := options.Copy(cfg.SpanOpts...) // make a copy of cfg.SpanOpts to avoid races
+	if cfg.Service != "" {
+		opts = append(opts, tracer.ServiceName(cfg.Service))
+	}
+	if cfg.Resource != "" {
+		opts = append(opts, tracer.ResourceName(cfg.Resource))
+	}
+	if cfg.Route != "" {
+		opts = append(opts, tracer.Tag(ext.HTTPRoute, cfg.Route))
+	}
+	span, ctx := StartRequestSpan(r, opts...)
+	rw, ddrw := wrapResponseWriter(w)
+	rt := r.WithContext(ctx)
+
+	closeSpan := func() {
+		FinishRequestSpan(span, ddrw.status, cfg.FinishOpts...)
+	}
+	afterHandle := closeSpan
+	handled := false
+	if appsec.Enabled() {
+		secW, secReq, secAfterHandle, secHandled := httpsec.BeforeHandle(rw, rt, span, cfg.RouteParams, nil)
+		afterHandle = func() {
+			secAfterHandle()
+			closeSpan()
+		}
+		rw = secW
+		rt = secReq
+		handled = secHandled
+	}
+	return rw, rt, afterHandle, handled
+}

contrib/internal/httptrace/response_writer.go

@@ -0,0 +1,51 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package httptrace
+
+//go:generate sh -c "go run make_responsewriter.go | gofmt > trace_gen.go"
+
+import "net/http"
+
+// responseWriter is a small wrapper around an http response writer that will
+// intercept and store the status of a request.
+type responseWriter struct {
+	http.ResponseWriter
+	status int
+}
+
+func newResponseWriter(w http.ResponseWriter) *responseWriter {
+	return &responseWriter{w, 0}
+}
+
+// Status returns the status code that was monitored.
+func (w *responseWriter) Status() int {
+	return w.status
+}
+
+// Write writes the data to the connection as part of an HTTP reply.
+// We explicitly call WriteHeader with the 200 status code
+// in order to get it reported into the span.
+func (w *responseWriter) Write(b []byte) (int, error) {
+	if w.status == 0 {
+		w.WriteHeader(http.StatusOK)
+	}
+	return w.ResponseWriter.Write(b)
+}
+
+// WriteHeader sends an HTTP response header with status code.
+// It also sets the status code to the span.
+func (w *responseWriter) WriteHeader(status int) {
+	if w.status != 0 {
+		return
+	}
+	w.ResponseWriter.WriteHeader(status)
+	w.status = status
+}
+
+// Unwrap returns the underlying wrapped http.ResponseWriter.
+func (w *responseWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}

contrib/internal/httptrace/response_writer_test.go

@@ -0,0 +1,36 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package httptrace
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_wrapResponseWriter(t *testing.T) {
+	// there doesn't appear to be an easy way to test http.Pusher support via an http request
+	// so we'll just confirm wrapResponseWriter preserves it
+	t.Run("Pusher", func(t *testing.T) {
+		var i struct {
+			http.ResponseWriter
+			http.Pusher
+		}
+		var w http.ResponseWriter = i
+		_, ok := w.(http.ResponseWriter)
+		assert.True(t, ok)
+		_, ok = w.(http.Pusher)
+		assert.True(t, ok)
+
+		w, _ = wrapResponseWriter(w)
+		_, ok = w.(http.ResponseWriter)
+		assert.True(t, ok)
+		_, ok = w.(http.Pusher)
+		assert.True(t, ok)
+	})
+
+}

contrib/julienschmidt/httprouter/httprouter.go

@@ -7,68 +7,66 @@
 package httprouter // import "gopkg.in/DataDog/dd-trace-go.v1/contrib/julienschmidt/httprouter"
 
 import (
-	"math"
 	"net/http"
-	"strings"
-
-	httptraceinternal "gopkg.in/DataDog/dd-trace-go.v1/contrib/internal/httptrace"
-	"gopkg.in/DataDog/dd-trace-go.v1/contrib/internal/options"
-	httptrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/net/http"
-	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
-	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
-	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
-	"gopkg.in/DataDog/dd-trace-go.v1/internal/telemetry"
 
 	"github.com/julienschmidt/httprouter"
-)
-
-const componentName = "julienschmidt/httprouter"
 
-func init() {
-	telemetry.LoadIntegration(componentName)
-	tracer.MarkIntegrationImported("github.com/julienschmidt/httprouter")
-}
+	"gopkg.in/DataDog/dd-trace-go.v1/contrib/julienschmidt/httprouter/internal/tracing"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+)
 
 // Router is a traced version of httprouter.Router.
 type Router struct {
 	*httprouter.Router
-	config *routerConfig
+	config *tracing.Config
 }
 
 // New returns a new router augmented with tracing.
 func New(opts ...RouterOption) *Router {
-	cfg := new(routerConfig)
-	defaults(cfg)
-	for _, fn := range opts {
-		fn(cfg)
-	}
-	if !math.IsNaN(cfg.analyticsRate) {
-		cfg.spanOpts = append(cfg.spanOpts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
-	}
-
-	cfg.spanOpts = append(cfg.spanOpts, tracer.Tag(ext.SpanKind, ext.SpanKindServer))
-	cfg.spanOpts = append(cfg.spanOpts, tracer.Tag(ext.Component, componentName))
-
+	cfg := tracing.NewConfig(opts...)
 	log.Debug("contrib/julienschmidt/httprouter: Configuring Router: %#v", cfg)
 	return &Router{httprouter.New(), cfg}
 }
 
 // ServeHTTP implements http.Handler.
 func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
-	// get the resource associated to this request
-	route := req.URL.Path
-	_, ps, _ := r.Router.Lookup(req.Method, route)
-	for _, param := range ps {
-		route = strings.Replace(route, param.Value, ":"+param.Key, 1)
+	tw, treq, afterHandle, handled := tracing.BeforeHandle(r.config, r.Router, wrapRouter, w, req)
+	defer afterHandle()
+	if handled {
+		return
+	}
+	r.Router.ServeHTTP(tw, treq)
+}
+
+type wRouter struct {
+	*httprouter.Router
+}
+
+func wrapRouter(r *httprouter.Router) tracing.Router {
+	return &wRouter{r}
+}
+
+func (w wRouter) Lookup(method string, path string) (any, []tracing.Param, bool) {
+	h, params, ok := w.Router.Lookup(method, path)
+	return h, wrapParams(params), ok
+}
+
+type wParam struct {
+	httprouter.Param
+}
+
+func wrapParams(params httprouter.Params) []tracing.Param {
+	wParams := make([]tracing.Param, len(params))
+	for i, p := range params {
+		wParams[i] = wParam{p}
 	}
-	resource := req.Method + " " + route
-	spanOpts := options.Copy(r.config.spanOpts...) // spanOpts must be a copy of r.config.spanOpts, locally scoped, to avoid races.
-	spanOpts = append(spanOpts, httptraceinternal.HeaderTagsFromRequest(req, r.config.headerTags))
+	return wParams
+}
+
+func (w wParam) GetKey() string {
+	return w.Key
+}
 
-	httptrace.TraceAndServe(r.Router, w, req, &httptrace.ServeConfig{
-		Service:  r.config.serviceName,
-		Resource: resource,
-		SpanOpts: spanOpts,
-		Route:    route,
-	})
+func (w wParam) GetValue() string {
+	return w.Value
 }