From 0154620e624db02944920681999e34b7149d7abe Mon Sep 17 00:00:00 2001
From: Weston Ruter <weston@xwp.co>
Date: Mon, 12 Mar 2018 00:33:05 -0700
Subject: [PATCH] Move any AMP component scripts from body into head

Call ensure_required_markup prior to sanitizing document
---
 includes/class-amp-theme-support.php   | 24 ++++++++++++++++++------
 tests/test-class-amp-theme-support.php |  9 ++++++++-
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/includes/class-amp-theme-support.php b/includes/class-amp-theme-support.php
index c08b84b904e..806eb01e163 100644
--- a/includes/class-amp-theme-support.php
+++ b/includes/class-amp-theme-support.php
@@ -750,10 +750,18 @@ public static function print_amp_styles() {
 	 * canonical URL by default if a singular post.
 	 *
 	 * @since 0.7
+	 * @todo All of this might be better placed inside of a sanitizer.
 	 *
 	 * @param DOMDocument $dom Doc.
 	 */
 	protected static function ensure_required_markup( DOMDocument $dom ) {
+		$xpath = new DOMXPath( $dom );
+
+		// First ensure the mandatory amp attribute is present on the html element, as otherwise it will be stripped entirely.
+		if ( ! $dom->documentElement->hasAttribute( 'amp' ) && ! $dom->documentElement->hasAttribute( '⚡️' ) ) {
+			$dom->documentElement->setAttribute( 'amp', '' );
+		}
+
 		$head = $dom->getElementsByTagName( 'head' )->item( 0 );
 		if ( ! $head ) {
 			$head = $dom->createElement( 'head' );
@@ -803,6 +811,15 @@ protected static function ensure_required_markup( DOMDocument $dom ) {
 			) );
 			$head->appendChild( $rel_canonical );
 		}
+
+		// Make sure scripts from the body get moved to the head.
+		$scripts = array();
+		foreach ( $xpath->query( '//body//script[ @custom-element or @custom-template ]' ) as $script ) {
+			$scripts[] = $script;
+		}
+		foreach ( $scripts as $script ) {
+			$head->appendChild( $script );
+		}
 	}
 
 	/**
@@ -939,15 +956,10 @@ public static function prepare_response( $response, $args = array() ) {
 		}
 		$dom = AMP_DOM_Utils::get_dom( $response );
 
-		// First ensure the mandatory amp attribute is present on the html element, as otherwise it will be stripped entirely.
-		if ( ! $dom->documentElement->hasAttribute( 'amp' ) && ! $dom->documentElement->hasAttribute( '⚡️' ) ) {
-			$dom->documentElement->setAttribute( 'amp', '' );
-		}
+		self::ensure_required_markup( $dom );
 
 		$assets = AMP_Content_Sanitizer::sanitize_document( $dom, self::$sanitizer_classes, $args );
 
-		self::ensure_required_markup( $dom );
-
 		// @todo If 'utf-8' is not the blog charset, then we'll need to do some character encoding conversation or "entityification".
 		if ( 'utf-8' !== strtolower( get_bloginfo( 'charset' ) ) ) {
 			/* translators: %s is the charset of the current site */
diff --git a/tests/test-class-amp-theme-support.php b/tests/test-class-amp-theme-support.php
index 7a8d1a60122..a6a3be40ce8 100644
--- a/tests/test-class-amp-theme-support.php
+++ b/tests/test-class-amp-theme-support.php
@@ -109,8 +109,14 @@ public function test_prepare_response() {
 		wp_widgets_init();
 
 		add_action( 'wp_enqueue_scripts', function() {
-			wp_enqueue_script( 'amp-mathml' );
+			wp_enqueue_script( 'amp-list' );
 		} );
+		add_action( 'wp_footer', function() {
+			wp_print_scripts( 'amp-mathml' );
+			?>
+			<amp-mathml layout="container" data-formula="\[x = {-b \pm \sqrt{b^2-4ac} \over 2a}.\]"></amp-mathml>
+			<?php
+		}, 1 );
 
 		ob_start();
 		?>
@@ -150,6 +156,7 @@ public function test_prepare_response() {
 		$this->assertContains( '<style amp-boilerplate>', $sanitized_html );
 		$this->assertContains( '<style amp-custom>body { background: black; }', $sanitized_html );
 		$this->assertContains( '<script type="text/javascript" src="https://cdn.ampproject.org/v0.js" async></script>', $sanitized_html ); // phpcs:ignore WordPress.WP.EnqueuedResources.NonEnqueuedScript
+		$this->assertContains( '<script type="text/javascript" src="https://cdn.ampproject.org/v0/amp-list-latest.js" async custom-element="amp-list"></script>', $sanitized_html ); // phpcs:ignore WordPress.WP.EnqueuedResources.NonEnqueuedScript
 		$this->assertContains( '<script type="text/javascript" src="https://cdn.ampproject.org/v0/amp-mathml-latest.js" async custom-element="amp-mathml"></script>', $sanitized_html ); // phpcs:ignore WordPress.WP.EnqueuedResources.NonEnqueuedScript
 		$this->assertContains( '<meta name="generator" content="AMP Plugin', $sanitized_html );