Small refactor to capture request data for route. (#765)

Author: kfswain

pkg/epp/handlers/request.go

@@ -35,11 +35,10 @@ import (
 func (s *StreamingServer) HandleRequestBody(
 	ctx context.Context,
 	reqCtx *RequestContext,
-	req *extProcPb.ProcessingRequest,
-	requestBodyMap map[string]interface{},
 ) (*RequestContext, error) {
 	var requestBodyBytes []byte
 	logger := log.FromContext(ctx)
+	requestBodyMap := reqCtx.Request.Body
 
 	// Resolve target models.
 	model, ok := requestBodyMap["model"].(string)
@@ -152,6 +151,15 @@ func (s *StreamingServer) HandleRequestHeaders(ctx context.Context, reqCtx *Requ
 		}
 		endpoint := pod.Address + ":" + strconv.Itoa(int(pool.Spec.TargetPortNumber))
 		s.populateRequestHeaderResponse(reqCtx, endpoint, 0)
+		return nil
+	}
+
+	for _, header := range req.RequestHeaders.Headers.Headers {
+		if header.RawValue != nil {
+			reqCtx.Request.Headers[header.Key] = string(header.RawValue)
+		} else {
+			reqCtx.Request.Headers[header.Key] = header.Value
+		}
 	}
 	return nil
 }

pkg/epp/handlers/server.go

@@ -82,6 +82,7 @@ type RequestContext struct {
 	ResponseComplete          bool
 	ResponseStatusCode        string
 	RequestRunning            bool
+	Request                   *Request
 
 	RequestState         StreamRequestState
 	modelServerStreaming bool
@@ -95,6 +96,10 @@ type RequestContext struct {
 	respTrailerResp *extProcPb.ProcessingResponse
 }
 
+type Request struct {
+	Headers map[string]string
+	Body    map[string]interface{}
+}
 type StreamRequestState int
 
 const (
@@ -118,10 +123,14 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 	// See https://github.com/envoyproxy/envoy/issues/17540.
 	reqCtx := &RequestContext{
 		RequestState: RequestReceived,
+		Request: &Request{
+			Headers: make(map[string]string),
+			Body:    make(map[string]interface{}),
+		},
 	}
 
 	var body []byte
-	var requestBody, responseBody map[string]interface{}
+	var responseBody map[string]interface{}
 
 	// Create error handling var as each request should only report once for
 	// error metrics. This doesn't cover the error "Cannot receive stream request" because
@@ -167,15 +176,17 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 			// Message is buffered, we can read and decode.
 			if v.RequestBody.EndOfStream {
 				loggerTrace.Info("decoding")
-				err = json.Unmarshal(body, &requestBody)
+				err = json.Unmarshal(body, &reqCtx.Request.Body)
 				if err != nil {
 					logger.V(logutil.DEFAULT).Error(err, "Error unmarshaling request body")
+					// TODO: short circuit and send the body back as is (this could be an envoy error), currently we drop
+					// whatever the body request would have been and send our immediate response instead.
 				}
 
 				// Body stream complete. Allocate empty slice for response to use.
 				body = []byte{}
 
-				reqCtx, err = s.HandleRequestBody(ctx, reqCtx, req, requestBody)
+				reqCtx, err = s.HandleRequestBody(ctx, reqCtx)
 				if err != nil {
 					logger.V(logutil.DEFAULT).Error(err, "Error handling body")
 				} else {
@@ -256,7 +267,7 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 					loggerTrace.Info("stream completed")
 					// Don't send a 500 on a response error. Just let the message passthrough and log our error for debugging purposes.
 					// We assume the body is valid JSON, err messages are not guaranteed to be json, and so capturing and sending a 500 obfuscates the response message.
-					// using the standard 'err' var will send an immediate error response back to the caller.
+					// Using the standard 'err' var will send an immediate error response back to the caller.
 					var responseErr error
 					responseErr = json.Unmarshal(body, &responseBody)
 					if responseErr != nil {