Add bool param to middleware for decoding

middleware.go

@@ -17,6 +17,7 @@
 package candlelight
 
 import (
+	"context"
 	"crypto/rand"
 	"encoding/base64"
 	"net/http"
@@ -62,34 +63,42 @@ func (traceConfig *TraceConfig) TraceMiddleware(delegate http.Handler) http.Hand
 	})
 }
 
-// EchoFirstNodeTraceInfo captures the trace information from a request and writes it
-// back in the response headers if the request is the first one in the trace path.
-func EchoFirstTraceNodeInfo(propagator propagation.TextMapPropagator) func(http.Handler) http.Handler {
+// EchoFirstNodeTraceInfo captures the trace information from a request, writes it
+// back in the response headers, and adds it to the request's context
+// It can also decode the request and save the resulting WRP object in the context if isDecodable is true
+func EchoFirstTraceNodeInfo(propagator propagation.TextMapPropagator, isDecodable bool) func(http.Handler) http.Handler {
 	return func(delegate http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 
-			if req, err := wrphttp.DecodeRequest(r, nil); err == nil {
-				r = req
-			}
+			var ctx context.Context
 
-			var traceHeaders []string
-			ctx := propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
-			if msg, ok := wrpcontext.Get[*wrp.Message](ctx); ok {
-				traceHeaders = msg.Headers
-			}
+			if isDecodable {
+				if req, err := wrphttp.DecodeRequest(r, nil); err == nil {
+					r = req
+				}
 
-			// Iterate through the trace headers (if any), format them, and add them to ctx
-			var tmp propagation.TextMapCarrier = propagation.MapCarrier{}
-			for _, f := range traceHeaders {
-				if f != "" {
-					parts := strings.Split(f, ":")
-					// Remove leading space if there's any
-					parts[1] = strings.Trim(parts[1], " ")
-					tmp.Set(parts[0], parts[1])
+				var traceHeaders []string
+				ctx = propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+				if msg, ok := wrpcontext.Get[*wrp.Message](ctx); ok {
+					traceHeaders = msg.Headers
 				}
+
+				// Iterate through the trace headers (if any), format them, and add them to ctx
+				var tmp propagation.TextMapCarrier = propagation.MapCarrier{}
+				for _, f := range traceHeaders {
+					if f != "" {
+						parts := strings.Split(f, ":")
+						// Remove leading space if there's any
+						parts[1] = strings.Trim(parts[1], " ")
+						tmp.Set(parts[0], parts[1])
+					}
+				}
+
+				ctx = propagation.TraceContext{}.Extract(ctx, tmp)
+			} else {
+				ctx = propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
 			}
 
-			ctx = propagation.TraceContext{}.Extract(ctx, tmp)
 			sc := trace.SpanContextFromContext(ctx)
 			if sc.IsValid() {
 				w.Header().Set(spanIDHeaderName, sc.SpanID().String())