Account for ob_start() being called during page execution

jritterbush · Apr 19, 2018 · e142186 · e142186
1 parent b5bb51e
commit e142186
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/includes/class-amp-theme-support.php b/includes/class-amp-theme-support.php
@@ -76,6 +76,14 @@ class AMP_Theme_Support {
 	 */
 	public static $headers_sent = array();
 
+	/**
+	 * Output buffering level when starting.
+	 *
+	 * @since 0.7
+	 * @var int
+	 */
+	protected static $initial_ob_level = 0;
+
 	/**
 	 * Initialize.
 	 */
@@ -969,6 +977,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 );
@@ -981,6 +990,12 @@ public static function start_output_buffering() {
 	 * @see AMP_Theme_Support::start_output_buffering()
 	 */
 	public static function finish_output_buffering() {
+
+		// 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.
 	}
 

diff --git a/tests/test-class-amp-theme-support.php b/tests/test-class-amp-theme-support.php
@@ -894,6 +894,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();
@@ -913,12 +914,24 @@ public function test_finish_output_buffering() {
 		// start first layer buffer.
 		ob_start();
 		AMP_Theme_Support::start_output_buffering();
+
 		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 );