Merge pull request #1102 from Automattic/fix/incorporate-shutdown-output

Move any content output during shutdown to be injected before closing body tag
ampproject · Jun 20, 2018 · c99e7de · c99e7de
2 parents 4364ae8 + b5948c5
commit c99e7de
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/includes/class-amp-theme-support.php b/includes/class-amp-theme-support.php
@@ -1191,6 +1191,22 @@ public static function prepare_response( $response, $args = array() ) {
 		$dom  = AMP_DOM_Utils::get_dom( $response );
 		$head = $dom->getElementsByTagName( 'head' )->item( 0 );
 
+		// Move anything after </html>, such as Query Monitor output added at shutdown, to be moved before </body>.
+		$body = $dom->getElementsByTagName( 'body' )->item( 0 );
+		if ( $body ) {
+			while ( $dom->documentElement->nextSibling ) {
+				// Trailing elements after </html> will get wrapped in additional <html> elements.
+				if ( 'html' === $dom->documentElement->nextSibling->nodeName ) {
+					while ( $dom->documentElement->nextSibling->firstChild ) {
+						$body->appendChild( $dom->documentElement->nextSibling->firstChild );
+					}
+					$dom->removeChild( $dom->documentElement->nextSibling );
+				} else {
+					$body->appendChild( $dom->documentElement->nextSibling );
+				}
+			}
+		}
+
 		// Make sure scripts from the body get moved to the head.
 		if ( isset( $head ) ) {
 			$xpath = new DOMXPath( $dom );

diff --git a/tests/test-class-amp-theme-support.php b/tests/test-class-amp-theme-support.php
@@ -1102,7 +1102,7 @@ public function test_prepare_response() {
 				<?php wp_head(); ?>
 				<script data-head>document.write('Illegal');</script>
 			</head>
-			<body>
+			<body><!-- </body></html> -->
 				<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"
@@ -1111,13 +1111,17 @@ public function test_prepare_response() {
 					data-aax_size="300x250"
 					data-aax_pubname="test123"
 					data-aax_src="302"></amp-ad>
+
 				<?php wp_footer(); ?>
 
 				<button onclick="alert('Illegal');">no-onclick</button>
 
 				<style>body { background: black; }</style>
 			</body>
 		</html>
+		<!--comment-after-html-->
+		<div id="after-html"></div>
+		<!--comment-end-html-->
 		<?php
 		$original_html  = trim( ob_get_clean() );
 		$sanitized_html = AMP_Theme_Support::prepare_response( $original_html );
@@ -1165,6 +1169,9 @@ public function test_prepare_response() {
 			$removed_nodes
 		);
 
+		// Make sure trailing content after </html> gets moved.
+		$this->assertRegExp( '#<!--comment-after-html-->\s*<div id="after-html"></div>\s*<!--comment-end-html-->\s*</body>\s*</html>\s*$#s', $sanitized_html );
+
 		$prepare_response_args = array(
 			'enable_response_caching' => true,
 		);