Merge pull request #947 from Automattic/fix/handling-response-for-non…

…-html Fix handling of response when it is not HTML
ampproject · Feb 9, 2018 · b8d9a2c · b8d9a2c
2 parents 9b46771 + 2e8b4a2
commit b8d9a2c
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 2 deletions.
diff --git a/includes/class-amp-theme-support.php b/includes/class-amp-theme-support.php
@@ -772,6 +772,15 @@ public static function finish_output_buffering() {
 	public static function prepare_response( $response, $args = array() ) {
 		global $content_width;
 
+		/*
+		 * Check if the response starts with HTML markup.
+		 * Without this check, JSON responses will be erroneously corrupted,
+		 * being wrapped in HTML documents.
+		 */
+		if ( '<' !== substr( ltrim( $response ), 0, 1 ) ) {
+			return $response;
+		}
+
 		$args = array_merge(
 			array(
 				'content_max_width'       => ! empty( $content_width ) ? $content_width : AMP_Post_Template::CONTENT_MAX_WIDTH, // Back-compat.

diff --git a/includes/sanitizers/class-amp-style-sanitizer.php b/includes/sanitizers/class-amp-style-sanitizer.php
@@ -196,7 +196,12 @@ public function sanitize() {
 			if ( ! $this->amp_custom_style_element ) {
 				$this->amp_custom_style_element = $this->dom->createElement( 'style' );
 				$this->amp_custom_style_element->setAttribute( 'amp-custom', '' );
-				$this->dom->getElementsByTagName( 'head' )->item( 0 )->appendChild( $this->amp_custom_style_element );
+				$head = $this->dom->getElementsByTagName( 'head' )->item( 0 );
+				if ( ! $head ) {
+					$head = $this->dom->createElement( 'head' );
+					$this->dom->documentElement->insertBefore( $head, $this->dom->documentElement->firstChild );
+				}
+				$head->appendChild( $this->amp_custom_style_element );
 			}
 
 			// Gather stylesheets to print as long as they don't surpass the limit.

diff --git a/tests/test-amp-style-sanitizer.php b/tests/test-amp-style-sanitizer.php
@@ -160,7 +160,7 @@ public function get_link_and_style_test_data() {
 					's {color:yellow}',
 				),
 			),
-			'style_eleemnts_with_link_elements' => array(
+			'style_elements_with_link_elements' => array(
 				sprintf(
 					'<html amp><head><meta charset="utf-8"><style type="text/css">strong.before-dashicon {color:green}</style><link rel="stylesheet" href="%s"><style type="text/css">strong.after-dashicon {color:green}</style></head><body><style>s {color:yellow !important}</style></body></html>', // phpcs:ignore WordPress.WP.EnqueuedResources.NonEnqueuedStylesheet
 					includes_url( 'css/dashicons.css' )
@@ -172,6 +172,12 @@ public function get_link_and_style_test_data() {
 					's {color:yellow}',
 				),
 			),
+			'style_with_no_head' => array(
+				'<html amp><body>Not good!<style>body{color:red}</style></body>',
+				array(
+					'body{color:red}',
+				),
+			),
 		);
 	}
 

diff --git a/tests/test-class-amp-theme-support.php b/tests/test-class-amp-theme-support.php
@@ -138,6 +138,29 @@ public function test_prepare_response() {
 		$this->assertContains( '<script async custom-element="amp-ad"', $sanitized_html );
 	}
 
+	/**
+	 * Test prepare_response for bad/non-HTML.
+	 *
+	 * @covers AMP_Theme_Support::prepare_response()
+	 */
+	public function test_prepare_response_bad_html() {
+		add_theme_support( 'amp' );
+		AMP_Theme_Support::init();
+
+		// JSON.
+		$input = '{"success":true}';
+		$this->assertEquals( $input, AMP_Theme_Support::prepare_response( $input ) );
+
+		// Nothing, for redirect.
+		$input = '';
+		$this->assertEquals( $input, AMP_Theme_Support::prepare_response( $input ) );
+
+		// HTML, but very stripped down.
+		$input  = '<html>Hello</html>';
+		$output = AMP_Theme_Support::prepare_response( $input );
+		$this->assertContains( '<html amp', $output );
+	}
+
 	/**
 	 * Test prepare_response to inject html[amp] attribute and ensure HTML5 doctype.
 	 *