Prepare AMP response at shutdown instead of output buffer callback (#931

) * Prepare AMP response at shutdown instead of output buffer callback
ampproject · Feb 5, 2018 · c293486 · c293486
1 parent 574b253
commit c293486
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 24 deletions.
diff --git a/includes/class-amp-theme-support.php b/includes/class-amp-theme-support.php
@@ -791,36 +791,53 @@ protected static function ensure_required_markup( DOMDocument $dom ) {
 
 	/**
 	 * Start output buffering.
+	 *
+	 * @since 0.7
+	 * @see AMP_Theme_Support::finish_output_buffering()
 	 */
 	public static function start_output_buffering() {
-		ob_start( array( __CLASS__, 'finish_output_buffering' ) );
+		ob_start();
+
+		// Note that the following must be at 0 because wp_ob_end_flush_all() runs at shutdown:1.
+		add_action( 'shutdown', array( __CLASS__, 'finish_output_buffering' ), 0 );
 	}
 
 	/**
 	 * Finish output buffering.
 	 *
-	 * @todo Do this in shutdown instead of output buffering callback? This will make it easier to debug things since printing can be done in shutdown function but cannot in output buffer callback.
+	 * @since 0.7
+	 * @see AMP_Theme_Support::start_output_buffering()
+	 */
+	public static function finish_output_buffering() {
+		echo self::prepare_response( ob_get_clean() ); // WPCS: xss ok.
+	}
+
+	/**
+	 * Process response to ensure AMP validity.
+	 *
+	 * @since 0.7
+	 *
+	 * @param string $response HTML document response.
+	 * @return string AMP document response.
 	 * @global int $content_width
-	 * @param string $output Buffered output.
-	 * @return string Finalized output.
 	 */
-	public static function finish_output_buffering( $output ) {
+	public static function prepare_response( $response ) {
 		global $content_width;
 
 		/*
 		 * Make sure that <meta charset> is present in output prior to parsing.
 		 * Note that the meta charset is supposed to appear within the first 1024 bytes.
 		 * See <https://www.w3.org/International/questions/qa-html-encoding-declarations>.
 		 */
-		if ( ! preg_match( '#<meta[^>]+charset=#i', substr( $output, 0, 1024 ) ) ) {
-			$output = preg_replace(
+		if ( ! preg_match( '#<meta[^>]+charset=#i', substr( $response, 0, 1024 ) ) ) {
+			$response = preg_replace(
 				'/(<head[^>]*>)/i',
 				'$1' . sprintf( '<meta charset="%s">', esc_attr( get_bloginfo( 'charset' ) ) ),
-				$output,
+				$response,
 				1
 			);
 		}
-		$dom  = AMP_DOM_Utils::get_dom( $output );
+		$dom  = AMP_DOM_Utils::get_dom( $response );
 		$args = array(
 			'content_max_width'    => ! empty( $content_width ) ? $content_width : AMP_Post_Template::CONTENT_MAX_WIDTH, // Back-compat.
 			'use_document_element' => true,
@@ -841,25 +858,25 @@ public static function finish_output_buffering( $output ) {
 			trigger_error( esc_html( sprintf( __( 'The database has the %s encoding when it needs to be utf-8 to work with AMP.', 'amp' ), get_bloginfo( 'charset' ) ), E_USER_WARNING ) ); // phpcs:ignore WordPress.PHP.DevelopmentFunctions.error_log_trigger_error
 		}
 
-		$output  = "<!DOCTYPE html>\n";
-		$output .= AMP_DOM_Utils::get_content_from_dom_node( $dom, $dom->documentElement );
+		$response  = "<!DOCTYPE html>\n";
+		$response .= AMP_DOM_Utils::get_content_from_dom_node( $dom, $dom->documentElement );
 
 		// Inject required scripts.
-		$output = preg_replace(
+		$response = preg_replace(
 			'#' . preg_quote( self::SCRIPTS_PLACEHOLDER, '#' ) . '#',
 			self::get_amp_scripts( $assets['scripts'] ),
-			$output,
+			$response,
 			1
 		);
 
 		// Inject styles.
-		$output = preg_replace(
+		$response = preg_replace(
 			'#' . preg_quote( self::STYLES_PLACEHOLDER, '#' ) . '#',
 			self::get_amp_styles( $assets['stylesheets'] ),
-			$output,
+			$response,
 			1
 		);
 
-		return $output;
+		return $response;
 	}
 }
diff --git a/tests/test-class-amp-theme-support.php b/tests/test-class-amp-theme-support.php
@@ -88,11 +88,11 @@ public function test_register_widgets() {
 	}
 
 	/**
-	 * Test finish_output_buffering.
+	 * Test prepare_response.
 	 *
-	 * @covers AMP_Theme_Support::finish_output_buffering()
+	 * @covers AMP_Theme_Support::prepare_response()
 	 */
-	public function test_finish_output_buffering() {
+	public function test_prepare_response() {
 		add_theme_support( 'amp' );
 		AMP_Theme_Support::init();
 		ob_start();
@@ -117,7 +117,7 @@ public function test_finish_output_buffering() {
 		</html>
 		<?php
 		$original_html  = trim( ob_get_clean() );
-		$sanitized_html = AMP_Theme_Support::finish_output_buffering( $original_html );
+		$sanitized_html = AMP_Theme_Support::prepare_response( $original_html );
 
 		$this->assertContains( '<meta charset="' . get_bloginfo( 'charset' ) . '">', $sanitized_html );
 		$this->assertContains( '<meta name="viewport" content="width=device-width,minimum-scale=1">', $sanitized_html );
@@ -137,11 +137,11 @@ public function test_finish_output_buffering() {
 	}
 
 	/**
-	 * Test finish_output_buffering to inject html[amp] attribute and ensure HTML5 doctype.
+	 * Test prepare_response to inject html[amp] attribute and ensure HTML5 doctype.
 	 *
-	 * @covers AMP_Theme_Support::finish_output_buffering()
+	 * @covers AMP_Theme_Support::prepare_response()
 	 */
-	public function test_finish_output_buffering_to_add_html5_doctype_and_amp_attribute() {
+	public function test_prepare_response_to_add_html5_doctype_and_amp_attribute() {
 		add_theme_support( 'amp' );
 		AMP_Theme_Support::init();
 		ob_start();
@@ -150,7 +150,7 @@ public function test_finish_output_buffering_to_add_html5_doctype_and_amp_attrib
 		<html><head><?php wp_head(); ?></head><body><?php wp_footer(); ?></body></html>
 		<?php
 		$original_html  = trim( ob_get_clean() );
-		$sanitized_html = AMP_Theme_Support::finish_output_buffering( $original_html );
+		$sanitized_html = AMP_Theme_Support::prepare_response( $original_html );
 
 		$this->assertStringStartsWith( '<!DOCTYPE html>', $sanitized_html );
 		$this->assertContains( '<html amp', $sanitized_html );