From 2534f647d389c78a3677992e21665ea19a6670b8 Mon Sep 17 00:00:00 2001
From: Alec Geatches <alec.geatches@automattic.com>
Date: Tue, 17 Dec 2024 08:50:09 -0700
Subject: [PATCH] Add auth header error handling (#250)

* Allow WP_Error return from get_request_headers()
* Bubble auth errors from get_request_headers() in Salesforce integration queries
* Early return error in Google Sheets' get_request_headers()
* Fix phpcs errors from "useless @inheritDoc."

---------

Co-authored-by: Max Schmeling <max.schmeling@automattic.com>
---
 docs/extending/data-source.md                   |  2 +-
 docs/extending/query.md                         |  2 +-
 .../queries/class-art-institute-data-source.php |  3 ++-
 inc/Config/DataSource/HttpDataSource.php        |  5 +----
 .../DataSource/HttpDataSourceInterface.php      |  4 +++-
 inc/Config/QueryContext/HttpQueryContext.php    |  3 ++-
 .../QueryContext/HttpQueryContextInterface.php  |  4 +++-
 inc/Config/QueryRunner/QueryRunner.php          |  5 +++++
 inc/ExampleApi/Queries/ExampleApiDataSource.php |  6 ++----
 .../Airtable/AirtableDataSource.php             |  3 ++-
 .../GenericHttp/GenericHttpDataSource.php       |  3 ++-
 inc/Integrations/GitHub/GitHubDataSource.php    |  3 ++-
 .../Google/Sheets/GoogleSheetsDataSource.php    |  6 +++++-
 .../SalesforceB2C/Auth/SalesforceB2CAuth.php    | 17 ++++++++++++-----
 .../Queries/SalesforceB2CGetProductQuery.php    | 15 +++++++--------
 .../SalesforceB2CSearchProductsQuery.php        | 15 +++++++--------
 .../SalesforceB2C/SalesforceB2CDataSource.php   |  3 ++-
 inc/Integrations/Shopify/ShopifyDataSource.php  |  3 ++-
 tests/inc/Mocks/MockDataSource.php              |  6 ++----
 19 files changed, 63 insertions(+), 45 deletions(-)

diff --git a/docs/extending/data-source.md b/docs/extending/data-source.md
index 80979eaf..c092a3ba 100644
--- a/docs/extending/data-source.md
+++ b/docs/extending/data-source.md
@@ -30,7 +30,7 @@ class ZipCodeDataSource extends HttpDataSource {
 		return 'https://api.zippopotam.us/us/';
 	}
 
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		return [
 			'Content-Type' => 'application/json',
 		];
diff --git a/docs/extending/query.md b/docs/extending/query.md
index c55b07b7..447722c8 100644
--- a/docs/extending/query.md
+++ b/docs/extending/query.md
@@ -163,7 +163,7 @@ By default, the `get_request_headers` method proxies to the `get_request_headers
 ### Example
 
 ```php
-public function get_request_headers( array $input_variables ): array {
+public function get_request_headers( array $input_variables ): array|WP_Error {
 	return array_merge(
 		$this->get_data_source()->get_request_headers(),
 		[ 'X-Product-ID' => $input_variables['product_id'] ]
diff --git a/example/rest-api/art-institute/inc/queries/class-art-institute-data-source.php b/example/rest-api/art-institute/inc/queries/class-art-institute-data-source.php
index b715174a..d8bdb6da 100644
--- a/example/rest-api/art-institute/inc/queries/class-art-institute-data-source.php
+++ b/example/rest-api/art-institute/inc/queries/class-art-institute-data-source.php
@@ -3,6 +3,7 @@
 namespace RemoteDataBlocks\Example\ArtInstituteOfChicago;
 
 use RemoteDataBlocks\Config\DataSource\HttpDataSource;
+use WP_Error;
 
 class ArtInstituteOfChicagoDataSource extends HttpDataSource {
 	public function get_display_name(): string {
@@ -13,7 +14,7 @@ public function get_endpoint(): string {
 		return 'https://api.artic.edu/api/v1/artworks';
 	}
 
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		return [
 			'Content-Type' => 'application/json',
 		];
diff --git a/inc/Config/DataSource/HttpDataSource.php b/inc/Config/DataSource/HttpDataSource.php
index ccafad53..64d5dd14 100644
--- a/inc/Config/DataSource/HttpDataSource.php
+++ b/inc/Config/DataSource/HttpDataSource.php
@@ -27,10 +27,7 @@ abstract public function get_display_name(): string;
 
 	abstract public function get_endpoint(): string;
 
-	/**
-	 * @inheritDoc
-	 */
-	abstract public function get_request_headers(): array;
+	abstract public function get_request_headers(): array|WP_Error;
 
 	public function get_image_url(): ?string {
 		return null;
diff --git a/inc/Config/DataSource/HttpDataSourceInterface.php b/inc/Config/DataSource/HttpDataSourceInterface.php
index 92715e31..0c126201 100644
--- a/inc/Config/DataSource/HttpDataSourceInterface.php
+++ b/inc/Config/DataSource/HttpDataSourceInterface.php
@@ -2,6 +2,8 @@
 
 namespace RemoteDataBlocks\Config\DataSource;
 
+use WP_Error;
+
 /**
  * HttpDataSourceInterface
  *
@@ -38,5 +40,5 @@ public function get_endpoint(): string;
 	 *
 	 * @return array Associative array of request headers.
 	 */
-	public function get_request_headers(): array;
+	public function get_request_headers(): array|WP_Error;
 }
diff --git a/inc/Config/QueryContext/HttpQueryContext.php b/inc/Config/QueryContext/HttpQueryContext.php
index 2abfcec1..d1d2b170 100644
--- a/inc/Config/QueryContext/HttpQueryContext.php
+++ b/inc/Config/QueryContext/HttpQueryContext.php
@@ -9,6 +9,7 @@
 use RemoteDataBlocks\Config\QueryRunner\QueryRunnerInterface;
 use RemoteDataBlocks\Validation\Validator;
 use RemoteDataBlocks\Validation\ValidatorInterface;
+use WP_Error;
 
 defined( 'ABSPATH' ) || exit();
 
@@ -175,7 +176,7 @@ public function get_request_method(): string {
 	 *
 	 * @param array $input_variables The input variables for this query.
 	 */
-	public function get_request_headers( array $input_variables ): array {
+	public function get_request_headers( array $input_variables ): array|WP_Error {
 		return $this->get_data_source()->get_request_headers();
 	}
 
diff --git a/inc/Config/QueryContext/HttpQueryContextInterface.php b/inc/Config/QueryContext/HttpQueryContextInterface.php
index b45f610d..0b13d770 100644
--- a/inc/Config/QueryContext/HttpQueryContextInterface.php
+++ b/inc/Config/QueryContext/HttpQueryContextInterface.php
@@ -2,6 +2,8 @@
 
 namespace RemoteDataBlocks\Config\QueryContext;
 
+use WP_Error;
+
 /**
  * HttpQueryContextInterface interface
  *
@@ -10,6 +12,6 @@ interface HttpQueryContextInterface {
 	public function get_cache_ttl( array $input_variables ): null|int;
 	public function get_endpoint( array $input_variables ): string;
 	public function get_request_method(): string;
-	public function get_request_headers( array $input_variables ): array;
+	public function get_request_headers( array $input_variables ): array|WP_Error;
 	public function get_request_body( array $input_variables ): array|null;
 }
diff --git a/inc/Config/QueryRunner/QueryRunner.php b/inc/Config/QueryRunner/QueryRunner.php
index b868cedf..9ccf78e0 100644
--- a/inc/Config/QueryRunner/QueryRunner.php
+++ b/inc/Config/QueryRunner/QueryRunner.php
@@ -39,6 +39,11 @@ public function __construct(
 	 */
 	protected function get_request_details( array $input_variables ): array|WP_Error {
 		$headers = $this->query_context->get_request_headers( $input_variables );
+
+		if ( is_wp_error( $headers ) ) {
+			return $headers;
+		}
+
 		$method = $this->query_context->get_request_method();
 		$body = $this->query_context->get_request_body( $input_variables );
 		$endpoint = $this->query_context->get_endpoint( $input_variables );
diff --git a/inc/ExampleApi/Queries/ExampleApiDataSource.php b/inc/ExampleApi/Queries/ExampleApiDataSource.php
index f49bbae2..d002bc52 100644
--- a/inc/ExampleApi/Queries/ExampleApiDataSource.php
+++ b/inc/ExampleApi/Queries/ExampleApiDataSource.php
@@ -3,6 +3,7 @@
 namespace RemoteDataBlocks\ExampleApi\Queries;
 
 use RemoteDataBlocks\Config\DataSource\HttpDataSource;
+use WP_Error;
 
 /**
  * This is a placeholder DataSource used only to represent the data source in the
@@ -21,10 +22,7 @@ public function get_endpoint(): string {
 		return '';
 	}
 
-	/**
-	 * @inheritDoc
-	 */
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		return [];
 	}
 }
diff --git a/inc/Integrations/Airtable/AirtableDataSource.php b/inc/Integrations/Airtable/AirtableDataSource.php
index 5a916f60..aa251685 100644
--- a/inc/Integrations/Airtable/AirtableDataSource.php
+++ b/inc/Integrations/Airtable/AirtableDataSource.php
@@ -3,6 +3,7 @@
 namespace RemoteDataBlocks\Integrations\Airtable;
 
 use RemoteDataBlocks\Config\DataSource\HttpDataSource;
+use WP_Error;
 
 class AirtableDataSource extends HttpDataSource {
 	protected const SERVICE_SCHEMA_VERSION = 1;
@@ -82,7 +83,7 @@ public function get_endpoint(): string {
 		return 'https://api.airtable.com/v0/' . $this->config['base']['id'];
 	}
 
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		return [
 			'Authorization' => sprintf( 'Bearer %s', $this->config['access_token'] ),
 			'Content-Type' => 'application/json',
diff --git a/inc/Integrations/GenericHttp/GenericHttpDataSource.php b/inc/Integrations/GenericHttp/GenericHttpDataSource.php
index 5636f414..335f699b 100644
--- a/inc/Integrations/GenericHttp/GenericHttpDataSource.php
+++ b/inc/Integrations/GenericHttp/GenericHttpDataSource.php
@@ -3,6 +3,7 @@
 namespace RemoteDataBlocks\Integrations\GenericHttp;
 
 use RemoteDataBlocks\Config\DataSource\HttpDataSource;
+use WP_Error;
 
 class GenericHttpDataSource extends HttpDataSource {
 	protected const SERVICE_NAME = REMOTE_DATA_BLOCKS_GENERIC_HTTP_SERVICE;
@@ -62,7 +63,7 @@ public function get_endpoint(): string {
 		return $this->config['url'];
 	}
 
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		return [
 			'Accept' => 'application/json',
 		];
diff --git a/inc/Integrations/GitHub/GitHubDataSource.php b/inc/Integrations/GitHub/GitHubDataSource.php
index 31a3ffda..be0e8785 100644
--- a/inc/Integrations/GitHub/GitHubDataSource.php
+++ b/inc/Integrations/GitHub/GitHubDataSource.php
@@ -3,6 +3,7 @@
 namespace RemoteDataBlocks\Integrations\GitHub;
 
 use RemoteDataBlocks\Config\DataSource\HttpDataSource;
+use WP_Error;
 
 class GitHubDataSource extends HttpDataSource {
 	protected const SERVICE_NAME = REMOTE_DATA_BLOCKS_GITHUB_SERVICE;
@@ -42,7 +43,7 @@ public function get_endpoint(): string {
 		);
 	}
 
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		return [
 			'Accept' => 'application/vnd.github+json',
 		];
diff --git a/inc/Integrations/Google/Sheets/GoogleSheetsDataSource.php b/inc/Integrations/Google/Sheets/GoogleSheetsDataSource.php
index bfec06cf..7ef75274 100644
--- a/inc/Integrations/Google/Sheets/GoogleSheetsDataSource.php
+++ b/inc/Integrations/Google/Sheets/GoogleSheetsDataSource.php
@@ -80,12 +80,16 @@ public function get_endpoint(): string {
 		return sprintf( 'https://sheets.googleapis.com/v4/spreadsheets/%s', $this->config['spreadsheet']['id'] );
 	}
 
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		$access_token = GoogleAuth::generate_token_from_service_account_key(
 			$this->config['credentials'],
 			GoogleAuth::GOOGLE_SHEETS_SCOPES
 		);
 
+		if ( is_wp_error( $access_token ) ) {
+			return $access_token;
+		}
+
 		return [
 			'Authorization' => sprintf( 'Bearer %s', $access_token ),
 			'Content-Type' => 'application/json',
diff --git a/inc/Integrations/SalesforceB2C/Auth/SalesforceB2CAuth.php b/inc/Integrations/SalesforceB2C/Auth/SalesforceB2CAuth.php
index 458d985d..7e733436 100644
--- a/inc/Integrations/SalesforceB2C/Auth/SalesforceB2CAuth.php
+++ b/inc/Integrations/SalesforceB2C/Auth/SalesforceB2CAuth.php
@@ -18,26 +18,28 @@ class SalesforceB2CAuth {
 	 * @param string $organization_id The organization ID for the data source.
 	 * @param string $client_id The client ID (a version 4 UUID).
 	 * @param string $client_secret The client secret.
-	 * @return WP_Error|string The token or an error.
+	 * @return string|WP_Error The token or an error.
 	 */
 	public static function generate_token(
 		string $endpoint,
 		string $organization_id,
 		string $client_id,
 		string $client_secret
-	): WP_Error|string {
+	): string|WP_Error {
 		$saved_access_token = self::get_saved_access_token( $organization_id, $client_id );
 
 		if ( null !== $saved_access_token ) {
 			return $saved_access_token;
 		}
 
+		$access_token = null;
+
 		$saved_refresh_token = self::get_saved_refresh_token( $organization_id, $client_id );
 		if ( null !== $saved_refresh_token ) {
 			$access_token = self::get_token_using_refresh_token( $saved_refresh_token, $client_id, $client_secret, $endpoint, $organization_id );
 		}
 
-		if ( null !== $access_token ) {
+		if ( null !== $access_token && ! is_wp_error( $access_token ) ) {
 			return $access_token;
 		}
 
@@ -123,8 +125,13 @@ public static function get_token_using_refresh_token(
 			],
 		]);
 
+		$refresh_token_error = new WP_Error(
+			'salesforce_b2c_auth_error_refresh_token',
+			__( 'Failed to refresh authorization with refresh token', 'remote-data-blocks' )
+		);
+
 		if ( is_wp_error( $client_auth_response ) ) {
-			return null;
+			return $refresh_token_error;
 		}
 
 		$response_code = wp_remote_retrieve_response_code( $client_auth_response );
@@ -132,7 +139,7 @@ public static function get_token_using_refresh_token(
 		$response_data = json_decode( $response_body, true );
 
 		if ( 400 === $response_code ) {
-			return null;
+			return $refresh_token_error;
 		}
 
 		$access_token = $response_data['access_token'];
diff --git a/inc/Integrations/SalesforceB2C/Queries/SalesforceB2CGetProductQuery.php b/inc/Integrations/SalesforceB2C/Queries/SalesforceB2CGetProductQuery.php
index 22e95118..d1ae366b 100644
--- a/inc/Integrations/SalesforceB2C/Queries/SalesforceB2CGetProductQuery.php
+++ b/inc/Integrations/SalesforceB2C/Queries/SalesforceB2CGetProductQuery.php
@@ -4,6 +4,7 @@
 
 use RemoteDataBlocks\Config\QueryContext\HttpQueryContext;
 use RemoteDataBlocks\Integrations\SalesforceB2C\Auth\SalesforceB2CAuth;
+use WP_Error;
 
 class SalesforceB2CGetProductQuery extends HttpQueryContext {
 
@@ -60,7 +61,7 @@ public function get_output_schema(): array {
 		];
 	}
 
-	public function get_request_headers( array $input_variables ): array {
+	public function get_request_headers( array $input_variables ): array|WP_Error {
 		$data_source_config = $this->get_data_source()->to_array();
 		$data_source_endpoint = $this->get_data_source()->get_endpoint();
 
@@ -71,16 +72,14 @@ public function get_request_headers( array $input_variables ): array {
 			$data_source_config['client_secret']
 		);
 
-		$headers = [
-			'Content-Type' => 'application/json',
-		];
-
 		if ( is_wp_error( $access_token ) ) {
-			return $headers;
+			return $access_token;
 		}
 
-		$headers['Authorization'] = sprintf( 'Bearer %s', $access_token );
-		return $headers;
+		return [
+			'Content-Type' => 'application/json',
+			'Authorization' => sprintf( 'Bearer %s', $access_token ),
+		];
 	}
 
 	public function get_endpoint( array $input_variables ): string {
diff --git a/inc/Integrations/SalesforceB2C/Queries/SalesforceB2CSearchProductsQuery.php b/inc/Integrations/SalesforceB2C/Queries/SalesforceB2CSearchProductsQuery.php
index 64ac9cb7..b465eb34 100644
--- a/inc/Integrations/SalesforceB2C/Queries/SalesforceB2CSearchProductsQuery.php
+++ b/inc/Integrations/SalesforceB2C/Queries/SalesforceB2CSearchProductsQuery.php
@@ -4,6 +4,7 @@
 
 use RemoteDataBlocks\Config\QueryContext\HttpQueryContext;
 use RemoteDataBlocks\Integrations\SalesforceB2C\Auth\SalesforceB2CAuth;
+use WP_Error;
 
 class SalesforceB2CSearchProductsQuery extends HttpQueryContext {
 	public function get_input_schema(): array {
@@ -55,7 +56,7 @@ public function get_endpoint( array $input_variables ): string {
 		);
 	}
 
-	public function get_request_headers( array $input_variables ): array {
+	public function get_request_headers( array $input_variables ): array|WP_Error {
 		$data_source_config = $this->get_data_source()->to_array();
 		$data_source_endpoint = $this->get_data_source()->get_endpoint();
 
@@ -66,16 +67,14 @@ public function get_request_headers( array $input_variables ): array {
 			$data_source_config['client_secret']
 		);
 
-		$headers = [
-			'Content-Type' => 'application/json',
-		];
-
 		if ( is_wp_error( $access_token ) ) {
-			return $headers;
+			return $access_token;
 		}
 
-		$headers['Authorization'] = sprintf( 'Bearer %s', $access_token );
-		return $headers;
+		return [
+			'Content-Type' => 'application/json',
+			'Authorization' => sprintf( 'Bearer %s', $access_token ),
+		];
 	}
 
 	public function get_query_name(): string {
diff --git a/inc/Integrations/SalesforceB2C/SalesforceB2CDataSource.php b/inc/Integrations/SalesforceB2C/SalesforceB2CDataSource.php
index b4985438..fb9133c6 100644
--- a/inc/Integrations/SalesforceB2C/SalesforceB2CDataSource.php
+++ b/inc/Integrations/SalesforceB2C/SalesforceB2CDataSource.php
@@ -3,6 +3,7 @@
 namespace RemoteDataBlocks\Integrations\SalesforceB2C;
 
 use RemoteDataBlocks\Config\DataSource\HttpDataSource;
+use WP_Error;
 
 use function plugins_url;
 
@@ -42,7 +43,7 @@ public function get_endpoint(): string {
 		return sprintf( 'https://%s.api.commercecloud.salesforce.com', $this->config['shortcode'] );
 	}
 
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		return [
 			'Content-Type' => 'application/json',
 		];
diff --git a/inc/Integrations/Shopify/ShopifyDataSource.php b/inc/Integrations/Shopify/ShopifyDataSource.php
index abe6bfc9..3e60e7ff 100644
--- a/inc/Integrations/Shopify/ShopifyDataSource.php
+++ b/inc/Integrations/Shopify/ShopifyDataSource.php
@@ -3,6 +3,7 @@
 namespace RemoteDataBlocks\Integrations\Shopify;
 
 use RemoteDataBlocks\Config\DataSource\HttpDataSource;
+use WP_Error;
 
 use function plugins_url;
 
@@ -40,7 +41,7 @@ public function get_endpoint(): string {
 		return 'https://' . $this->config['store_name'] . '.myshopify.com/api/2024-04/graphql.json';
 	}
 
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		return [
 			'Content-Type' => 'application/json',
 			'X-Shopify-Storefront-Access-Token' => $this->config['access_token'],
diff --git a/tests/inc/Mocks/MockDataSource.php b/tests/inc/Mocks/MockDataSource.php
index b6beb2c1..a6af2560 100644
--- a/tests/inc/Mocks/MockDataSource.php
+++ b/tests/inc/Mocks/MockDataSource.php
@@ -3,6 +3,7 @@
 namespace RemoteDataBlocks\Tests\Mocks;
 
 use RemoteDataBlocks\Config\DataSource\HttpDataSource;
+use WP_Error;
 
 class MockDataSource extends HttpDataSource {
 	private $endpoint = 'https://example.com/api';
@@ -31,10 +32,7 @@ public function get_endpoint(): string {
 		return $this->endpoint;
 	}
 
-	/**
-	 * @inheritDoc
-	 */
-	public function get_request_headers(): array {
+	public function get_request_headers(): array|WP_Error {
 		return $this->headers;
 	}