Merge e699ad5 into 62dc8b2

propagation/b3/http.go

@@ -21,6 +21,31 @@ import (
 	"github.com/openzipkin/zipkin-go/propagation"
 )
 
+type InjectOption func(opts *InjectOptions)
+
+type InjectOptions struct {
+	shouldInjectSingleHeader bool
+	shouldInjectMultiHeader  bool
+}
+
+// WithSingleAndMultiHeader allows to include both single and multiple
+// headers in the context injection
+func WithSingleAndMultiHeader() InjectOption {
+	return func(opts *InjectOptions) {
+		opts.shouldInjectSingleHeader = true
+		opts.shouldInjectMultiHeader = true
+	}
+}
+
+// WithSingleHeaderOnly allows to include only single header in the context
+// injection
+func WithSingleHeaderOnly() InjectOption {
+	return func(opts *InjectOptions) {
+		opts.shouldInjectSingleHeader = true
+		opts.shouldInjectMultiHeader = false
+	}
+}
+
 // ExtractHTTP will extract a span.Context from the HTTP Request if found in
 // B3 header format.
 func ExtractHTTP(r *http.Request) propagation.Extractor {
@@ -31,42 +56,72 @@ func ExtractHTTP(r *http.Request) propagation.Extractor {
 			parentSpanIDHeader = r.Header.Get(ParentSpanID)
 			sampledHeader      = r.Header.Get(Sampled)
 			flagsHeader        = r.Header.Get(Flags)
+			singleHeader       = r.Header.Get(Context)
+		)
+
+		var (
+			sc   *model.SpanContext
+			sErr error
+			mErr error
 		)
+		if singleHeader != "" {
+			sc, sErr = ParseSingleHeader(singleHeader)
+			if sErr == nil {
+				return sc, nil
+			}
+		}
 
-		return ParseHeaders(
-			traceIDHeader, spanIDHeader, parentSpanIDHeader, sampledHeader,
-			flagsHeader,
+		sc, mErr = ParseHeaders(
+			traceIDHeader, spanIDHeader, parentSpanIDHeader,
+			sampledHeader, flagsHeader,
 		)
+
+		if mErr != nil && sErr != nil {
+			return nil, sErr
+		}
+
+		return sc, mErr
 	}
 }
 
 // InjectHTTP will inject a span.Context into a HTTP Request
-func InjectHTTP(r *http.Request) propagation.Injector {
+func InjectHTTP(r *http.Request, opts ...InjectOption) propagation.Injector {
+	options := InjectOptions{shouldInjectMultiHeader: true}
+	for _, opt := range opts {
+		opt(&options)
+	}
+
 	return func(sc model.SpanContext) error {
 		if (model.SpanContext{}) == sc {
 			return ErrEmptyContext
 		}
 
-		if sc.Debug {
-			r.Header.Set(Flags, "1")
-		} else if sc.Sampled != nil {
-			// Debug is encoded as X-B3-Flags: 1. Since Debug implies Sampled,
-			// so don't also send "X-B3-Sampled: 1".
-			if *sc.Sampled {
-				r.Header.Set(Sampled, "1")
-			} else {
-				r.Header.Set(Sampled, "0")
+		if options.shouldInjectMultiHeader {
+			if sc.Debug {
+				r.Header.Set(Flags, "1")
+			} else if sc.Sampled != nil {
+				// Debug is encoded as X-B3-Flags: 1. Since Debug implies Sampled,
+				// so don't also send "X-B3-Sampled: 1".
+				if *sc.Sampled {
+					r.Header.Set(Sampled, "1")
+				} else {
+					r.Header.Set(Sampled, "0")
+				}
 			}
-		}
 
-		if !sc.TraceID.Empty() && sc.ID > 0 {
-			r.Header.Set(TraceID, sc.TraceID.String())
-			r.Header.Set(SpanID, sc.ID.String())
-			if sc.ParentID != nil {
-				r.Header.Set(ParentSpanID, sc.ParentID.String())
+			if !sc.TraceID.Empty() && sc.ID > 0 {
+				r.Header.Set(TraceID, sc.TraceID.String())
+				r.Header.Set(SpanID, sc.ID.String())
+				if sc.ParentID != nil {
+					r.Header.Set(ParentSpanID, sc.ParentID.String())
+				}
 			}
 		}
 
+		if options.shouldInjectSingleHeader {
+			r.Header.Set(Context, BuildSingleHeader(sc))
+		}
+
 		return nil
 	}
 }

propagation/b3/http_test.go

@@ -257,7 +257,35 @@ func TestHTTPExtractInvalidParentIDError(t *testing.T) {
 	if want, have := b3.ErrInvalidParentSpanIDHeader, err; want != have {
 		t.Errorf("ExtractHTTP Error want %+v, have %+v", want, have)
 	}
+}
+
+func TestHTTPExtractSingleFailsAndMultipleFallsbackSuccessfully(t *testing.T) {
+	r := newHTTPRequest(t)
+
+	r.Header.Set(b3.Context, "invalid")
+	r.Header.Set(b3.TraceID, "1")
+	r.Header.Set(b3.SpanID, "2")
+
+	_, err := b3.ExtractHTTP(r)()
+
+	if err != nil {
+		t.Errorf("ExtractHTTP Unexpected error %+v", err)
+	}
+}
+
+func TestHTTPExtractSingleFailsAndMultipleFallsbackFailing(t *testing.T) {
+	r := newHTTPRequest(t)
+
+	r.Header.Set(b3.Context, "0000000000000001-0000000000000002-x")
+	r.Header.Set(b3.TraceID, "1")
+	r.Header.Set(b3.SpanID, "2")
+	r.Header.Set(b3.ParentSpanID, "invalid_data")
+
+	_, err := b3.ExtractHTTP(r)()
 
+	if want, have := b3.ErrInvalidSampledByte, err; want != have {
+		t.Errorf("HTTPExtract Error want %+v, have %+v", want, have)
+	}
 }
 
 func TestHTTPInjectEmptyContextError(t *testing.T) {
@@ -336,6 +364,49 @@ func TestHTTPInjectSampledAndDebugTrace(t *testing.T) {
 	}
 }
 
+func TestHTTPInjectWithSingleOnlyHeaders(t *testing.T) {
+	r := newHTTPRequest(t)
+
+	sampled := true
+	sc := model.SpanContext{
+		TraceID: model.TraceID{Low: 1},
+		ID:      model.ID(2),
+		Debug:   true,
+		Sampled: &sampled,
+	}
+
+	b3.InjectHTTP(r, b3.WithSingleHeaderOnly())(sc)
+
+	if want, have := "", r.Header.Get(b3.TraceID); want != have {
+		t.Errorf("TraceID want empty, have %s", have)
+	}
+
+	if want, have := "0000000000000001-0000000000000002-d", r.Header.Get(b3.Context); want != have {
+		t.Errorf("Context want %s, have %s", want, have)
+	}
+}
+func TestHTTPInjectWithBothSingleAndMultipleHeaders(t *testing.T) {
+	r := newHTTPRequest(t)
+
+	sampled := true
+	sc := model.SpanContext{
+		TraceID: model.TraceID{Low: 1},
+		ID:      model.ID(2),
+		Debug:   true,
+		Sampled: &sampled,
+	}
+
+	b3.InjectHTTP(r, b3.WithSingleAndMultiHeader())(sc)
+
+	if want, have := "0000000000000001", r.Header.Get(b3.TraceID); want != have {
+		t.Errorf("Trace ID want %s, have %s", want, have)
+	}
+
+	if want, have := "0000000000000001-0000000000000002-d", r.Header.Get(b3.Context); want != have {
+		t.Errorf("Context want %s, have %s", want, have)
+	}
+}
+
 func newHTTPRequest(t *testing.T) *http.Request {
 	r, err := http.NewRequest("test", "", nil)
 	if err != nil {

propagation/b3/shared.go

@@ -18,14 +18,19 @@ import "errors"
 
 // Common Header Extraction / Injection errors
 var (
+	ErrInvalidSampledByte        = errors.New("invalid B3 Sampled found")
 	ErrInvalidSampledHeader      = errors.New("invalid B3 Sampled header found")
 	ErrInvalidFlagsHeader        = errors.New("invalid B3 Flags header found")
 	ErrInvalidTraceIDHeader      = errors.New("invalid B3 TraceID header found")
 	ErrInvalidSpanIDHeader       = errors.New("invalid B3 SpanID header found")
 	ErrInvalidParentSpanIDHeader = errors.New("invalid B3 ParentSpanID header found")
 	ErrInvalidScope              = errors.New("require either both TraceID and SpanID or none")
 	ErrInvalidScopeParent        = errors.New("ParentSpanID requires both TraceID and SpanID to be available")
+	ErrInvalidScopeParentSingle  = errors.New("ParentSpanID requires TraceID, SpanID and Sampled to be available")
 	ErrEmptyContext              = errors.New("empty request context")
+	ErrInvalidTraceIDValue       = errors.New("invalid B3 TraceID value found")
+	ErrInvalidSpanIDValue        = errors.New("invalid B3 SpanID value found")
+	ErrInvalidParentSpanIDValue  = errors.New("invalid B3 ParentSpanID value found")
 )
 
 // Default B3 Header keys
@@ -35,4 +40,5 @@ const (
 	ParentSpanID = "x-b3-parentspanid"
 	Sampled      = "x-b3-sampled"
 	Flags        = "x-b3-flags"
+	Context      = "b3"
 )

propagation/b3/spancontext.go

@@ -88,3 +88,115 @@ func ParseHeaders(
 
 	return sc, nil
 }
+
+// ParseSingleHeader takes values found from B3 Single Header and tries to reconstruct a
+// SpanContext.
+func ParseSingleHeader(contextHeader string) (*model.SpanContext, error) {
+	if contextHeader == "" {
+		return nil, ErrEmptyContext
+	}
+
+	var (
+		sc       = model.SpanContext{}
+		sampling string
+	)
+
+	headerLen := len(contextHeader)
+
+	if headerLen == 1 {
+		sampling = contextHeader
+	} else if headerLen == 16 || headerLen == 32 {
+		return nil, ErrInvalidScope
+	} else if headerLen >= 16+16+1 {
+		var high, low uint64
+		pos := 0
+		if string(contextHeader[16]) != "-" {
+			// traceID must be 128 bits
+			var err error
+			high, err = strconv.ParseUint(contextHeader[0:16], 16, 64)
+			if err != nil {
+				return nil, ErrInvalidTraceIDValue
+			}
+			pos = 16
+		}
+
+		low, err := strconv.ParseUint(contextHeader[pos+1:pos+16], 16, 64)
+		if err != nil {
+			return nil, ErrInvalidTraceIDValue
+		}
+
+		sc.TraceID = model.TraceID{High: high, Low: low}
+
+		rawID, err := strconv.ParseUint(contextHeader[pos+16+1:pos+16+1+16], 16, 64)
+		if err != nil {
+			return nil, ErrInvalidSpanIDValue
+		}
+
+		sc.ID = model.ID(rawID)
+
+		if headerLen > pos+16+1+16 {
+			if headerLen == pos+16+1+16+1 {
+				return nil, ErrInvalidSampledByte
+			}
+
+			if headerLen == pos+16+1+16+1+1 {
+				sampling = string(contextHeader[pos+16+1+16+1])
+			} else if headerLen == pos+16+1+16+1+16 {
+				return nil, ErrInvalidScopeParentSingle
+			} else if headerLen == pos+16+1+16+1+1+1+16 {
+				sampling = string(contextHeader[pos+16+1+16+1])
+
+				rawParentID, err := strconv.ParseUint(contextHeader[pos+16+1+16+1+1+1:], 16, 64)
+				if err != nil {
+					return nil, ErrInvalidParentSpanIDValue
+				}
+
+				parentID := model.ID(rawParentID)
+				sc.ParentID = &parentID
+			} else {
+				return nil, ErrInvalidParentSpanIDValue
+			}
+		}
+	} else {
+		return nil, ErrInvalidTraceIDValue
+	}
+	switch sampling {
+	case "d":
+		sc.Debug = true
+	case "1":
+		trueVal := true
+		sc.Sampled = &trueVal
+	case "0":
+		falseVal := false
+		sc.Sampled = &falseVal
+	case "":
+	default:
+		return nil, ErrInvalidSampledByte
+	}
+
+	return &sc, nil
+}
+
+// BuildSingleHeader takes the values from the SpanContext and builds the B3 header
+func BuildSingleHeader(sc model.SpanContext) string {
+	header := []string{}
+	if !sc.TraceID.Empty() && sc.ID > 0 {
+		header = append(header, sc.TraceID.String(), sc.ID.String())
+	}
+
+	if sc.Debug {
+		header = append(header, "d")
+	} else if sc.Sampled != nil {
+		if *sc.Sampled {
+			header = append(header, "1")
+		} else {
+			header = append(header, "0")
+		}
+	}
+
+	if sc.ParentID != nil {
+		header = append(header, sc.ParentID.String())
+	}
+
+	return strings.Join(header, "-")
+}