Sanitize entire HTML output when theme support is present

includes/class-amp-theme-support.php

@@ -191,8 +191,6 @@ public static function register_hooks() {
 		 */
 		add_action( 'template_redirect', array( __CLASS__, 'start_output_buffering' ), 0 );
 
-		add_filter( 'the_content', array( __CLASS__, 'filter_the_content' ), PHP_INT_MAX );
-
 		// @todo Add character conversion.
 	}
 
@@ -422,32 +420,12 @@ public static function get_amp_custom_styles() {
 		return $css;
 	}
 
-	/**
-	 * Filter the content to be valid AMP.
-	 *
-	 * @param string $content Content.
-	 * @return string Amplified content.
-	 */
-	public static function filter_the_content( $content ) {
-		$args = array(
-			'content_max_width' => ! empty( $content_width ) ? $content_width : AMP_Post_Template::CONTENT_MAX_WIDTH, // Back-compat.
-		);
-
-		list( $sanitized_content, $scripts, $styles ) = AMP_Content_Sanitizer::sanitize( $content, self::$sanitizer_classes, $args );
-
-		self::$amp_scripts = array_merge( self::$amp_scripts, $scripts );
-		self::$amp_styles  = array_merge( self::$amp_styles, $styles );
-
-		return $sanitized_content;
-	}
-
 	/**
 	 * Determine required AMP scripts.
 	 *
-	 * @param string $html Output HTML.
 	 * @return string Scripts to inject into the HEAD.
 	 */
-	public static function get_amp_component_scripts( $html ) {
+	public static function get_amp_component_scripts() {
 		$amp_scripts = self::$amp_scripts;
 
 		foreach ( self::$embed_handlers as $embed_handler ) {
@@ -492,10 +470,32 @@ public static function start_output_buffering() {
 	 * Finish output buffering.
 	 *
 	 * @todo Do this in shutdown instead of output buffering callback?
+	 * @global int $content_width
 	 * @param string $output Buffered output.
 	 * @return string Finalized output.
 	 */
 	public static function finish_output_buffering( $output ) {
+		global $content_width;
+
+		$dom  = AMP_DOM_Utils::get_dom( $output );
+		$args = array(
+			'content_max_width' => ! empty( $content_width ) ? $content_width : AMP_Post_Template::CONTENT_MAX_WIDTH, // Back-compat.
+		);
+
+		/*
+		 * @todo The sanitize method needs to be updated to sanitize the entire HTML element and not just the BODY.
+		 * This will require updating mandatory_parent_blacklist in amphtml-update.py to include elements that appear in the HEAD.
+		 * This will ensure that the scripts and styles that plugins output via wp_head() will be sanitized as well. However,
+		 * since the the old paired mode is sending content from the *body* we'll need to be able to filter out the elements
+		 * from outside the body from being part of the whitelist sanitizer when it runs when theme support is not present,
+		 * as otherwise elements from the HEAD could get added to the BODY.
+		 */
+		list( $sanitized_inner_body, $scripts, $styles ) = AMP_Content_Sanitizer::sanitize( $dom, self::$sanitizer_classes, $args );
+
+		self::$amp_scripts = array_merge( self::$amp_scripts, $scripts );
+		self::$amp_styles  = array_merge( self::$amp_styles, $styles );
+
+		$output = preg_replace( '#(<body.*?>)(.+)(</body>)#si', '$1' . $sanitized_inner_body . '$3', $output );
 
 		// Inject required scripts.
 		$output = preg_replace(
@@ -513,7 +513,6 @@ public static function finish_output_buffering( $output ) {
 			1
 		);
 
-		// @todo Add more validation checking and potentially the whitelist sanitizer.
 		return $output;
 	}
 }

includes/templates/class-amp-content-sanitizer.php

@@ -13,16 +13,20 @@ class AMP_Content_Sanitizer {
 	/**
 	 * Sanitize.
 	 *
-	 * @param string   $content           Content.
-	 * @param string[] $sanitizer_classes Sanitizer classes.
-	 * @param array    $global_args       Global args.
+	 * @param string|DOMDocument $content HTML content string or DOM document.
+	 * @param string[]           $sanitizer_classes Sanitizer classes.
+	 * @param array              $global_args       Global args.
 	 *
 	 * @return array
 	 */
 	public static function sanitize( $content, array $sanitizer_classes, $global_args = array() ) {
 		$scripts = array();
 		$styles  = array();
-		$dom     = AMP_DOM_Utils::get_dom_from_content( $content );
+		if ( $content instanceof DOMDocument ) {
+			$dom = $content;
+		} else {
+			$dom = AMP_DOM_Utils::get_dom_from_content( $content );
+		}
 
 		foreach ( $sanitizer_classes as $sanitizer_class => $args ) {
 			if ( ! class_exists( $sanitizer_class ) ) {
@@ -31,6 +35,11 @@ public static function sanitize( $content, array $sanitizer_classes, $global_arg
 				continue;
 			}
 
+			/**
+			 * Sanitizer.
+			 *
+			 * @type AMP_Base_Sanitizer $sanitizer
+			 */
 			$sanitizer = new $sanitizer_class( $dom, array_merge( $global_args, $args ) );
 
 			if ( ! is_subclass_of( $sanitizer, 'AMP_Base_Sanitizer' ) ) {

includes/utils/class-amp-dom-utils.php

@@ -88,6 +88,8 @@ public static function get_content_from_dom( $dom ) {
 
 		/**
 		 * We only want children of the body tag, since we have a subset of HTML.
+		 *
+		 * @todo We will want to get the full HTML eventually.
 		 */
 		$body = $dom->getElementsByTagName( 'body' )->item( 0 );
 

tests/test-class-amp-theme-support.php

@@ -67,4 +67,53 @@ public function test_is_paired_available() {
 		$this->assertTrue( is_search() );
 		$this->assertFalse( AMP_Theme_Support::is_paired_available() );
 	}
+
+	/**
+	 * Test finish_output_buffering.
+	 *
+	 * @covers AMP_Theme_Support::finish_output_buffering()
+	 */
+	public function test_finish_output_buffering() {
+		add_theme_support( 'amp' );
+		AMP_Theme_Support::init();
+		ob_start();
+		?>
+		<!DOCTYPE html>
+		<html amp <?php language_attributes(); ?>>
+			<head>
+				<?php wp_head(); ?>
+				<script data-head>document.write('TODO: This needs to be sanitized as well once.');</script>
+			</head>
+			<body>
+				<img width="100" height="100" src="https://example.com/test.png">
+				<audio width="400" height="300" src="https://example.com/audios/myaudio.mp3"></audio>
+				<amp-ad type="a9"
+					width="300"
+					height="250"
+					data-aax_size="300x250"
+					data-aax_pubname="test123"
+					data-aax_src="302"></amp-ad>
+				<?php wp_footer(); ?>
+			</body>
+		</html>
+		<?php
+		$original_html  = trim( ob_get_clean() );
+		$sanitized_html = AMP_Theme_Support::finish_output_buffering( $original_html );
+
+		$this->assertContains( '<meta charset="utf-8">', $sanitized_html );
+		$this->assertContains( '<meta name="viewport" content="width=device-width,minimum-scale=1">', $sanitized_html );
+		$this->assertContains( '<style amp-boilerplate>', $sanitized_html );
+		$this->assertContains( '<style amp-custom>', $sanitized_html );
+		$this->assertContains( '<script async src="https://cdn.ampproject.org/v0.js"', $sanitized_html ); // phpcs:ignore WordPress.WP.EnqueuedResources.NonEnqueuedScript
+		$this->assertContains( '<meta name="generator" content="AMP Plugin', $sanitized_html );
+
+		$this->assertNotContains( '<img', $sanitized_html );
+		$this->assertContains( '<amp-img', $sanitized_html );
+
+		$this->assertNotContains( '<audio', $sanitized_html );
+		$this->assertContains( '<amp-audio', $sanitized_html );
+		$this->assertContains( '<script async custom-element="amp-audio"', $sanitized_html );
+
+		$this->assertContains( '<script async custom-element="amp-ad"', $sanitized_html );
+	}
 }