From 6737d670d38e2c3c73b6f36ec3e10bf260719964 Mon Sep 17 00:00:00 2001 From: Serena Zamarripa Date: Mon, 5 Jun 2023 11:37:14 -0500 Subject: [PATCH 1/3] Add bool param to middleware for decoding --- middleware.go | 49 +++++++++++++++++++++++++++++-------------------- 1 file changed, 29 insertions(+), 20 deletions(-) diff --git a/middleware.go b/middleware.go index d74c376..6011285 100644 --- a/middleware.go +++ b/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()) From 9da3452feb85ee36756c969fa833c0986f61dfa8 Mon Sep 17 00:00:00 2001 From: Serena Zamarripa Date: Tue, 13 Jun 2023 15:41:07 -0500 Subject: [PATCH 2/3] Add Check for traceheader, ensure tracing --- middleware.go | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/middleware.go b/middleware.go index 6011285..ee52f62 100644 --- a/middleware.go +++ b/middleware.go @@ -76,29 +76,31 @@ func EchoFirstTraceNodeInfo(propagator propagation.TextMapPropagator, isDecodabl if req, err := wrphttp.DecodeRequest(r, nil); err == nil { r = req } + } - var traceHeaders []string - ctx = propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header)) - if msg, ok := wrpcontext.Get[*wrp.Message](ctx); ok { - traceHeaders = msg.Headers - } + var traceHeaders []string + ctx = propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header)) + if msg, ok := wrpcontext.Get[*wrp.Message](ctx); ok { + traceHeaders = msg.Headers + } else if headers := r.Header.Values("X-midt-Headers"); len(headers) != 0 { + traceHeaders = 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, ":") + // 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, ":") + if len(parts) > 1 { // 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()) From 8bc6dc32d806ecdb46e35d7e57e1919a8638ab6d Mon Sep 17 00:00:00 2001 From: Serena Zamarripa Date: Tue, 13 Jun 2023 15:49:13 -0500 Subject: [PATCH 3/3] Correct Xmidt header --- middleware.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/middleware.go b/middleware.go index ee52f62..90f48a7 100644 --- a/middleware.go +++ b/middleware.go @@ -82,7 +82,7 @@ func EchoFirstTraceNodeInfo(propagator propagation.TextMapPropagator, isDecodabl ctx = propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header)) if msg, ok := wrpcontext.Get[*wrp.Message](ctx); ok { traceHeaders = msg.Headers - } else if headers := r.Header.Values("X-midt-Headers"); len(headers) != 0 { + } else if headers := r.Header.Values("X-Xmidt-Headers"); len(headers) != 0 { traceHeaders = headers }