From e570410159ebfc6b82eff3ca1c60c3586b8493f8 Mon Sep 17 00:00:00 2001
From: Greg Marshall <greg@alley.co>
Date: Wed, 1 May 2024 19:44:07 -0500
Subject: [PATCH] set response pieces correctly

---
 src/class-service.php | 39 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/class-service.php b/src/class-service.php
index 54e3267..40fe450 100644
--- a/src/class-service.php
+++ b/src/class-service.php
@@ -250,7 +250,44 @@ protected function get_response( WP_REST_Request $request, string $url, array $a
 		 */
 		$response = apply_filters( 'wp_proxy_service_response_after_request', $response, $request, $url );
 
-		return rest_ensure_response( $response );
+		// Check for a successful response.
+		if ( is_wp_error( $response ) ) {
+			return new WP_REST_Response(
+				[
+					'message' => 'Failed to fetch data',
+					'error'   => $response->get_error_message(),
+				],
+				500
+			);
+		}
+
+		// Get the response body.
+		$data = wp_remote_retrieve_body( $response );
+
+		// Check if decode was successful.
+		if ( null === $data ) {
+			return new WP_REST_Response(
+				[
+					'message' => 'Error getting body',
+				],
+				500
+			);
+		}
+
+		// Create a WP_REST_Response object.
+		$rest_response = new WP_REST_Response( $data );
+
+		// Set status code from the original HTTP response.
+		$rest_response->set_status( wp_remote_retrieve_response_code( $response ) );
+
+		// Set headers.
+		$headers = wp_remote_retrieve_headers( $response );
+		foreach ( $headers as $name => $value ) {
+			$rest_response->header( $name, $value );
+		}
+		$rest_response->header( 'X-Proxied-By', 'wp-proxy-service' );
+
+		return rest_ensure_response( $rest_response );
 	}
 
 	/**