context/tracing/otel.go

package tracing

import (
	"context"
	"strings"
)

// TraceParent holds the decomposed W3C traceparent header. Unlike the X-Ray
// header, traceparent uses a fixed positional format.
type TraceParent struct {
	Version    string
	TraceID    string
	ParentID   string
	TraceFlags string
}

// IsValid reports whether the fields look reasonable. It does not validate
// the cryptographic integrity of the ids, only that they have the right
// shapes.
func (p TraceParent) IsValid() bool {
	return p.Version == "00" &&
		len(p.TraceID) == 32 &&
		len(p.ParentID) == 16 &&
		len(p.TraceFlags) == 2
}

// Sampled reports whether the sampled bit is set in the trace flags.
func (p TraceParent) Sampled() bool {
	if len(p.TraceFlags) < 2 {
		return false
	}
	// The flags are a two-character hex byte. Sampled is bit 0.
	b := fromHex(p.TraceFlags[0])<<4 | fromHex(p.TraceFlags[1])
	return b&0x01 == 0x01
}

// ParseTraceParent parses a W3C traceparent header.
// Example: 00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01.
func ParseTraceParent(raw string) (TraceParent, bool) {
	parts := strings.Split(raw, "-")
	if len(parts) != 4 {
		return TraceParent{}, false
	}
	p := TraceParent{Version: parts[0], TraceID: parts[1], ParentID: parts[2], TraceFlags: parts[3]}
	return p, p.IsValid()
}

type otelKey struct{}

var otelCtxKey otelKey

// WithTraceParent attaches p to ctx.
func WithTraceParent(ctx context.Context, p TraceParent) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, otelCtxKey, p)
}

// TraceParentFromContext returns the parent attached by WithTraceParent.
func TraceParentFromContext(ctx context.Context) (TraceParent, bool) {
	if ctx == nil {
		return TraceParent{}, false
	}
	p, ok := ctx.Value(otelCtxKey).(TraceParent)
	return p, ok
}

// LogFields returns "trace_id", "span_id", and "sampled" entries suitable
// for correlating log lines with traces.
func (p TraceParent) LogFields() []Field {
	if !p.IsValid() {
		return nil
	}
	return []Field{
		{Key: "trace_id", Value: p.TraceID},
		{Key: "span_id", Value: p.ParentID},
		{Key: "sampled", Value: p.Sampled()},
	}
}

// FromXRay performs a best-effort conversion from the AWS X-Ray header form
// into a W3C traceparent. AWS's mapping guidance says to take the 32
// trailing hex chars of the Root field as TraceID.
func FromXRay(x XRayHeader) (TraceParent, bool) {
	root := x.Root
	if i := strings.LastIndex(root, "-"); i >= 0 && i+1 < len(root) {
		root = root[i+1:]
	}
	root = strings.TrimPrefix(root, "1-")
	root = strings.ReplaceAll(root, "-", "")
	if len(root) != 32 || x.Parent == "" {
		return TraceParent{}, false
	}
	flags := "00"
	if x.Sampled == "1" {
		flags = "01"
	}
	p := TraceParent{Version: "00", TraceID: root, ParentID: x.Parent, TraceFlags: flags}
	return p, p.IsValid()
}

func fromHex(c byte) byte {
	switch {
	case c >= '0' && c <= '9':
		return c - '0'
	case c >= 'a' && c <= 'f':
		return c - 'a' + 10
	case c >= 'A' && c <= 'F':
		return c - 'A' + 10
	}
	return 0
}