Merge pull request ampproject#1066 from Automattic/fix/premature-ob-f…

…lush Detect when ob_flush() is called prematurely in the course of output-buffering document for AMP theme support
jritterbush · Apr 13, 2018 · f8f3ef5 · f8f3ef5
2 parents 6fe62ed + adf3b5d
commit f8f3ef5
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 1 deletion.
diff --git a/includes/class-amp-theme-support.php b/includes/class-amp-theme-support.php
@@ -1034,6 +1034,15 @@ 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 = ! empty( $_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
@@ -912,7 +912,7 @@ 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>';
+		echo '<html><img src="test.png"><script data-test>document.write(\'Illegal\');</script></html>';
 		AMP_Theme_Support::finish_output_buffering();
 		// get first layer buffer.
 		$output = ob_get_clean();
@@ -1070,6 +1070,25 @@ 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.
 	 *