1package mocktracer
2
3import (
4	"fmt"
5	"net/url"
6	"strconv"
7	"strings"
8
9	"github.com/opentracing/opentracing-go"
10)
11
12const mockTextMapIdsPrefix = "mockpfx-ids-"
13const mockTextMapBaggagePrefix = "mockpfx-baggage-"
14
15var emptyContext = MockSpanContext{}
16
17// Injector is responsible for injecting SpanContext instances in a manner suitable
18// for propagation via a format-specific "carrier" object. Typically the
19// injection will take place across an RPC boundary, but message queues and
20// other IPC mechanisms are also reasonable places to use an Injector.
21type Injector interface {
22	// Inject takes `SpanContext` and injects it into `carrier`. The actual type
23	// of `carrier` depends on the `format` passed to `Tracer.Inject()`.
24	//
25	// Implementations may return opentracing.ErrInvalidCarrier or any other
26	// implementation-specific error if injection fails.
27	Inject(ctx MockSpanContext, carrier interface{}) error
28}
29
30// Extractor is responsible for extracting SpanContext instances from a
31// format-specific "carrier" object. Typically the extraction will take place
32// on the server side of an RPC boundary, but message queues and other IPC
33// mechanisms are also reasonable places to use an Extractor.
34type Extractor interface {
35	// Extract decodes a SpanContext instance from the given `carrier`,
36	// or (nil, opentracing.ErrSpanContextNotFound) if no context could
37	// be found in the `carrier`.
38	Extract(carrier interface{}) (MockSpanContext, error)
39}
40
41// TextMapPropagator implements Injector/Extractor for TextMap and HTTPHeaders formats.
42type TextMapPropagator struct {
43	HTTPHeaders bool
44}
45
46// Inject implements the Injector interface
47func (t *TextMapPropagator) Inject(spanContext MockSpanContext, carrier interface{}) error {
48	writer, ok := carrier.(opentracing.TextMapWriter)
49	if !ok {
50		return opentracing.ErrInvalidCarrier
51	}
52	// Ids:
53	writer.Set(mockTextMapIdsPrefix+"traceid", strconv.Itoa(spanContext.TraceID))
54	writer.Set(mockTextMapIdsPrefix+"spanid", strconv.Itoa(spanContext.SpanID))
55	writer.Set(mockTextMapIdsPrefix+"sampled", fmt.Sprint(spanContext.Sampled))
56	// Baggage:
57	for baggageKey, baggageVal := range spanContext.Baggage {
58		safeVal := baggageVal
59		if t.HTTPHeaders {
60			safeVal = url.QueryEscape(baggageVal)
61		}
62		writer.Set(mockTextMapBaggagePrefix+baggageKey, safeVal)
63	}
64	return nil
65}
66
67// Extract implements the Extractor interface
68func (t *TextMapPropagator) Extract(carrier interface{}) (MockSpanContext, error) {
69	reader, ok := carrier.(opentracing.TextMapReader)
70	if !ok {
71		return emptyContext, opentracing.ErrInvalidCarrier
72	}
73	rval := MockSpanContext{0, 0, true, nil}
74	err := reader.ForeachKey(func(key, val string) error {
75		lowerKey := strings.ToLower(key)
76		switch {
77		case lowerKey == mockTextMapIdsPrefix+"traceid":
78			// Ids:
79			i, err := strconv.Atoi(val)
80			if err != nil {
81				return err
82			}
83			rval.TraceID = i
84		case lowerKey == mockTextMapIdsPrefix+"spanid":
85			// Ids:
86			i, err := strconv.Atoi(val)
87			if err != nil {
88				return err
89			}
90			rval.SpanID = i
91		case lowerKey == mockTextMapIdsPrefix+"sampled":
92			b, err := strconv.ParseBool(val)
93			if err != nil {
94				return err
95			}
96			rval.Sampled = b
97		case strings.HasPrefix(lowerKey, mockTextMapBaggagePrefix):
98			// Baggage:
99			if rval.Baggage == nil {
100				rval.Baggage = make(map[string]string)
101			}
102			safeVal := val
103			if t.HTTPHeaders {
104				// unescape errors are ignored, nothing can be done
105				if rawVal, err := url.QueryUnescape(val); err == nil {
106					safeVal = rawVal
107				}
108			}
109			rval.Baggage[lowerKey[len(mockTextMapBaggagePrefix):]] = safeVal
110		}
111		return nil
112	})
113	if rval.TraceID == 0 || rval.SpanID == 0 {
114		return emptyContext, opentracing.ErrSpanContextNotFound
115	}
116	if err != nil {
117		return emptyContext, err
118	}
119	return rval, nil
120}
121