span_v05.go

// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016-present Datadog, Inc.

// Package idx is used to unmarshal v1.0 Trace payloads
package idx

import (
	"errors"
	"fmt"

	"github.com/tinylib/msgp/msgp"
)

// buildStringTable builds a string table from a list of unique v05 strings
// However, unlike the v05 array, the string table expects that the 0 index is always the empty string
// To get around this we swap whatever string is at index 0 with the location of the empty string (if present), or append it to the end
// We then return the string table and the new location for string ref `0` (0 if unchanged)
func buildStringTable(v05Strings []string) (*StringTable, uint32) {
	if len(v05Strings) == 0 {
		return NewStringTable(), 0
	}
	if v05Strings[0] == "" {
		// The empty string is already at index 0, so we can just use the string table as is
		return StringTableFromArray(v05Strings), 0
	}
	emptyStringIndex := uint32(0)
	for i, str := range v05Strings {
		if str == "" {
			emptyStringIndex = uint32(i)
			break
		}
	}
	newZeroRef := emptyStringIndex
	if emptyStringIndex != 0 {
		v05Strings[0], v05Strings[emptyStringIndex] = v05Strings[emptyStringIndex], v05Strings[0]
	} else {
		// The empty string is not present, so we append the 0th string to the end and set the first index to the empty string
		v05Strings = append(v05Strings, v05Strings[0])
		v05Strings[0] = ""
		newZeroRef = uint32(len(v05Strings) - 1)
	}
	return StringTableFromArray(v05Strings), newZeroRef
}

// UnmarshalMsgDictionary decodes a InternalTracerPayload using the specification from the v0.5 endpoint.
// For details, see the documentation for endpoint v0.5 in pkg/trace/api/version.go
func (tp *InternalTracerPayload) UnmarshalMsgDictionary(bts []byte) error {
	var err error
	if _, bts, err = safeReadHeaderBytes(bts, msgp.ReadArrayHeaderBytes); err != nil {
		return err
	}
	// read dictionary
	var sz uint32
	if sz, bts, err = safeReadHeaderBytes(bts, msgp.ReadArrayHeaderBytes); err != nil {
		return err
	}
	dict := make([]string, sz)
	for i := range dict {
		var str string
		str, bts, err = parseStringBytes(bts)
		if err != nil {
			return err
		}
		dict[i] = str
	}
	dictSize := sz
	stringTable, newZeroRef := buildStringTable(dict)
	tp.Strings = stringTable
	// read num chunks
	sz, bts, err = safeReadHeaderBytes(bts, msgp.ReadArrayHeaderBytes)
	if err != nil {
		return err
	}
	if cap(tp.Chunks) >= int(sz) {
		tp.Chunks = tp.Chunks[:sz]
	} else {
		tp.Chunks = make([]*InternalTraceChunk, sz)
	}
	chunkConvertedFields := ChunkConvertedFields{}
	convertedTagSet := false
	for i := range tp.Chunks {
		sz, bts, err = safeReadHeaderBytes(bts, msgp.ReadArrayHeaderBytes)
		if err != nil {
			return err
		}
		if tp.Chunks[i] == nil {
			tp.Chunks[i] = &InternalTraceChunk{Strings: stringTable}
		}
		if cap(tp.Chunks[i].Spans) >= int(sz) {
			tp.Chunks[i].Spans = tp.Chunks[i].Spans[:sz]
		} else {
			tp.Chunks[i].Spans = make([]*InternalSpan, sz)
		}
		convertedFields := NewSpanConvertedFields()
		var rootSampling RootSamplingMergeState
		for j := range tp.Chunks[i].Spans {
			if tp.Chunks[i].Spans[j] == nil {
				tp.Chunks[i].Spans[j] = NewInternalSpan(stringTable, &Span{})
			}
			if bts, err = tp.Chunks[i].Spans[j].UnmarshalMsgDictionaryConverted(bts, convertedFields, dictSize, newZeroRef); err != nil {
				return err
			}
			if !convertedTagSet {
				// _dd.convertedv1 marks that this payload was converted from the v0.5
				// wire format. It is a debugging aid, so we only tag the first span of
				// the payload rather than paying the allocation on every span.
				tp.Chunks[i].Spans[j].SetStringAttribute("_dd.convertedv1", "v05")
				convertedTagSet = true
			}
			rootSampling.ReconcileSamplingPriorityAfterChunkSpan(convertedFields, tp.Chunks[i].Spans[j].ParentID())
		}
		tp.Chunks[i].ApplyPromotedFields(convertedFields, &chunkConvertedFields)
	}
	tp.ApplyPromotedFields(&chunkConvertedFields)
	return nil
}

// spanPropertyCount specifies the number of top-level properties that a span
// has.
const spanPropertyCount = 12

func readV05StringRef(dictSize uint32, newZeroRef uint32, bts []byte) (uint32, []byte, error) {
	var parsedRef uint32
	var err error
	parsedRef, bts, err = msgp.ReadUint32Bytes(bts)
	if err != nil {
		return 0, bts, err
	}
	if parsedRef >= dictSize {
		return 0, bts, fmt.Errorf("dictionary index %d out of range", parsedRef)
	}
	if parsedRef == 0 && newZeroRef != 0 {
		// This string was moved from index 0 to index newZeroRef, so we return the new index
		return newZeroRef, bts, nil
	}
	return parsedRef, bts, nil
}

// UnmarshalMsgDictionaryConverted decodes a v05 span directly into an InternalSpan, for details, see the documentation for endpoint v0.5
// in pkg/trace/api/version.go
// The provided InternalSpan must have a pre-populated Strings field
// newZeroRef is the new location for string ref `0` (0 if unchanged) see buildStringTable for more details
func (s *InternalSpan) UnmarshalMsgDictionaryConverted(bts []byte, convertedFields *SpanConvertedFields, dictSize uint32, newZeroRef uint32) ([]byte, error) {
	var (
		sz  uint32
		err error
	)
	sz, bts, err = safeReadHeaderBytes(bts, msgp.ReadArrayHeaderBytes)
	if err != nil {
		return bts, err
	}
	if sz != spanPropertyCount {
		return bts, errors.New("encoded span needs exactly 12 elements in array")
	}
	// Service (0)
	s.span.ServiceRef, bts, err = readV05StringRef(dictSize, newZeroRef, bts)
	if err != nil {
		return bts, err
	}
	// Name (1)
	s.span.NameRef, bts, err = readV05StringRef(dictSize, newZeroRef, bts)
	if err != nil {
		return bts, err
	}
	// Resource (2)
	s.span.ResourceRef, bts, err = readV05StringRef(dictSize, newZeroRef, bts)
	if err != nil {
		return bts, err
	}
	// TraceID (3)
	convertedFields.TraceIDLower, bts, err = parseUint64Bytes(bts)
	if err != nil {
		return bts, err
	}
	// SpanID (4)
	s.span.SpanID, bts, err = parseUint64Bytes(bts)
	if err != nil {
		return bts, err
	}
	// ParentID (5)
	s.span.ParentID, bts, err = parseUint64Bytes(bts)
	if err != nil {
		return bts, err
	}
	// Start (6)
	var spanStart int64
	spanStart, bts, err = parseInt64Bytes(bts)
	s.span.Start = uint64(spanStart)
	if err != nil {
		return bts, err
	}
	// Duration (7)
	var spanDuration int64
	spanDuration, bts, err = parseInt64Bytes(bts)
	s.span.Duration = uint64(spanDuration)
	if err != nil {
		return bts, err
	}
	// Error (8)
	var spanError int32
	spanError, bts, err = parseInt32Bytes(bts)
	s.span.Error = spanError != 0
	if err != nil {
		return bts, err
	}
	// Meta (9)
	sz, bts, err = safeReadHeaderBytes(bts, msgp.ReadMapHeaderBytes)
	if err != nil {
		return bts, err
	}
	if s.span.Attributes == nil && sz > 0 {
		s.span.Attributes = make(map[uint32]*AnyValue, sz)
	}
	if sz > 0 {
		if err = checkSlabCount(sz, bts); err != nil {
			return bts, err
		}
		// Slab-allocate the AnyValue containers and their string-ref oneof wrappers
		// for every meta entry in two allocations, rather than two per entry. The map
		// holds pointers into these backing arrays, which live as long as the span.
		values := make([]AnyValue, sz)
		refs := make([]AnyValue_StringValueRef, sz)
		for i := uint32(0); i < sz; i++ {
			var key, val uint32
			key, bts, err = readV05StringRef(dictSize, newZeroRef, bts)
			if err != nil {
				return bts, err
			}
			val, bts, err = readV05StringRef(dictSize, newZeroRef, bts)
			if err != nil {
				return bts, err
			}
			s.handlePromotedMetaFields(key, val, convertedFields)
			refs[i].StringValueRef = val
			values[i].Value = &refs[i]
			s.span.Attributes[key] = &values[i]
		}
	}
	// Metrics (10)
	sz, bts, err = safeReadHeaderBytes(bts, msgp.ReadMapHeaderBytes)
	if err != nil {
		return bts, err
	}
	if s.span.Attributes == nil && sz > 0 {
		s.span.Attributes = make(map[uint32]*AnyValue, sz)
	}
	if sz > 0 {
		if err = checkSlabCount(sz, bts); err != nil {
			return bts, err
		}
		// Slab-allocate the AnyValue containers and their double oneof wrappers for
		// every metric in two allocations, rather than two per metric.
		values := make([]AnyValue, sz)
		doubles := make([]AnyValue_DoubleValue, sz)
		for i := uint32(0); i < sz; i++ {
			var (
				key uint32
				val float64
			)
			key, bts, err = readV05StringRef(dictSize, newZeroRef, bts)
			if err != nil {
				return bts, err
			}
			val, bts, err = parseFloat64Bytes(bts)
			if err != nil {
				return bts, err
			}
			s.handlePromotedMetricsFields(key, val, convertedFields)
			doubles[i].DoubleValue = val
			values[i].Value = &doubles[i]
			s.span.Attributes[key] = &values[i]
		}
	}
	// Type (11)
	s.span.TypeRef, bts, err = readV05StringRef(dictSize, newZeroRef, bts)
	if err != nil {
		return bts, err
	}
	return bts, nil
}

// safeReadHeaderBytes wraps msgp header readers (typically ReadArrayHeaderBytes and ReadMapHeaderBytes).
// It enforces the dictionary max size of 25MB and protects the caller from making unbounded allocations through `make(any, sz)`.
func safeReadHeaderBytes(b []byte, read func([]byte) (uint32, []byte, error)) (uint32, []byte, error) {
	sz, bts, err := read(b)
	if err != nil {
		return 0, nil, err
	}
	if sz > 25*1e6 {
		// Dictionary can't be larger than 25 MB
		return 0, nil, errors.New("too long payload")
	}
	return sz, bts, err
}

// minBytesPerSlabEntry is a conservative lower bound on the wire size of one meta/metrics/meta_struct
// entry: every entry reads at least two msgpack values (e.g. two 1-byte nils).
const minBytesPerSlabEntry = 2

// checkSlabCount guards a slab pre-allocation (make([]AnyValue, n) and its oneof-wrapper sibling)
// against a claimed entry count that couldn't possibly be backed by the remaining bytes. Without this,
// a tiny malicious payload could set a map header to millions of entries and force a large allocation
// before decoding ever reaches the missing bytes and fails naturally.
func checkSlabCount(n uint32, remaining []byte) error {
	if uint64(n)*minBytesPerSlabEntry > uint64(len(remaining)) {
		return fmt.Errorf("not enough data for %d entries", n)
	}
	return nil
}