veza/veza-backend-api/internal/tracing/trace_context.go

package tracing

import (
	"context"
	"fmt"
	"net/http"
	"strings"

	"github.com/google/uuid"
)

const (
	// W3C Trace Context headers
	TraceParentHeader = "traceparent"
	TraceStateHeader  = "tracestate"

	// Legacy headers (for backward compatibility)
	LegacyTraceIDHeader = "X-Trace-ID"
	LegacySpanIDHeader  = "X-Span-ID"
)

// TraceContext représente un contexte de tracing distribué (BE-SVC-018)
type TraceContext struct {
	TraceID    string
	SpanID     string
	ParentSpan string // ID du span parent (pour hiérarchie)
	Sampled    bool   // Indique si le trace doit être échantillonné
}

// NewTraceContext crée un nouveau contexte de tracing
func NewTraceContext() *TraceContext {
	return &TraceContext{
		TraceID: uuid.New().String(),
		SpanID:  uuid.New().String(),
		Sampled: true, // Par défaut, on échantillonne tous les traces
	}
}

// NewChildSpan crée un nouveau span enfant à partir d'un contexte parent
func (tc *TraceContext) NewChildSpan() *TraceContext {
	return &TraceContext{
		TraceID:    tc.TraceID, // Même trace ID
		SpanID:     uuid.New().String(),
		ParentSpan: tc.SpanID, // Le span actuel devient le parent
		Sampled:    tc.Sampled,
	}
}

// ToW3CTraceParent convertit le contexte en format W3C Trace Context (traceparent)
// Format: version-trace_id-parent_id-trace_flags
// Exemple: 00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01
func (tc *TraceContext) ToW3CTraceParent() string {
	version := "00" // Version 00 (actuelle)
	traceID := formatTraceID(tc.TraceID)
	parentID := formatSpanID(tc.ParentSpan)
	if parentID == "" {
		parentID = "0000000000000000" // Pas de parent
	}
	flags := "01" // Sampled flag
	if !tc.Sampled {
		flags = "00"
	}
	return fmt.Sprintf("%s-%s-%s-%s", version, traceID, parentID, flags)
}

// FromW3CTraceParent parse un header traceparent W3C
func FromW3CTraceParent(traceParent string) (*TraceContext, error) {
	parts := strings.Split(traceParent, "-")
	if len(parts) != 4 {
		return nil, fmt.Errorf("invalid traceparent format")
	}

	version := parts[0]
	if version != "00" {
		return nil, fmt.Errorf("unsupported traceparent version: %s", version)
	}

	traceID := parts[1]
	parentID := parts[2]
	flags := parts[3]

	// Parser les flags (dernier caractère = sampled flag)
	sampled := flags[len(flags)-1] == '1'

	tc := &TraceContext{
		TraceID:    traceID,
		SpanID:     uuid.New().String(), // Nouveau span ID pour ce service
		ParentSpan: parentID,
		Sampled:    sampled,
	}

	return tc, nil
}

// formatTraceID formate un trace ID pour W3C (32 caractères hex)
func formatTraceID(id string) string {
	// Supprimer les tirets si c'est un UUID
	id = strings.ReplaceAll(id, "-", "")
	// S'assurer que c'est exactement 32 caractères hex
	if len(id) > 32 {
		id = id[:32]
	} else if len(id) < 32 {
		// Padding avec des zéros
		id = id + strings.Repeat("0", 32-len(id))
	}
	return id
}

// formatSpanID formate un span ID pour W3C (16 caractères hex)
func formatSpanID(id string) string {
	if id == "" {
		return ""
	}
	// Supprimer les tirets si c'est un UUID
	id = strings.ReplaceAll(id, "-", "")
	// S'assurer que c'est exactement 16 caractères hex
	if len(id) > 16 {
		id = id[:16]
	} else if len(id) < 16 {
		// Padding avec des zéros
		id = id + strings.Repeat("0", 16-len(id))
	}
	return id
}

// InjectTraceContext injecte le contexte de tracing dans les headers HTTP
func InjectTraceContext(req *http.Request, tc *TraceContext) {
	if tc == nil {
		return
	}

	// Injecter le format W3C Trace Context (standard)
	req.Header.Set(TraceParentHeader, tc.ToW3CTraceParent())

	// Injecter aussi les headers legacy pour compatibilité
	req.Header.Set(LegacyTraceIDHeader, tc.TraceID)
	req.Header.Set(LegacySpanIDHeader, tc.SpanID)
}

// ExtractTraceContext extrait le contexte de tracing depuis les headers HTTP
func ExtractTraceContext(req *http.Request) *TraceContext {
	// Essayer d'abord le format W3C
	if traceParent := req.Header.Get(TraceParentHeader); traceParent != "" {
		if tc, err := FromW3CTraceParent(traceParent); err == nil {
			return tc
		}
	}

	// Fallback sur les headers legacy
	traceID := req.Header.Get(LegacyTraceIDHeader)
	spanID := req.Header.Get(LegacySpanIDHeader)

	if traceID == "" && spanID == "" {
		// Aucun contexte, créer un nouveau
		return NewTraceContext()
	}

	// Créer un contexte à partir des headers legacy
	if traceID == "" {
		traceID = uuid.New().String()
	}
	if spanID == "" {
		spanID = uuid.New().String()
	}

	return &TraceContext{
		TraceID: traceID,
		SpanID:  spanID,
		Sampled: true,
	}
}

// ContextKey est le type pour la clé de contexte
type contextKey string

const traceContextKey contextKey = "trace_context"

// WithTraceContext ajoute un contexte de tracing au contexte Go
func WithTraceContext(ctx context.Context, tc *TraceContext) context.Context {
	return context.WithValue(ctx, traceContextKey, tc)
}

// FromContext extrait le contexte de tracing depuis le contexte Go
func FromContext(ctx context.Context) *TraceContext {
	if tc, ok := ctx.Value(traceContextKey).(*TraceContext); ok {
		return tc
	}
	return nil
}

// GetOrCreateTraceContext obtient le contexte de tracing depuis le contexte Go, ou en crée un nouveau
func GetOrCreateTraceContext(ctx context.Context) *TraceContext {
	if tc := FromContext(ctx); tc != nil {
		return tc
	}
	return NewTraceContext()
}