1// Copyright (c) 2017 Uber Technologies, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package zipkin
16
17import (
18	"strconv"
19	"strings"
20
21	opentracing "github.com/opentracing/opentracing-go"
22
23	"github.com/uber/jaeger-client-go"
24)
25
26// Option is a function that sets an option on Propagator
27type Option func(propagator *Propagator)
28
29// BaggagePrefix is a function that sets baggage prefix on Propagator
30func BaggagePrefix(prefix string) Option {
31	return func(propagator *Propagator) {
32		propagator.baggagePrefix = prefix
33	}
34}
35
36// Propagator is an Injector and Extractor
37type Propagator struct {
38	baggagePrefix string
39}
40
41// NewZipkinB3HTTPHeaderPropagator creates a Propagator for extracting and injecting
42// Zipkin HTTP B3 headers into SpanContexts. Baggage is by default enabled and uses prefix
43// 'baggage-'.
44func NewZipkinB3HTTPHeaderPropagator(opts ...Option) Propagator {
45	p := Propagator{baggagePrefix: "baggage-"}
46	for _, opt := range opts {
47		opt(&p)
48	}
49	return p
50}
51
52// Inject conforms to the Injector interface for decoding Zipkin HTTP B3 headers
53func (p Propagator) Inject(
54	sc jaeger.SpanContext,
55	abstractCarrier interface{},
56) error {
57	textMapWriter, ok := abstractCarrier.(opentracing.TextMapWriter)
58	if !ok {
59		return opentracing.ErrInvalidCarrier
60	}
61
62	textMapWriter.Set("x-b3-traceid", sc.TraceID().String())
63	if sc.ParentID() != 0 {
64		textMapWriter.Set("x-b3-parentspanid", strconv.FormatUint(uint64(sc.ParentID()), 16))
65	}
66	textMapWriter.Set("x-b3-spanid", strconv.FormatUint(uint64(sc.SpanID()), 16))
67	if sc.IsSampled() {
68		textMapWriter.Set("x-b3-sampled", "1")
69	} else {
70		textMapWriter.Set("x-b3-sampled", "0")
71	}
72	sc.ForeachBaggageItem(func(k, v string) bool {
73		textMapWriter.Set(p.baggagePrefix+k, v)
74		return true
75	})
76	return nil
77}
78
79// Extract conforms to the Extractor interface for encoding Zipkin HTTP B3 headers
80func (p Propagator) Extract(abstractCarrier interface{}) (jaeger.SpanContext, error) {
81	textMapReader, ok := abstractCarrier.(opentracing.TextMapReader)
82	if !ok {
83		return jaeger.SpanContext{}, opentracing.ErrInvalidCarrier
84	}
85	var traceID jaeger.TraceID
86	var spanID uint64
87	var parentID uint64
88	sampled := false
89	var baggage map[string]string
90	err := textMapReader.ForeachKey(func(rawKey, value string) error {
91		key := strings.ToLower(rawKey) // TODO not necessary for plain TextMap
92		var err error
93		if key == "x-b3-traceid" {
94			traceID, err = jaeger.TraceIDFromString(value)
95		} else if key == "x-b3-parentspanid" {
96			parentID, err = strconv.ParseUint(value, 16, 64)
97		} else if key == "x-b3-spanid" {
98			spanID, err = strconv.ParseUint(value, 16, 64)
99		} else if key == "x-b3-sampled" && (value == "1" || value == "true") {
100			sampled = true
101		} else if strings.HasPrefix(key, p.baggagePrefix) {
102			if baggage == nil {
103				baggage = make(map[string]string)
104			}
105			baggage[key[len(p.baggagePrefix):]] = value
106		}
107		return err
108	})
109
110	if err != nil {
111		return jaeger.SpanContext{}, err
112	}
113	if !traceID.IsValid() {
114		return jaeger.SpanContext{}, opentracing.ErrSpanContextNotFound
115	}
116	return jaeger.NewSpanContext(
117		traceID,
118		jaeger.SpanID(spanID),
119		jaeger.SpanID(parentID),
120		sampled, baggage), nil
121}
122