Parse request x-request-id and expose it in contextual logger (#746)

* parse x-request-id and expose in contextual logger

Signed-off-by: Hang Yin <[email protected]>

* make ExtractHeaderValue a general func

---------

Signed-off-by: Hang Yin <[email protected]>

Author: delavet

pkg/bbr/handlers/server.go

@@ -28,6 +28,7 @@ import (
 	"google.golang.org/grpc/status"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+	requtil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/request"
 )
 
 func NewServer(streaming bool) *Server {
@@ -74,6 +75,11 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 				// If streaming and the body is not empty, then headers are handled when processing request body.
 				loggerVerbose.Info("Received headers, passing off header processing until body arrives...")
 			} else {
+				if requestId := requtil.ExtractHeaderValue(v, requtil.RequestIdHeaderKey); len(requestId) > 0 {
+					logger = logger.WithValues(requtil.RequestIdHeaderKey, requestId)
+					loggerVerbose = logger.V(logutil.VERBOSE)
+					ctx = log.IntoContext(ctx, logger)
+				}
 				responses, err = s.HandleRequestHeaders(req.GetRequestHeaders())
 			}
 		case *extProcPb.ProcessingRequest_RequestBody:

pkg/epp/handlers/server.go

@@ -40,6 +40,7 @@ import (
 	schedulingtypes "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 	errutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/error"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+	requtil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/request"
 )
 
 func NewStreamingServer(scheduler Scheduler, destinationEndpointHintMetadataNamespace, destinationEndpointHintKey string, datastore datastore.Datastore) *StreamingServer {
@@ -167,6 +168,11 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 
 		switch v := req.Request.(type) {
 		case *extProcPb.ProcessingRequest_RequestHeaders:
+			if requestId := requtil.ExtractHeaderValue(v, requtil.RequestIdHeaderKey); len(requestId) > 0 {
+				logger = logger.WithValues(requtil.RequestIdHeaderKey, requestId)
+				loggerTrace = logger.V(logutil.TRACE)
+				ctx = log.IntoContext(ctx, logger)
+			}
 			err = s.HandleRequestHeaders(ctx, reqCtx, v)
 		case *extProcPb.ProcessingRequest_RequestBody:
 			loggerTrace.Info("Incoming body chunk", "EoS", v.RequestBody.EndOfStream)

pkg/epp/util/request/headers.go

@@ -0,0 +1,24 @@
+package request
+
+import (
+	"strings"
+
+	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+)
+
+const (
+	RequestIdHeaderKey = "x-request-id"
+)
+
+func ExtractHeaderValue(req *extProcPb.ProcessingRequest_RequestHeaders, headerKey string) string {
+	// header key should be case insensitive
+	headerKeyInLower := strings.ToLower(headerKey)
+	if req != nil && req.RequestHeaders != nil && req.RequestHeaders.Headers != nil {
+		for _, headerKv := range req.RequestHeaders.Headers.Headers {
+			if strings.ToLower(headerKv.Key) == headerKeyInLower {
+				return string(headerKv.RawValue)
+			}
+		}
+	}
+	return ""
+}

pkg/epp/util/request/headers_test.go

@@ -0,0 +1,59 @@
+package request
+
+import (
+	"testing"
+
+	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
+	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+)
+
+func TestExtractHeaderValue(t *testing.T) {
+	tests := []struct {
+		name     string
+		headers  []*corev3.HeaderValue
+		key      string
+		expected string
+	}{
+		{
+			name: "Exact match",
+			headers: []*corev3.HeaderValue{
+				{Key: "x-request-id", RawValue: []byte("123")},
+			},
+			key:      "x-request-id",
+			expected: "123",
+		},
+		{
+			name: "Case-insensitive match",
+			headers: []*corev3.HeaderValue{
+				{Key: "X-Request-ID", RawValue: []byte("456")},
+			},
+			key:      "x-request-id",
+			expected: "456",
+		},
+		{
+			name: "Non-existent key",
+			headers: []*corev3.HeaderValue{
+				{Key: "other-header", RawValue: []byte("abc")},
+			},
+			key:      "x-request-id",
+			expected: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := &extProcPb.ProcessingRequest_RequestHeaders{
+				RequestHeaders: &extProcPb.HttpHeaders{
+					Headers: &corev3.HeaderMap{
+						Headers: tt.headers,
+					},
+				},
+			}
+
+			result := ExtractHeaderValue(req, tt.key)
+			if result != tt.expected {
+				t.Errorf("ExtractHeaderValue() = %v, want %v", result, tt.expected)
+			}
+		})
+	}
+}