优化

internal/server/middleware/request.go

@@ -3,15 +3,14 @@ package middleware
 import (
 	"crypto/rand"
 	"encoding/hex"
-	"regexp"
 	"strings"
 
+	"kra/pkg/logging"
+
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 )
 
-var traceParentPattern = regexp.MustCompile(`^00-([0-9a-f]{32})-([0-9a-f]{16})-[0-9a-f]{2}$`)
-
 func randomHex(bytes int) string {
 	value := make([]byte, bytes)
 	_, _ = rand.Read(value)
@@ -21,13 +20,13 @@ func randomHex(bytes int) string {
 func RequestMeta() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		requestID := c.GetHeader("X-Request-Id")
-		if requestID == "" || len(requestID) > 128 || strings.ContainsAny(requestID, "\r\n") {
+		if !saneHeaderID(requestID) {
 			requestID = uuid.NewString()
 		}
 		traceID, parentSpanID := "", ""
-		if match := traceParentPattern.FindStringSubmatch(strings.ToLower(c.GetHeader("traceparent"))); len(match) == 3 {
-			traceID, parentSpanID = match[1], match[2]
-		} else if candidate := c.GetHeader("X-Trace-Id"); len(candidate) <= 128 && !strings.ContainsAny(candidate, "\r\n") {
+		if upstreamTraceID, upstreamSpanID, ok := parseTraceparent(c.GetHeader("traceparent")); ok {
+			traceID, parentSpanID = upstreamTraceID, upstreamSpanID
+		} else if candidate := c.GetHeader("X-Trace-Id"); saneHeaderID(candidate) {
 			traceID = candidate
 		}
 		if traceID == "" {
@@ -36,13 +35,68 @@ func RequestMeta() gin.HandlerFunc {
 		spanID := randomHex(8)
 		c.Header("X-Request-Id", requestID)
 		c.Header("X-Trace-Id", traceID)
-		if len(traceID) == 32 {
+		if validTraceID(traceID) {
 			c.Header("traceparent", "00-"+traceID+"-"+spanID+"-01")
 		}
 		c.Set("request_id", requestID)
 		c.Set("trace_id", traceID)
 		c.Set("span_id", spanID)
 		c.Set("parent_span_id", parentSpanID)
+		c.Request = c.Request.WithContext(logging.WithContextFields(c.Request.Context(), &logging.ContextFields{
+			RequestID: requestID, TraceID: traceID, SpanID: spanID, ParentSpanID: parentSpanID,
+			DeviceID: c.GetHeader("X-Device-Id"), ClientIP: c.ClientIP(),
+			HTTPMethod: c.Request.Method, HTTPPath: c.Request.URL.Path,
+		}))
 		c.Next()
 	}
 }
+
+func parseTraceparent(value string) (traceID, parentSpanID string, ok bool) {
+	parts := strings.Split(value, "-")
+	if len(parts) < 4 {
+		return "", "", false
+	}
+	version, traceID, parentSpanID, flags := parts[0], parts[1], parts[2], parts[3]
+	if !lowerHex(version, 2) || version == "ff" || version == "00" && len(parts) != 4 ||
+		!validTraceID(traceID) || !lowerHex(parentSpanID, 16) || allZero(parentSpanID) || !lowerHex(flags, 2) {
+		return "", "", false
+	}
+	return traceID, parentSpanID, true
+}
+
+func saneHeaderID(value string) bool {
+	if value == "" || len(value) > 64 {
+		return false
+	}
+	for index := 0; index < len(value); index++ {
+		if value[index] <= 0x20 || value[index] > 0x7e {
+			return false
+		}
+	}
+	return true
+}
+
+func validTraceID(value string) bool { return lowerHex(value, 32) && !allZero(value) }
+
+func lowerHex(value string, length int) bool {
+	if len(value) != length {
+		return false
+	}
+	for index := 0; index < len(value); index++ {
+		if current := value[index]; current < '0' || current > '9' {
+			if current < 'a' || current > 'f' {
+				return false
+			}
+		}
+	}
+	return true
+}
+
+func allZero(value string) bool {
+	for index := 0; index < len(value); index++ {
+		if value[index] != '0' {
+			return false
+		}
+	}
+	return true
+}