Merge branch '0.7' of https://github.com/Automattic/amp-wp into develop

ampproject · Apr 19, 2018 · 31016b1 · 31016b1
2 parents ac53ee7 + d65578d
commit 31016b1
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 29 deletions.
diff --git a/includes/class-amp-theme-support.php b/includes/class-amp-theme-support.php
@@ -75,6 +75,14 @@ class AMP_Theme_Support {
 	 */
 	public static $init_start_time;
 
+	/**
+	 * Output buffering level when starting.
+	 *
+	 * @since 0.7
+	 * @var int
+	 */
+	protected static $initial_ob_level = 0;
+
 	/**
 	 * Initialize.
 	 *
@@ -933,6 +941,7 @@ public static function start_output_buffering() {
 		}
 
 		ob_start();
+		self::$initial_ob_level = ob_get_level();
 
 		// 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 );
@@ -946,6 +955,12 @@ public static function start_output_buffering() {
 	 */
 	public static function finish_output_buffering() {
 		AMP_Response_Headers::send_server_timing( 'amp_output_buffer', -self::$init_start_time, 'AMP Output Buffer' );
+
+		// Flush output buffer stack until we get to the output buffer we started.
+		while ( ob_get_level() > self::$initial_ob_level ) {
+			ob_end_flush();
+		}
+
 		echo self::prepare_response( ob_get_clean() ); // WPCS: xss ok.
 	}
 
@@ -1000,15 +1015,6 @@ public static function prepare_response( $response, $args = array() ) {
 			return $response;
 		}
 
-		// Account for case where ob_flush() was called prematurely.
-		if ( false === strpos( $response, '<html' ) ) {
-			$error = sprintf(
-				'<div style="color:red; background: white; padding: 0.5em; position: fixed; z-index: 100000; bottom: 0; border: dashed 1px red;">%s</div>',
-				wp_kses_post( __( '<strong>AMP Plugin Error</strong>: It appears that your WordPress install prematurely flushed the output buffer. You will need to disable AMP theme support until that is fixed.', 'amp' ) )
-			);
-			return $error . $response;
-		}
-
 		$is_validation_debug_mode = isset( $_REQUEST[ AMP_Validation_Utils::DEBUG_QUERY_VAR ] ); // WPCS: csrf ok.
 
 		$args = array_merge(

diff --git a/tests/test-class-amp-theme-support.php b/tests/test-class-amp-theme-support.php
@@ -875,6 +875,7 @@ function newrelic_disable_autorum() {
 
 		$this->assertEquals( 0, has_action( 'shutdown', array( self::TESTED_CLASS, 'finish_output_buffering' ) ) );
 		$this->assertTrue( ob_get_level() > $ob_level );
+
 		// End output buffer.
 		if ( ob_get_level() > $ob_level ) {
 			ob_get_clean();
@@ -894,12 +895,24 @@ public function test_finish_output_buffering() {
 		// start first layer buffer.
 		ob_start();
 		AMP_Theme_Support::start_output_buffering();
-		echo '<html><img src="test.png"><script data-test>document.write(\'Illegal\');</script></html>';
+
+		echo '<img src="test.png"><script data-test>document.write(\'Illegal\');</script>';
+
+		// Additional nested output bufferings which aren't getting closed.
+		ob_start();
+		echo 'foo';
+		ob_start( function( $response ) {
+			return strtoupper( $response );
+		} );
+		echo 'bar';
+
 		AMP_Theme_Support::finish_output_buffering();
 		// get first layer buffer.
 		$output = ob_get_clean();
 
 		$this->assertContains( '<html amp', $output );
+		$this->assertContains( 'foo', $output );
+		$this->assertContains( 'BAR', $output );
 		$this->assertContains( '<amp-img src="test.png"', $output );
 		$this->assertNotContains( '<script data-test', $output );
 
@@ -1051,25 +1064,6 @@ public function test_prepare_response_to_add_html5_doctype_and_amp_attribute() {
 		$this->assertContains( '<html amp', $sanitized_html );
 	}
 
-	/**
-	 * Test preparing an incomplete response due to ob_flush().
-	 *
-	 * @covers AMP_Theme_Support::prepare_response()
-	 */
-	public function test_prepare_response_premature_ob_flush() {
-		add_theme_support( 'amp' );
-		AMP_Theme_Support::init();
-		ob_start();
-		wp_footer();
-		echo '</body></html>';
-		$original_html  = trim( ob_get_clean() );
-		$sanitized_html = AMP_Theme_Support::prepare_response( $original_html );
-
-		$this->assertContains( 'AMP Plugin Error', $sanitized_html );
-		$this->assertNotContains( '<!DOCTYPE html>', $sanitized_html );
-		$this->assertNotContains( '<html amp', $sanitized_html );
-	}
-
 	/**
 	 * Data provider for test_ensure_required_markup.
 	 *