# Copyright (c) 2016 Uber Technologies, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import struct from opentracing import ( InvalidCarrierException, SpanContextCorruptedException, ) from .constants import ( BAGGAGE_HEADER_KEY, BAGGAGE_HEADER_PREFIX, DEBUG_ID_HEADER_KEY, TRACE_ID_HEADER, ) from .span_context import SpanContext from .constants import SAMPLED_FLAG, DEBUG_FLAG import urllib.parse class Codec(object): def inject(self, span_context, carrier): raise NotImplementedError() def extract(self, carrier): raise NotImplementedError() class TextCodec(Codec): def __init__(self, url_encoding=False, trace_id_header=TRACE_ID_HEADER, baggage_header_prefix=BAGGAGE_HEADER_PREFIX, debug_id_header=DEBUG_ID_HEADER_KEY, baggage_header=BAGGAGE_HEADER_KEY): self.url_encoding = url_encoding self.trace_id_header = trace_id_header.lower().replace('_', '-') self.baggage_prefix = baggage_header_prefix.lower().replace('_', '-') self.debug_id_header = debug_id_header.lower().replace('_', '-') self.baggage_header = baggage_header self.prefix_length = len(baggage_header_prefix) def inject(self, span_context, carrier): if not isinstance(carrier, dict): raise InvalidCarrierException('carrier not a collection') # Note: we do not url-encode the trace ID because the ':' separator # is not a problem for HTTP header values carrier[self.trace_id_header] = span_context_to_string( trace_id=span_context.trace_id, span_id=span_context.span_id, parent_id=span_context.parent_id, flags=span_context.flags) baggage = span_context.baggage if baggage: for key, value in baggage.items(): encoded_key = key if isinstance(key, (bytes,)): encoded_key = str(key, 'utf-8') if self.url_encoding: encoded_value = urllib.parse.quote(value) else: if isinstance(key, (bytes,)): encoded_value = str(value, 'utf-8') else: encoded_value = value # Leave the below print(), you will thank me next time you debug unicode strings # print('adding baggage', key, '=>', value, 'as', encoded_key, '=>', encoded_value) header_key = '%s%s' % (self.baggage_prefix, encoded_key) carrier[header_key] = encoded_value def extract(self, carrier): if not hasattr(carrier, 'items'): raise InvalidCarrierException('carrier not a collection') trace_id, span_id, parent_id, flags = None, None, None, None baggage = None debug_id = None for key, value in carrier.items(): uc_key = key.lower() if uc_key == self.trace_id_header: if self.url_encoding: value = urllib.parse.unquote(value) trace_id, span_id, parent_id, flags = \ span_context_from_string(value) elif uc_key.startswith(self.baggage_prefix): if self.url_encoding: value = urllib.parse.unquote(value) attr_key = key[self.prefix_length:] if baggage is None: baggage = {attr_key.lower(): value} else: baggage[attr_key.lower()] = value elif uc_key == self.debug_id_header: if self.url_encoding: value = urllib.parse.unquote(value) debug_id = value elif uc_key == self.baggage_header: if self.url_encoding: value = urllib.parse.unquote(value) baggage = self._parse_baggage_header(value, baggage) if not trace_id or not span_id: # reset all IDs trace_id, span_id, parent_id, flags = None, None, None, None if not trace_id and not debug_id and not baggage: return None return SpanContext(trace_id=trace_id, span_id=span_id, parent_id=parent_id, flags=flags, baggage=baggage, debug_id=debug_id) def _parse_baggage_header(self, header, baggage): for part in header.split(','): kv = part.strip().split('=') if len(kv) == 2: if not baggage: baggage = {} baggage[kv[0]] = kv[1] return baggage class BinaryCodec(Codec): """ Implements inject/extract of SpanContext to/from binary that compatible with golang implementation https://github.com/jaegertracing/jaeger-client-go/blob/master/propagation.go#L177-L290 Supports propagation of trace_id, span_id, flags and baggage """ def inject(self, span_context, carrier): if not isinstance(carrier, bytearray): raise InvalidCarrierException('carrier not a bytearray') # check if we have 128 bit trace_id, break it into two 64 units max_int64 = 0xFFFFFFFFFFFFFFFF if span_context.trace_id > max_int64: high = (span_context.trace_id >> 64) & max_int64 low = span_context.trace_id & max_int64 else: high = 0 low = span_context.trace_id carrier += struct.pack('>QQQQBI', high, low, span_context.span_id or 0, span_context.parent_id or 0, span_context.flags, len(span_context.baggage)) for k, v in span_context.baggage.items(): carrier += self._pack_baggage_item(k, v) def extract(self, carrier): if not isinstance(carrier, bytearray): raise InvalidCarrierException('carrier not a bytearray') baggage = {} high_trace_id, low_trace_id, span_id, parent_id, flags, baggage_count = \ struct.unpack('>QQQQBI', carrier[:37]) # if high_trace_id isn't 0, then we are dealing with 128bit trace id integer, # therefore unpack into 1 number if high_trace_id: trace_id = (high_trace_id << 64) | low_trace_id else: trace_id = low_trace_id if baggage_count != 0: baggage_data = carrier[37:] for _ in range(baggage_count): key, value, bytes_read = self._unpack_baggage_item(baggage_data) baggage[key] = value baggage_data = baggage_data[bytes_read:] return SpanContext(trace_id=trace_id, span_id=span_id, parent_id=parent_id, flags=flags, baggage=baggage) def _pack_baggage_item(self, key, value): baggage = bytearray() if not isinstance(key, bytes): key = key.encode('utf-8') baggage += struct.pack('>I', len(key)) baggage += key if not isinstance(value, bytes): value = value.encode('utf-8') baggage += struct.pack('>I', len(value)) baggage += value return baggage def _unpack_baggage_item(self, baggage): bytes_read = 0 key, b_read = self._read_kv(baggage) bytes_read += b_read value, b_read = self._read_kv(baggage[bytes_read:]) bytes_read += b_read return key, value, bytes_read def _read_kv(self, data): data_len = struct.unpack('>i', data[:4])[0] data_value = struct.unpack('>' + 'c' * data_len, data[4:4 + data_len]) bytes_read = 4 + data_len return b''.join(data_value).decode('utf-8'), bytes_read def span_context_to_string(trace_id, span_id, parent_id, flags): """ Serialize span ID to a string {trace_id}:{span_id}:{parent_id}:{flags} Numbers are encoded as variable-length lower-case hex strings. If parent_id is None, it is written as 0. :param trace_id: :param span_id: :param parent_id: :param flags: """ parent_id = parent_id or 0 return '{:x}:{:x}:{:x}:{:x}'.format(trace_id, span_id, parent_id, flags) def span_context_from_string(value): """ Decode span ID from a string into a TraceContext. Returns None if the string value is malformed. :param value: formatted {trace_id}:{span_id}:{parent_id}:{flags} """ if type(value) is list and len(value) > 0: # sometimes headers are presented as arrays of values if len(value) > 1: raise SpanContextCorruptedException( 'trace context must be a string or array of 1: "%s"' % value) value = value[0] if not isinstance(value, (str,)): raise SpanContextCorruptedException( 'trace context not a string "%s"' % value) parts = value.split(':') if len(parts) != 4: raise SpanContextCorruptedException( 'malformed trace context "%s"' % value) try: trace_id = int(parts[0], 16) span_id = int(parts[1], 16) parent_id = int(parts[2], 16) flags = int(parts[3], 16) if trace_id < 1 or span_id < 1 or parent_id < 0 or flags < 0: raise SpanContextCorruptedException( 'malformed trace context "%s"' % value) if parent_id == 0: parent_id = None return trace_id, span_id, parent_id, flags except ValueError as e: raise SpanContextCorruptedException( 'malformed trace context "%s": %s' % (value, e)) # String constants identifying the interop format. ZipkinSpanFormat = 'zipkin-span-format' class ZipkinCodec(Codec): """ ZipkinCodec handles ZipkinSpanFormat, which is an interop format used by TChannel. """ def inject(self, span_context, carrier): if not isinstance(carrier, dict): raise InvalidCarrierException('carrier not a dictionary') carrier['trace_id'] = span_context.trace_id carrier['span_id'] = span_context.span_id carrier['parent_id'] = span_context.parent_id carrier['traceflags'] = span_context.flags def extract(self, carrier): if isinstance(carrier, dict): trace_id = carrier.get('trace_id') span_id = carrier.get('span_id') parent_id = carrier.get('parent_id') flags = carrier.get('traceflags') else: if hasattr(carrier, 'trace_id'): trace_id = getattr(carrier, 'trace_id') else: raise InvalidCarrierException('carrier has no trace_id') if hasattr(carrier, 'span_id'): span_id = getattr(carrier, 'span_id') else: raise InvalidCarrierException('carrier has no span_id') if hasattr(carrier, 'parent_id'): parent_id = getattr(carrier, 'parent_id') else: raise InvalidCarrierException('carrier has no parent_id') if hasattr(carrier, 'traceflags'): flags = getattr(carrier, 'traceflags') else: raise InvalidCarrierException('carrier has no traceflags') if not trace_id: return None return SpanContext(trace_id=trace_id, span_id=span_id, parent_id=parent_id, flags=flags, baggage=None) def header_to_hex(header): if not isinstance(header, (str,)): raise SpanContextCorruptedException( 'malformed trace context "%s", expected hex string' % header) try: return int(header, 16) except ValueError: raise SpanContextCorruptedException( 'malformed trace context "%s", expected hex string' % header) class B3Codec(Codec): """ Support B3 header properties https://github.com/openzipkin/b3-propagation """ trace_header = 'X-B3-TraceId' _trace_header_lc = trace_header.lower() span_header = 'X-B3-SpanId' _span_header_lc = span_header.lower() parent_span_header = 'X-B3-ParentSpanId' _parent_span_header_lc = parent_span_header.lower() sampled_header = 'X-B3-Sampled' _sampled_header_lc = sampled_header.lower() flags_header = 'X-B3-Flags' _flags_header_lc = flags_header.lower() def __init__(self, generate_128bit_trace_id=False): self.generate_128bit_trace_id = generate_128bit_trace_id def inject(self, span_context, carrier): if not isinstance(carrier, dict): raise InvalidCarrierException('carrier not a dictionary') if self.generate_128bit_trace_id: carrier[self.trace_header] = format(span_context.trace_id, 'x').zfill(32) else: carrier[self.trace_header] = format(span_context.trace_id, 'x').zfill(16) carrier[self.span_header] = format(span_context.span_id, 'x').zfill(16) if span_context.parent_id is not None: carrier[self.parent_span_header] = format(span_context.parent_id, 'x').zfill(16) if span_context.flags & DEBUG_FLAG == DEBUG_FLAG: carrier[self.flags_header] = '1' elif span_context.flags & SAMPLED_FLAG == SAMPLED_FLAG: carrier[self.sampled_header] = '1' def extract(self, carrier): if not isinstance(carrier, dict): raise InvalidCarrierException('carrier not a dictionary') trace_id = span_id = parent_id = None flags = 0x00 for header_key, header_value in carrier.items(): if header_value is None: continue lower_key = header_key.lower() if lower_key == self._trace_header_lc: trace_id = header_to_hex(header_value) elif lower_key == self._span_header_lc: span_id = header_to_hex(header_value) elif lower_key == self._parent_span_header_lc: parent_id = header_to_hex(header_value) elif lower_key == self._sampled_header_lc and header_value == '1': flags |= SAMPLED_FLAG elif lower_key == self._flags_header_lc and header_value == '1': flags |= DEBUG_FLAG if not trace_id or not span_id: return None return SpanContext(trace_id=trace_id, span_id=span_id, parent_id=parent_id, flags=flags, baggage=None)