Support full duplex streaming in body-based routing extension

cmd/body-based-routing/main.go

@@ -44,14 +44,16 @@ import (
 var (
 	grpcPort = flag.Int(
 		"grpcPort",
-		runserver.DefaultGrpcPort,
+		9004,
 		"The gRPC port used for communicating with Envoy proxy")
 	grpcHealthPort = flag.Int(
 		"grpcHealthPort",
 		9005,
 		"The port used for gRPC liveness and readiness probes")
 	metricsPort = flag.Int(
 		"metricsPort", 9090, "The metrics port")
+	streaming = flag.Bool(
+		"streaming", false, "Enables streaming support for Envoy full-duplex streaming mode")
 	logVerbosity = flag.Int("v", logging.DEFAULT, "number for the log level verbosity")
 
 	setupLog = ctrl.Log.WithName("setup")
@@ -92,7 +94,7 @@ func run() error {
 	ctx := ctrl.SetupSignalHandler()
 
 	// Setup runner.
-	serverRunner := &runserver.ExtProcServerRunner{GrpcPort: *grpcPort}
+	serverRunner := runserver.NewDefaultExtProcServerRunner(*grpcPort, *streaming)
 
 	// Register health server.
 	if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), *grpcHealthPort); err != nil {

pkg/body-based-routing/handlers/request.go

@@ -23,55 +23,93 @@ import (
 
 	basepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	eppb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/body-based-routing/metrics"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
+const modelHeader = "X-Gateway-Model-Name"
+
 // HandleRequestBody handles request bodies.
-func (s *Server) HandleRequestBody(ctx context.Context, body *eppb.HttpBody) (*eppb.ProcessingResponse, error) {
+func (s *Server) HandleRequestBody(ctx context.Context, data map[string]any) ([]*eppb.ProcessingResponse, error) {
 	logger := log.FromContext(ctx)
+	var ret []*eppb.ProcessingResponse
 
-	var data map[string]any
-	if err := json.Unmarshal(body.GetBody(), &data); err != nil {
+	requestBodyBytes, err := json.Marshal(data)
+	if err != nil {
 		return nil, err
 	}
 
 	modelVal, ok := data["model"]
 	if !ok {
 		metrics.RecordModelNotInBodyCounter()
 		logger.V(logutil.DEFAULT).Info("Request body does not contain model parameter")
-		return &eppb.ProcessingResponse{
-			Response: &eppb.ProcessingResponse_RequestBody{
-				RequestBody: &eppb.BodyResponse{},
-			},
-		}, nil
+		if s.streaming {
+			ret = append(ret, &eppb.ProcessingResponse{
+				Response: &eppb.ProcessingResponse_RequestHeaders{
+					RequestHeaders: &eppb.HeadersResponse{},
+				},
+			})
+			ret = addStreamedBodyResponse(ret, requestBodyBytes)
+			return ret, nil
+		} else {
+			ret = append(ret, &eppb.ProcessingResponse{
+				Response: &eppb.ProcessingResponse_RequestBody{
+					RequestBody: &eppb.BodyResponse{},
+				},
+			})
+		}
+		return ret, nil
 	}
 
 	modelStr, ok := modelVal.(string)
 	if !ok {
 		metrics.RecordModelNotParsedCounter()
 		logger.V(logutil.DEFAULT).Info("Model parameter value is not a string")
-		return &eppb.ProcessingResponse{
-			Response: &eppb.ProcessingResponse_RequestBody{
-				RequestBody: &eppb.BodyResponse{},
-			},
-		}, fmt.Errorf("the model parameter value %v is not a string", modelVal)
+		return nil, fmt.Errorf("the model parameter value %v is not a string", modelVal)
 	}
 
 	metrics.RecordSuccessCounter()
-	return &eppb.ProcessingResponse{
-		Response: &eppb.ProcessingResponse_RequestBody{
-			RequestBody: &eppb.BodyResponse{
-				Response: &eppb.CommonResponse{
-					// Necessary so that the new headers are used in the routing decision.
-					ClearRouteCache: true,
-					HeaderMutation: &eppb.HeaderMutation{
-						SetHeaders: []*basepb.HeaderValueOption{
-							{
-								Header: &basepb.HeaderValue{
-									Key:      "X-Gateway-Model-Name",
-									RawValue: []byte(modelStr),
+
+	if s.streaming {
+		ret = append(ret, &eppb.ProcessingResponse{
+			Response: &eppb.ProcessingResponse_RequestHeaders{
+				RequestHeaders: &eppb.HeadersResponse{
+					Response: &eppb.CommonResponse{
+						ClearRouteCache: true,
+						HeaderMutation: &eppb.HeaderMutation{
+							SetHeaders: []*basepb.HeaderValueOption{
+								{
+									Header: &basepb.HeaderValue{
+										Key:      modelHeader,
+										RawValue: []byte(modelStr),
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		})
+		ret = addStreamedBodyResponse(ret, requestBodyBytes)
+		return ret, nil
+	}
+
+	return []*eppb.ProcessingResponse{
+		{
+			Response: &eppb.ProcessingResponse_RequestBody{
+				RequestBody: &eppb.BodyResponse{
+					Response: &eppb.CommonResponse{
+						// Necessary so that the new headers are used in the routing decision.
+						ClearRouteCache: true,
+						HeaderMutation: &eppb.HeaderMutation{
+							SetHeaders: []*basepb.HeaderValueOption{
+								{
+									Header: &basepb.HeaderValue{
+										Key:      modelHeader,
+										RawValue: []byte(modelStr),
+									},
 								},
 							},
 						},
@@ -82,20 +120,43 @@ func (s *Server) HandleRequestBody(ctx context.Context, body *eppb.HttpBody) (*e
 	}, nil
 }
 
+func addStreamedBodyResponse(responses []*eppb.ProcessingResponse, requestBodyBytes []byte) []*eppb.ProcessingResponse {
+	return append(responses, &extProcPb.ProcessingResponse{
+		Response: &extProcPb.ProcessingResponse_RequestBody{
+			RequestBody: &extProcPb.BodyResponse{
+				Response: &extProcPb.CommonResponse{
+					BodyMutation: &extProcPb.BodyMutation{
+						Mutation: &extProcPb.BodyMutation_StreamedResponse{
+							StreamedResponse: &extProcPb.StreamedBodyResponse{
+								Body:        requestBodyBytes,
+								EndOfStream: true,
+							},
+						},
+					},
+				},
+			},
+		},
+	})
+}
+
 // HandleRequestHeaders handles request headers.
-func (s *Server) HandleRequestHeaders(headers *eppb.HttpHeaders) (*eppb.ProcessingResponse, error) {
-	return &eppb.ProcessingResponse{
-		Response: &eppb.ProcessingResponse_RequestHeaders{
-			RequestHeaders: &eppb.HeadersResponse{},
+func (s *Server) HandleRequestHeaders(headers *eppb.HttpHeaders) ([]*eppb.ProcessingResponse, error) {
+	return []*eppb.ProcessingResponse{
+		{
+			Response: &eppb.ProcessingResponse_RequestHeaders{
+				RequestHeaders: &eppb.HeadersResponse{},
+			},
 		},
 	}, nil
 }
 
 // HandleRequestTrailers handles request trailers.
-func (s *Server) HandleRequestTrailers(trailers *eppb.HttpTrailers) (*eppb.ProcessingResponse, error) {
-	return &eppb.ProcessingResponse{
-		Response: &eppb.ProcessingResponse_RequestTrailers{
-			RequestTrailers: &eppb.TrailersResponse{},
+func (s *Server) HandleRequestTrailers(trailers *eppb.HttpTrailers) ([]*eppb.ProcessingResponse, error) {
+	return []*eppb.ProcessingResponse{
+		{
+			Response: &eppb.ProcessingResponse_RequestTrailers{
+				RequestTrailers: &eppb.TrailersResponse{},
+			},
 		},
 	}, nil
 }