go-omni/telemetry/headers.go

package telemetry

import (
	"context"
	"errors"
	"net/http"
	"sort"
	"strings"
	"unicode"

	"git.yuntuo.hk/k2safe/go-omni/logx"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"
)

// HeaderPropagationConfig maps a legacy or gateway header pair to the
// standard OpenTelemetry trace context. TraceIDHeader and SpanIDHeader must
// carry a 32- and 16-hexadecimal ID respectively. RequestIDHeader is optional
// and is copied to/from the request correlation context.
//
// Legacy headers do not carry a sampling flag. The adapter treats a valid pair
// as sampled so the trace remains visible to the configured exporter. Services
// that need upstream sampling decisions should use the standard traceparent
// header, which always takes precedence.
type HeaderPropagationConfig struct {
	TraceIDHeader   string `json:"traceIdHeader"`
	SpanIDHeader    string `json:"spanIdHeader"`
	RequestIDHeader string `json:"requestIdHeader,omitempty"`
}

// HeaderPropagator accepts standard W3C trace context and one configured
// legacy header pair. W3C context wins when both formats are present. It also
// propagates the configured request ID without making it part of the span
// parent. Use it only at a boundary that is authorized to supply these IDs.
type HeaderPropagator struct {
	config HeaderPropagationConfig
	trace  propagation.TraceContext
}

// NewHeaderPropagator creates a generic compatibility propagator. Header names
// are canonicalized and compared case-insensitively. It does not trust or
// inspect a Host header; the caller decides which ingress is trusted.
func NewHeaderPropagator(config HeaderPropagationConfig) (*HeaderPropagator, error) {
	traceHeader, err := canonicalPropagationHeader(config.TraceIDHeader)
	if err != nil {
		return nil, errors.Join(errors.New("telemetry: invalid trace ID header"), err)
	}
	spanHeader, err := canonicalPropagationHeader(config.SpanIDHeader)
	if err != nil {
		return nil, errors.Join(errors.New("telemetry: invalid span ID header"), err)
	}
	requestHeader := ""
	if strings.TrimSpace(config.RequestIDHeader) != "" {
		requestHeader, err = canonicalPropagationHeader(config.RequestIDHeader)
		if err != nil {
			return nil, errors.Join(errors.New("telemetry: invalid request ID header"), err)
		}
	}
	if strings.EqualFold(traceHeader, spanHeader) ||
		(requestHeader != "" && (strings.EqualFold(traceHeader, requestHeader) || strings.EqualFold(spanHeader, requestHeader))) {
		return nil, errors.New("telemetry: propagation headers must be distinct")
	}
	for _, name := range []string{traceHeader, spanHeader, requestHeader} {
		if strings.EqualFold(name, "Traceparent") || strings.EqualFold(name, "Tracestate") {
			return nil, errors.New("telemetry: custom headers cannot replace W3C trace context")
		}
	}
	config.TraceIDHeader = traceHeader
	config.SpanIDHeader = spanHeader
	config.RequestIDHeader = requestHeader
	return &HeaderPropagator{config: config}, nil
}

func canonicalPropagationHeader(name string) (string, error) {
	name = strings.TrimSpace(name)
	if name == "" {
		return "", errors.New("header name is required")
	}
	canonical := http.CanonicalHeaderKey(name)
	if canonical == "" || strings.IndexFunc(name, unicode.IsSpace) >= 0 || strings.ContainsAny(name, "\r\n:") {
		return "", errors.New("header name is malformed")
	}
	return canonical, nil
}

// Inject writes W3C trace context first and then the configured compatibility
// headers. Existing caller values are replaced so all propagated IDs describe
// the same current span.
func (p *HeaderPropagator) Inject(ctx context.Context, carrier propagation.TextMapCarrier) {
	if p == nil || carrier == nil {
		return
	}
	p.trace.Inject(ctx, carrier)
	sc := trace.SpanContextFromContext(ctx)
	if sc.IsValid() {
		carrier.Set(p.config.TraceIDHeader, sc.TraceID().String())
		carrier.Set(p.config.SpanIDHeader, sc.SpanID().String())
	}
	if p.config.RequestIDHeader != "" {
		if requestID := validCorrelationID(logx.RequestID(ctx)); requestID != "" {
			carrier.Set(p.config.RequestIDHeader, requestID)
		}
	}
}

// Extract prefers a valid W3C parent. If it is absent or malformed, a valid
// configured legacy pair becomes a remote parent. Invalid legacy IDs are
// ignored, never copied into logs or a new outbound request.
func (p *HeaderPropagator) Extract(ctx context.Context, carrier propagation.TextMapCarrier) context.Context {
	if p == nil || carrier == nil {
		return ctx
	}
	ctx = p.trace.Extract(ctx, carrier)
	if p.config.RequestIDHeader != "" {
		if requestID := validCorrelationID(carrierValue(carrier, p.config.RequestIDHeader)); requestID != "" {
			ctx = logx.WithRequestID(ctx, requestID)
		}
	}
	if trace.SpanContextFromContext(ctx).IsValid() {
		return ctx
	}
	traceID, err := trace.TraceIDFromHex(carrierValue(carrier, p.config.TraceIDHeader))
	if err != nil || !traceID.IsValid() {
		return ctx
	}
	spanID, err := trace.SpanIDFromHex(carrierValue(carrier, p.config.SpanIDHeader))
	if err != nil || !spanID.IsValid() {
		return ctx
	}
	return trace.ContextWithRemoteSpanContext(ctx, trace.NewSpanContext(trace.SpanContextConfig{
		TraceID: traceID, SpanID: spanID, TraceFlags: trace.FlagsSampled, Remote: true,
	}))
}

// Fields reports both W3C and compatibility fields in stable order.
func (p *HeaderPropagator) Fields() []string {
	if p == nil {
		return nil
	}
	fields := []string{"traceparent", "tracestate", p.config.TraceIDHeader, p.config.SpanIDHeader}
	if p.config.RequestIDHeader != "" {
		fields = append(fields, p.config.RequestIDHeader)
	}
	sort.Strings(fields)
	return fields
}

func carrierValue(carrier propagation.TextMapCarrier, name string) string {
	if value := carrier.Get(name); value != "" {
		return value
	}
	for _, key := range carrier.Keys() {
		if strings.EqualFold(key, name) {
			return carrier.Get(key)
		}
	}
	return ""
}

func validCorrelationID(value string) string {
	if value == "" || value != strings.TrimSpace(value) || len(value) > 128 {
		return ""
	}
	for _, r := range value {
		if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '-' || r == '_' || r == '.' || r == ':' {
			continue
		}
		return ""
	}
	return value
}

var _ propagation.TextMapPropagator = (*HeaderPropagator)(nil)