75 lines
1.9 KiB
Go
75 lines
1.9 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/google/uuid"
|
|
|
|
"git.threesix.ai/jordan/sp4-verify-1770325799/pkg/httpcontext"
|
|
"git.threesix.ai/jordan/sp4-verify-1770325799/pkg/logging"
|
|
)
|
|
|
|
// Trace ID headers in priority order.
|
|
const (
|
|
TraceIDHeader = "X-Trace-ID"
|
|
CloudTraceHeader = "X-Cloud-Trace-Context"
|
|
)
|
|
|
|
// Tracing returns middleware that extracts or generates trace IDs.
|
|
//
|
|
// Checks headers in order:
|
|
// 1. X-Trace-ID - direct trace ID
|
|
// 2. X-Cloud-Trace-Context - GCP format "TRACE_ID/SPAN_ID;o=OPTIONS"
|
|
// 3. Generates a new UUID if none found
|
|
//
|
|
// The trace ID is stored in context via httpcontext.SetTraceID and
|
|
// logging.WithTraceID, and set in the X-Trace-ID response header.
|
|
//
|
|
// Usage:
|
|
//
|
|
// r.Use(middleware.RequestID())
|
|
// r.Use(middleware.Tracing())
|
|
// r.Use(middleware.RequestLogger(logger))
|
|
func Tracing() func(http.Handler) http.Handler {
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
traceID := extractTraceID(r)
|
|
|
|
if traceID == "" {
|
|
traceID = uuid.New().String()
|
|
}
|
|
|
|
// Store in context
|
|
ctx := httpcontext.SetTraceID(r.Context(), traceID)
|
|
ctx = logging.WithTraceID(ctx, traceID)
|
|
|
|
// Set response header
|
|
w.Header().Set(TraceIDHeader, traceID)
|
|
|
|
next.ServeHTTP(w, r.WithContext(ctx))
|
|
})
|
|
}
|
|
}
|
|
|
|
// extractTraceID tries to extract a trace ID from known headers.
|
|
func extractTraceID(r *http.Request) string {
|
|
// X-Trace-ID takes priority
|
|
if traceID := r.Header.Get(TraceIDHeader); traceID != "" {
|
|
return traceID
|
|
}
|
|
|
|
// X-Cloud-Trace-Context format: "TRACE_ID/SPAN_ID;o=OPTIONS"
|
|
if cloudTrace := r.Header.Get(CloudTraceHeader); cloudTrace != "" {
|
|
if idx := strings.IndexByte(cloudTrace, '/'); idx > 0 {
|
|
return cloudTrace[:idx]
|
|
}
|
|
if idx := strings.IndexByte(cloudTrace, ';'); idx > 0 {
|
|
return cloudTrace[:idx]
|
|
}
|
|
return cloudTrace
|
|
}
|
|
|
|
return ""
|
|
}
|