Refactor Store API hydration logic and prevent fatal errors from session class usage

src/Assets/AssetDataRegistry.php

@@ -2,7 +2,7 @@
 namespace Automattic\WooCommerce\Blocks\Assets;
 
 use Automattic\WooCommerce\Blocks\Package;
-
+use Automattic\WooCommerce\Blocks\Domain\Services\Hydration;
 use Exception;
 use InvalidArgumentException;
 
@@ -317,16 +317,40 @@ public function add( $key, $data, $check_key_exists = false ) {
 	}
 
 	/**
-	 * Hydrate from API.
+	 * Hydrate from the API.
 	 *
 	 * @param string $path REST API path to preload.
 	 */
 	public function hydrate_api_request( $path ) {
 		if ( ! isset( $this->preloaded_api_requests[ $path ] ) ) {
-			$this->preloaded_api_requests = rest_preload_api_request( $this->preloaded_api_requests, $path );
+			$this->preloaded_api_requests[ $path ] = Package::container()->get( Hydration::class )->get_rest_api_response_data( $path );
 		}
 	}
 
+	/**
+	 * Hydrate some data from the API.
+	 *
+	 * @param string  $key  The key used to reference the data being registered.
+	 * @param string  $path REST API path to preload.
+	 * @param boolean $check_key_exists If set to true, duplicate data will be ignored if the key exists.
+	 *                                  If false, duplicate data will cause an exception.
+	 *
+	 * @throws InvalidArgumentException  Only throws when site is in debug mode. Always logs the error.
+	 */
+	public function hydrate_data_from_api_request( $key, $path, $check_key_exists = false ) {
+		$this->add(
+			$key,
+			function() use ( $path ) {
+				if ( isset( $this->preloaded_api_requests[ $path ], $this->preloaded_api_requests[ $path ]['body'] ) ) {
+					return $this->preloaded_api_requests[ $path ]['body'];
+				}
+				$response = Package::container()->get( Hydration::class )->get_rest_api_response_data( $path );
+				return $response['body'] ?? '';
+			},
+			$check_key_exists
+		);
+	}
+
 	/**
 	 * Adds a page permalink to the data registry.
 	 *

src/BlockTypes/AllProducts.php

@@ -31,20 +31,12 @@ protected function enqueue_data( array $attributes = [] ) {
 		$this->asset_data_registry->add( 'max_rows', wc_get_theme_support( 'product_blocks::max_rows', 6 ), true );
 		$this->asset_data_registry->add( 'default_rows', wc_get_theme_support( 'product_blocks::default_rows', 3 ), true );
 
-		// Hydrate the following data depending on admin or frontend context.
+		// Hydrate the All Product block with data from the API. This is for the add to cart buttons which show current quantity in cart, and events.
 		if ( ! is_admin() && ! WC()->is_rest_api_request() ) {
-			$this->hydrate_from_api();
+			$this->asset_data_registry->hydrate_api_request( '/wc/store/v1/cart' );
 		}
 	}
 
-	/**
-	 * Hydrate the All Product block with data from the API. This is for the add to cart buttons which show current
-	 * quantity in cart, and events.
-	 */
-	protected function hydrate_from_api() {
-		$this->asset_data_registry->hydrate_api_request( '/wc/store/v1/cart' );
-	}
-
 	/**
 	 * It is necessary to register and enqueue assets during the render phase because we want to load assets only if the block has the content.
 	 */

src/BlockTypes/Cart.php

@@ -239,7 +239,7 @@ protected function enqueue_data( array $attributes = [] ) {
 
 		// Hydrate the following data depending on admin or frontend context.
 		if ( ! is_admin() && ! WC()->is_rest_api_request() ) {
-			$this->hydrate_from_api();
+			$this->asset_data_registry->hydrate_api_request( '/wc/store/v1/cart' );
 		}
 
 		/**
@@ -250,19 +250,6 @@ protected function enqueue_data( array $attributes = [] ) {
 		do_action( 'woocommerce_blocks_cart_enqueue_data' );
 	}
 
-	/**
-	 * Hydrate the cart block with data from the API.
-	 */
-	protected function hydrate_from_api() {
-		// Cache existing notices now, otherwise they are caught by the Cart Controller and converted to exceptions.
-		$old_notices = WC()->session->get( 'wc_notices', array() );
-		wc_clear_notices();
-
-		$this->asset_data_registry->hydrate_api_request( '/wc/store/v1/cart' );
-
-		// Restore notices.
-		WC()->session->set( 'wc_notices', $old_notices );
-	}
 	/**
 	 * Register script and style assets for the block type before it is registered.
 	 *

src/BlockTypes/Checkout.php

@@ -323,7 +323,8 @@ function( $acc, $method ) {
 		}
 
 		if ( ! is_admin() && ! WC()->is_rest_api_request() ) {
-			$this->hydrate_from_api();
+			$this->asset_data_registry->hydrate_api_request( '/wc/store/v1/cart' );
+			$this->asset_data_registry->hydrate_data_from_api_request( 'checkoutData', '/wc/store/v1/checkout' );
 			$this->hydrate_customer_payment_methods();
 		}
 
@@ -391,25 +392,6 @@ function( $saved_payment_method ) use ( $payment_gateways ) {
 		remove_filter( 'woocommerce_payment_methods_list_item', [ $this, 'include_token_id_with_payment_methods' ], 10, 2 );
 	}
 
-	/**
-	 * Hydrate the checkout block with data from the API.
-	 */
-	protected function hydrate_from_api() {
-		// Cache existing notices now, otherwise they are caught by the Cart Controller and converted to exceptions.
-		$old_notices = WC()->session->get( 'wc_notices', array() );
-		wc_clear_notices();
-
-		$this->asset_data_registry->hydrate_api_request( '/wc/store/v1/cart' );
-
-		add_filter( 'woocommerce_store_api_disable_nonce_check', '__return_true' );
-		$rest_preload_api_requests = rest_preload_api_request( [], '/wc/store/v1/checkout' );
-		$this->asset_data_registry->add( 'checkoutData', $rest_preload_api_requests['/wc/store/v1/checkout']['body'] ?? [] );
-		remove_filter( 'woocommerce_store_api_disable_nonce_check', '__return_true' );
-
-		// Restore notices.
-		WC()->session->set( 'wc_notices', $old_notices );
-	}
-
 	/**
 	 * Callback for woocommerce_payment_methods_list_item filter to add token id
 	 * to the generated list.

src/BlockTypesController.php

@@ -1,7 +1,6 @@
 <?php
 namespace Automattic\WooCommerce\Blocks;
 
-use Automattic\WooCommerce\Blocks\BlockTypes\AtomicBlock;
 use Automattic\WooCommerce\Blocks\Package;
 use Automattic\WooCommerce\Blocks\Assets\AssetDataRegistry;
 use Automattic\WooCommerce\Blocks\Assets\Api as AssetApi;
@@ -61,10 +60,10 @@ public function register_blocks() {
 		$block_types = $this->get_block_types();
 
 		foreach ( $block_types as $block_type ) {
-			$block_type_class    = __NAMESPACE__ . '\\BlockTypes\\' . $block_type;
-			$block_type_instance = new $block_type_class( $this->asset_api, $this->asset_data_registry, new IntegrationRegistry() );
-		}
+			$block_type_class = __NAMESPACE__ . '\\BlockTypes\\' . $block_type;
 
+			new $block_type_class( $this->asset_api, $this->asset_data_registry, new IntegrationRegistry() );
+		}
 	}
 
 	/**

src/Domain/Bootstrap.php

@@ -12,6 +12,7 @@
 use Automattic\WooCommerce\Blocks\Domain\Services\DraftOrders;
 use Automattic\WooCommerce\Blocks\Domain\Services\FeatureGating;
 use Automattic\WooCommerce\Blocks\Domain\Services\GoogleAnalytics;
+use Automattic\WooCommerce\Blocks\Domain\Services\Hydration;
 use Automattic\WooCommerce\Blocks\InboxNotifications;
 use Automattic\WooCommerce\Blocks\Installer;
 use Automattic\WooCommerce\Blocks\Migration;
@@ -353,6 +354,12 @@ function( Container $container ) {
 				return new Notices( $container->get( Package::class ) );
 			}
 		);
+		$this->container->register(
+			Hydration::class,
+			function( Container $container ) {
+				return new Hydration( $container->get( AssetDataRegistry::class ) );
+			}
+		);
 		$this->container->register(
 			PaymentsApi::class,
 			function ( Container $container ) {

src/Domain/Services/Hydration.php

@@ -0,0 +1,97 @@
+<?php
+namespace Automattic\WooCommerce\Blocks\Domain\Services;
+
+use Automattic\WooCommerce\Blocks\Assets\AssetDataRegistry;
+
+/**
+ * Service class that handles hydration of API data for blocks.
+ */
+class Hydration {
+	/**
+	 * Instance of the asset data registry.
+	 *
+	 * @var AssetDataRegistry
+	 */
+	protected $asset_data_registry;
+
+	/**
+	 * Cached notices to restore after hydrating the API.
+	 *
+	 * @var array
+	 */
+	protected $cached_store_notices = [];
+
+	/**
+	 * Constructor.
+	 *
+	 * @param AssetDataRegistry $asset_data_registry Instance of the asset data registry.
+	 */
+	public function __construct( AssetDataRegistry $asset_data_registry ) {
+		$this->asset_data_registry = $asset_data_registry;
+	}
+
+	/**
+	 * Hydrates the asset data registry with data from the API. Disables notices and nonces so requests contain valid
+	 * data that is not polluted by the current session.
+	 *
+	 * @param array $path API paths to hydrate e.g. '/wc/store/v1/cart'.
+	 * @return array Response data.
+	 */
+	public function get_rest_api_response_data( $path = '' ) {
+		$this->cache_store_notices();
+		$this->disable_nonce_check();
+
+		// Preload the request and add it to the array. It will be $preloaded_requests['path']  and contain 'body' and 'headers'.
+		$preloaded_requests = rest_preload_api_request( [], $path );
+
+		$this->restore_cached_store_notices();
+		$this->restore_nonce_check();
+
+		// Returns just the single preloaded request.
+		return $preloaded_requests[ $path ];
+	}
+
+	/**
+	 * Disable the nonce check temporarily.
+	 */
+	protected function disable_nonce_check() {
+		add_filter( 'woocommerce_store_api_disable_nonce_check', [ $this, 'disable_nonce_check_callback' ] );
+	}
+
+	/**
+	 * Callback to disable the nonce check. While we could use `__return_true`, we use a custom named callback so that
+	 * we can remove it later without affecting other filters.
+	 */
+	public function disable_nonce_check_callback() {
+		return true;
+	}
+
+	/**
+	 * Restore the nonce check.
+	 */
+	protected function restore_nonce_check() {
+		remove_filter( 'woocommerce_store_api_disable_nonce_check', [ $this, 'disable_nonce_check_callback' ] );
+	}
+
+	/**
+	 * Cache notices before hydrating the API if the customer has a session.
+	 */
+	protected function cache_store_notices() {
+		if ( ! did_action( 'woocommerce_init' ) || null === WC()->session ) {
+			return;
+		}
+		$this->cached_store_notices = WC()->session->get( 'wc_notices', array() );
+		WC()->session->set( 'wc_notices', null );
+	}
+
+	/**
+	 * Restore notices into current session from cache.
+	 */
+	protected function restore_cached_store_notices() {
+		if ( ! did_action( 'woocommerce_init' ) || null === WC()->session ) {
+			return;
+		}
+		WC()->session->set( 'wc_notices', $this->cached_store_notices );
+		$this->cached_store_notices = [];
+	}
+}