Introduce AMP_Content_Sanitizer::sanitize_document() to bypass saniti…

…ze always returning body content

includes/class-amp-theme-support.php

@@ -482,6 +482,11 @@ public static function finish_output_buffering( $output ) {
 			'content_max_width' => ! empty( $content_width ) ? $content_width : AMP_Post_Template::CONTENT_MAX_WIDTH, // Back-compat.
 		);
 
+		$assets = AMP_Content_Sanitizer::sanitize_document( $dom, self::$sanitizer_classes, $args );
+
+		self::$amp_scripts = array_merge( self::$amp_scripts, $assets['scripts'] );
+		self::$amp_styles  = array_merge( self::$amp_styles, $assets['styles'] );
+
 		/*
 		 * @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.
@@ -490,11 +495,7 @@ public static function finish_output_buffering( $output ) {
 		 * 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 );
-
+		$sanitized_inner_body = AMP_DOM_Utils::get_content_from_dom( $dom );
 		$output = preg_replace( '#(<body.*?>)(.+)(</body>)#si', '$1' . $sanitized_inner_body . '$3', $output );
 
 		// Inject required scripts.

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

@@ -7,27 +7,50 @@
 
 /**
  * Class AMP_Content_Sanitizer
+ *
+ * @since 0.4.1
  */
 class AMP_Content_Sanitizer {
 
 	/**
-	 * Sanitize.
+	 * Sanitize _content_.
 	 *
-	 * @param string|DOMDocument $content HTML content string or DOM document.
-	 * @param string[]           $sanitizer_classes Sanitizer classes.
-	 * @param array              $global_args       Global args.
+	 * @since 0.4.1
 	 *
-	 * @return array
+	 * @param string   $content HTML content string or DOM document.
+	 * @param string[] $sanitizer_classes Sanitizer classes.
+	 * @param array    $global_args       Global args.
+	 * @return array Tuple containing sanitized HTML, scripts array, and styles array.
 	 */
 	public static function sanitize( $content, array $sanitizer_classes, $global_args = array() ) {
+		$dom = AMP_DOM_Utils::get_dom_from_content( $content );
+
+		$results = self::sanitize_document( $dom, $sanitizer_classes, $global_args );
+		return array(
+			AMP_DOM_Utils::get_content_from_dom( $dom ),
+			$results['scripts'],
+			$results['styles'],
+		);
+	}
+
+	/**
+	 * Sanitize document.
+	 *
+	 * @since 0.7
+	 *
+	 * @param DOMDocument $dom               HTML document.
+	 * @param string[]    $sanitizer_classes Sanitizer classes.
+	 * @param array       $global_args       Global args passed into .
+	 * @return array {
+	 *     Scripts and styles needed by sanitizers.
+	 *
+	 *     @type array $scripts Scripts.
+	 *     @type array $styles  Styles.
+	 * }
+	 */
+	public static function sanitize_document( &$dom, $sanitizer_classes, $global_args ) {
 		$scripts = array();
 		$styles  = array();
-		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 ) ) {
 				/* translators: %s is sanitizer class */
@@ -54,9 +77,7 @@ public static function sanitize( $content, array $sanitizer_classes, $global_arg
 			$styles  = array_merge( $styles, $sanitizer->get_styles() );
 		}
 
-		$sanitized_content = AMP_DOM_Utils::get_content_from_dom( $dom );
-
-		return array( $sanitized_content, $scripts, $styles );
+		return compact( 'scripts', 'styles' );
 	}
 }
 

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

@@ -17,7 +17,7 @@ class AMP_DOM_Utils {
 	 *
 	 * Not all are valid AMP, but we include them for completeness.
 	 *
-	 * @since 0.6
+	 * @since 0.7
 	 * @link https://www.w3.org/TR/html5/syntax.html#serializing-html-fragments
 	 * @var array
 	 */