diff --git a/internal/record/recording_https_proxy.go b/internal/record/recording_https_proxy.go index 3f8449e..22a6541 100644 --- a/internal/record/recording_https_proxy.go +++ b/internal/record/recording_https_proxy.go @@ -125,6 +125,7 @@ func (r *RecordingHTTPSProxy) redactRequest(req *http.Request) (*store.RecordedR // Redacts secrets from header values r.redactor.Headers(recordedRequest.Headers) recordedRequest.Request = r.redactor.String(recordedRequest.Request) + recordedRequest.URL = r.redactor.String(recordedRequest.URL) var redactedBodySegments []map[string]any for _, bodySegment := range recordedRequest.BodySegments { redactedBodySegments = append(redactedBodySegments, r.redactor.Map(bodySegment)) diff --git a/internal/replay/replay_http_server.go b/internal/replay/replay_http_server.go index 0beaa89..ebb42dc 100644 --- a/internal/replay/replay_http_server.go +++ b/internal/replay/replay_http_server.go @@ -131,6 +131,7 @@ func (r *ReplayHTTPServer) createRedactedRequest(req *http.Request) (*store.Reco // Redacts secrets from header values r.redactor.Headers(recordedRequest.Headers) recordedRequest.Request = r.redactor.String(recordedRequest.Request) + recordedRequest.URL = r.redactor.String(recordedRequest.URL) var redactedBodySegments []map[string]any for _, bodySegment := range recordedRequest.BodySegments { redactedBodySegments = append(redactedBodySegments, r.redactor.Map(bodySegment)) @@ -179,13 +180,34 @@ func (r *ReplayHTTPServer) writeResponse(w http.ResponseWriter, resp *store.Reco w.WriteHeader(int(resp.StatusCode)) - jsonBytes, err := json.Marshal(resp.BodySegments[0]) - if err != nil { + if len(resp.BodySegments) == 1 { + jsonBytes, err := json.Marshal(resp.BodySegments[0]) + if err != nil { + return err + } + + _, err = w.Write(jsonBytes) return err + } else { + respToWrite := []byte{} + for _, bodySegment := range resp.BodySegments { + jsonBytes, err := json.Marshal(bodySegment) + if err != nil { + return err + } + + line := append([]byte("data: "), jsonBytes...) + line = append(line, []byte("\n\n")...) + + respToWrite = append(respToWrite, line...) + } + + if _, err := w.Write(respToWrite); err != nil { + return err + } } - _, err = w.Write(jsonBytes) - return err + return nil } func extractNumber(i *int, content string) (int, error) { diff --git a/internal/store/store.go b/internal/store/store.go index 1476caa..85a2237 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -17,6 +17,7 @@ limitations under the License. package store import ( + "bufio" "bytes" "compress/gzip" "crypto/sha256" @@ -175,16 +176,45 @@ func NewRecordedResponse(resp *http.Response, body []byte) (*RecordedResponse, e } + var bodySegments []map[string]any var bodySegment map[string]any err := json.Unmarshal(body, &bodySegment) if err != nil { - return nil, err + // Attempt to process streamed response. + prefix := []byte("data: ") + + reader := bytes.NewReader(body) + scanner := bufio.NewScanner(reader) + + for scanner.Scan() { + lineBytes := scanner.Bytes() + if len(lineBytes) == 0 { + continue + } + + if jsonBytes, ok := bytes.CutPrefix(lineBytes, prefix); ok { + var jsonMap map[string]any + if err := json.Unmarshal(jsonBytes, &jsonMap); err != nil { + log.Printf("Error unmarshaling JSON: %v", err) + continue + } + + bodySegments = append(bodySegments, jsonMap) + } + } + + if err := scanner.Err(); err != nil { + log.Fatalf("Error reading input bytes: %v", err) + return nil, err + } + } else { + bodySegments = append(bodySegments, bodySegment) } recordedResponse := &RecordedResponse{ StatusCode: int32(resp.StatusCode), Headers: GetHeadersMap(&resp.Header), - BodySegments: []map[string]any{bodySegment}, + BodySegments: bodySegments, } return recordedResponse, nil }