From 739a8f752d8483256dd9d33fa0921a9f33713e59 Mon Sep 17 00:00:00 2001
From: Weston Ruter <westonruter@google.com>
Date: Fri, 20 Dec 2019 12:26:29 -0800
Subject: [PATCH 1/3] Introduce ampdevmode data flag on script/style
 dependencies to indicate data-ampdevmode needed

---
 includes/amp-helper-functions.php          | 22 +++++++-
 includes/class-amp-theme-support.php       | 66 ++++++++++++++--------
 tests/php/test-amp-helper-functions.php    |  1 -
 tests/php/test-class-amp-theme-support.php | 56 ++++++++++++------
 4 files changed, 101 insertions(+), 44 deletions(-)

diff --git a/includes/amp-helper-functions.php b/includes/amp-helper-functions.php
index 0bfca629a3c..e7b3d9cc27a 100644
--- a/includes/amp-helper-functions.php
+++ b/includes/amp-helper-functions.php
@@ -960,8 +960,28 @@ function amp_get_content_sanitizers( $post = null ) {
 		if ( is_admin_bar_showing() ) {
 			$dev_mode_xpaths[] = '//*[ @id = "wpadminbar" ]';
 			$dev_mode_xpaths[] = '//*[ @id = "wpadminbar" ]//*';
-			$dev_mode_xpaths[] = '//style[ @id = "admin-bar-inline-css" ]';
 		}
+
+		// Ensure script localization data gets flagged for dev mode. This only applies to wp_localize_script() as
+		// inline scripts added via wp_add_inline_script() get filtered by script_loader_tag and thus will have the
+		// data-ampdevmode attribute added via AMP_Theme_Support::filter_script_loader_tag_for_dev_mode().
+		foreach ( wp_scripts()->done as $script_handle ) { // @todo Fail! This is currently being called _before_ the template is rendered, so it will always be empty.
+			if ( ! AMP_Theme_Support::dependency_needs_dev_mode( wp_scripts(), $script_handle ) ) {
+				continue;
+			}
+			$data = wp_scripts()->get_data( $script_handle, 'data' );
+			if ( preg_match( '/(\bvar\s*\w+\s+=)/', $data, $matches ) ) {
+				$dev_mode_xpaths[] = sprintf( '//script[ not( @src ) ][ contains( text(), "%s" ) ]', $matches[1] );
+			}
+		}
+
+		// Ensure all inline styles added via wp_add_inline_style() get the data-ampdevmode attribute.
+		foreach ( wp_styles()->done as $style_handle ) { // @todo Fail! This is currently being called _before_ the template is rendered, so it will always be empty.
+			if ( AMP_Theme_Support::dependency_needs_dev_mode( wp_styles(), $style_handle ) ) {
+				$dev_mode_xpaths[] = sprintf( '//style[ @id = "%s" ]', "$style_handle-inline-css" );
+			}
+		}
+
 		$sanitizers = array_merge(
 			[
 				'AMP_Dev_Mode_Sanitizer' => [
diff --git a/includes/class-amp-theme-support.php b/includes/class-amp-theme-support.php
index b9fe763e303..a2bd4017e7d 100644
--- a/includes/class-amp-theme-support.php
+++ b/includes/class-amp-theme-support.php
@@ -1098,6 +1098,8 @@ static function( $html ) {
 		);
 
 		add_action( 'admin_bar_init', [ __CLASS__, 'init_admin_bar' ] );
+		add_filter( 'style_loader_tag', [ __CLASS__, 'filter_style_loader_tag_for_dev_mode' ], 10, 2 );
+		add_filter( 'script_loader_tag', [ __CLASS__, 'filter_script_loader_tag_for_dev_mode' ], 10, 2 );
 		add_action( 'wp_head', 'amp_add_generator_metadata', 20 );
 		add_action( 'wp_enqueue_scripts', [ __CLASS__, 'enqueue_assets' ], 0 ); // Enqueue before theme's styles.
 		add_action( 'wp_enqueue_scripts', [ __CLASS__, 'dequeue_customize_preview_scripts' ], 1000 );
@@ -1427,8 +1429,6 @@ public static function filter_cancel_comment_reply_link( $formatted_link, $link,
 	 * @since 1.0
 	 */
 	public static function init_admin_bar() {
-		add_filter( 'style_loader_tag', [ __CLASS__, 'filter_admin_bar_style_loader_tag' ], 10, 2 );
-		add_filter( 'script_loader_tag', [ __CLASS__, 'filter_admin_bar_script_loader_tag' ], 10, 2 );
 
 		// Inject the data-ampdevmode attribute into the admin bar bump style. See \WP_Admin_Bar::initialize().
 		if ( current_theme_supports( 'admin-bar' ) ) {
@@ -1551,40 +1551,53 @@ protected static function is_exclusively_dependent( WP_Dependencies $dependencie
 	}
 
 	/**
-	 * Add data-ampdevmode attribute to any enqueued style that depends on the admin-bar.
+	 * Determine whether a given dependency handle needs dev mode.
 	 *
-	 * @since 1.3
+	 * @since 1.5
+	 *
+	 * @param WP_Dependencies $dependencies Dependencies (wither WP_Scripts or WP_Styles).
+	 * @param string          $handle       Dependency handle (for script or style).
+	 * @return bool Whether the <script>, <link>, or <style> needs dev mode.
+	 */
+	public static function dependency_needs_dev_mode( WP_Dependencies $dependencies, $handle ) {
+		return (
+			$dependencies->get_data( $handle, 'ampdevmode' )
+			||
+			(
+			in_array( $handle, $dependencies->registered['admin-bar']->deps, true ) ?
+				self::is_exclusively_dependent( $dependencies, $handle, 'admin-bar' ) :
+				self::has_dependency( $dependencies, $handle, 'admin-bar' )
+			)
+		);
+	}
+
+	/**
+	 * Add data-ampdevmode attribute to any enqueued style that is flagged for dev mode or which depends on the admin-bar.
+	 *
+	 * @since 1.5
 	 *
 	 * @param string $tag    The link tag for the enqueued style.
 	 * @param string $handle The style's registered handle.
 	 * @return string Tag.
 	 */
-	public static function filter_admin_bar_style_loader_tag( $tag, $handle ) {
-		if (
-			in_array( $handle, wp_styles()->registered['admin-bar']->deps, true ) ?
-				self::is_exclusively_dependent( wp_styles(), $handle, 'admin-bar' ) :
-				self::has_dependency( wp_styles(), $handle, 'admin-bar' )
-		) {
+	public static function filter_style_loader_tag_for_dev_mode( $tag, $handle ) {
+		if ( self::dependency_needs_dev_mode( wp_styles(), $handle ) ) {
 			$tag = preg_replace( '/(?<=<link)(?=\s|>)/i', ' ' . AMP_Rule_Spec::DEV_MODE_ATTRIBUTE, $tag );
 		}
 		return $tag;
 	}
 
 	/**
-	 * Add data-ampdevmode attribute to any enqueued script that depends on the admin-bar.
+	 * Add data-ampdevmode attribute to any enqueued script that is flagged for dev mode or which depends on the admin-bar.
 	 *
-	 * @since 1.3
+	 * @since 1.5
 	 *
 	 * @param string $tag    The `<script>` tag for the enqueued script.
 	 * @param string $handle The script's registered handle.
 	 * @return string Tag.
 	 */
-	public static function filter_admin_bar_script_loader_tag( $tag, $handle ) {
-		if (
-			in_array( $handle, wp_scripts()->registered['admin-bar']->deps, true ) ?
-				self::is_exclusively_dependent( wp_scripts(), $handle, 'admin-bar' ) :
-				self::has_dependency( wp_scripts(), $handle, 'admin-bar' )
-		) {
+	public static function filter_script_loader_tag_for_dev_mode( $tag, $handle ) {
+		if ( self::dependency_needs_dev_mode( wp_scripts(), $handle ) ) {
 			$tag = preg_replace( '/(?<=<script)(?=\s|>)/i', ' ' . AMP_Rule_Spec::DEV_MODE_ATTRIBUTE, $tag );
 		}
 		return $tag;
@@ -2467,23 +2480,26 @@ public static function setup_paired_browsing_client() {
 		$asset        = require $asset_file;
 		$dependencies = $asset['dependencies'];
 		$version      = $asset['version'];
+		$handle       = 'amp-paired-browsing-client';
 
 		wp_enqueue_script(
-			'amp-paired-browsing-client',
+			$handle,
 			amp_get_asset_url( '/js/amp-paired-browsing-client.js' ),
 			$dependencies,
 			$version,
 			true
 		);
+		wp_script_add_data( $handle, 'ampdevmode', true );
+		foreach ( $dependencies as $dependency ) {
+			wp_script_add_data( $dependency, 'ampdevmode', true );
+		}
 
-		// Whitelist enqueued script for AMP dev mdoe so that it is not removed.
-		// @todo Revisit with <https://github.com/google/site-kit-wp/pull/505#discussion_r348683617>.
 		add_filter(
 			'script_loader_tag',
-			static function( $tag, $handle ) {
-				if ( is_amp_endpoint() && self::has_dependency( wp_scripts(), 'amp-paired-browsing-client', $handle ) ) {
-					$attrs = [ AMP_Rule_Spec::DEV_MODE_ATTRIBUTE, 'async' ];
-					$tag   = preg_replace( '/(?<=<script)(?=\s|>)/i', ' ' . implode( ' ', $attrs ), $tag );
+			static function( $tag, $filtered_handle ) use ( $handle ) {
+				if ( is_amp_endpoint() && self::has_dependency( wp_scripts(), $handle, $filtered_handle ) ) {
+					// Inject async attribute into script tag.
+					$tag = preg_replace( '/(<script[^>]*)(?=\ssrc=)/i', '$1 async ', $tag );
 				}
 				return $tag;
 			},
diff --git a/tests/php/test-amp-helper-functions.php b/tests/php/test-amp-helper-functions.php
index 1b58739a657..3d8f6167438 100644
--- a/tests/php/test-amp-helper-functions.php
+++ b/tests/php/test-amp-helper-functions.php
@@ -794,7 +794,6 @@ function ( $xpaths ) use ( $element_xpaths ) {
 				[
 					'//*[ @id = "wpadminbar" ]',
 					'//*[ @id = "wpadminbar" ]//*',
-					'//style[ @id = "admin-bar-inline-css" ]',
 				]
 			),
 			$sanitizers['AMP_Dev_Mode_Sanitizer']['element_xpaths']
diff --git a/tests/php/test-class-amp-theme-support.php b/tests/php/test-class-amp-theme-support.php
index 77e3a2486c1..6550aad82cd 100644
--- a/tests/php/test-class-amp-theme-support.php
+++ b/tests/php/test-class-amp-theme-support.php
@@ -996,6 +996,8 @@ static function( $templates ) {
 	 * @covers AMP_Theme_Support::add_hooks()
 	 */
 	public function test_add_hooks() {
+		add_theme_support( 'amp' );
+		$this->go_to( '/' );
 		AMP_Theme_Support::add_hooks();
 		$this->assertFalse( has_action( 'wp_head', 'wp_post_preview_js' ) );
 		$this->assertFalse( has_action( 'wp_head', 'wp_oembed_add_host_js' ) );
@@ -1016,6 +1018,8 @@ public function test_add_hooks() {
 		$this->assertEquals( 10, has_filter( 'customize_partial_render', [ self::TESTED_CLASS, 'filter_customize_partial_render' ] ) );
 		$this->assertEquals( 10, has_action( 'wp_footer', 'amp_print_analytics' ) );
 		$this->assertEquals( 10, has_action( 'admin_bar_init', [ self::TESTED_CLASS, 'init_admin_bar' ] ) );
+		$this->assertEquals( 10, has_filter( 'style_loader_tag', [ 'AMP_Theme_Support', 'filter_style_loader_tag_for_dev_mode' ] ) );
+		$this->assertEquals( 10, has_filter( 'script_loader_tag', [ 'AMP_Theme_Support', 'filter_script_loader_tag_for_dev_mode' ] ) );
 		$priority = defined( 'PHP_INT_MIN' ) ? PHP_INT_MIN : ~PHP_INT_MAX; // phpcs:ignore PHPCompatibility.Constants.NewConstants.php_int_minFound
 		$this->assertEquals( $priority, has_action( 'template_redirect', [ self::TESTED_CLASS, 'start_output_buffering' ] ) );
 
@@ -1237,8 +1241,8 @@ public function test_filter_cancel_comment_reply_link() {
 	 * Test init_admin_bar.
 	 *
 	 * @covers \AMP_Theme_Support::init_admin_bar()
-	 * @covers \AMP_Theme_Support::filter_admin_bar_style_loader_tag()
-	 * @covers \AMP_Theme_Support::filter_admin_bar_script_loader_tag()
+	 * @covers \AMP_Theme_Support::filter_style_loader_tag_for_dev_mode()
+	 * @covers \AMP_Theme_Support::filter_script_loader_tag_for_dev_mode()
 	 */
 	public function test_init_admin_bar() {
 		require_once ABSPATH . WPINC . '/class-wp-admin-bar.php';
@@ -1269,8 +1273,7 @@ function() {
 		$this->assertEquals( 10, has_action( 'wp_head', $callback ) );
 
 		AMP_Theme_Support::init_admin_bar();
-		$this->assertEquals( 10, has_filter( 'style_loader_tag', [ 'AMP_Theme_Support', 'filter_admin_bar_style_loader_tag' ] ) );
-		$this->assertEquals( 10, has_filter( 'script_loader_tag', [ 'AMP_Theme_Support', 'filter_admin_bar_script_loader_tag' ] ) );
+		AMP_Theme_Support::add_hooks();
 		$this->assertFalse( has_action( 'wp_head', $callback ) );
 		ob_start();
 		wp_print_styles();
@@ -1310,7 +1313,10 @@ public function assert_queried_element_exists( DOMXPath $xpath, $query ) {
 	public function assert_dev_mode_is_on_queried_element( DOMXPath $xpath, $query ) {
 		$element = $xpath->query( $query )->item( 0 );
 		$this->assertInstanceOf( 'DOMElement', $element, 'Expected element for query: ' . $query );
-		$this->assertTrue( $element->hasAttribute( AMP_Rule_Spec::DEV_MODE_ATTRIBUTE ), 'Expected dev mode to be enabled on element for query: ' . $query );
+		$this->assertTrue(
+			$element->hasAttribute( AMP_Rule_Spec::DEV_MODE_ATTRIBUTE ),
+			'Expected dev mode to be enabled on element for query: ' . $query . "\nDocument: " . $element->ownerDocument->saveHTML()
+		);
 	}
 
 	/**
@@ -1326,7 +1332,7 @@ public function assert_dev_mode_is_not_on_queried_element( DOMXPath $xpath, $que
 	}
 
 	/**
-	 * Get data to test AMP_Theme_Support::filter_admin_bar_style_loader_tag().
+	 * Get data to test AMP_Theme_Support::filter_style_loader_tag_for_dev_mode().
 	 *
 	 * @return array
 	 */
@@ -1396,10 +1402,10 @@ function ( DOMXPath $xpath ) {
 	}
 
 	/**
-	 * Test filter_admin_bar_style_loader_tag.
+	 * Test filter_style_loader_tag_for_dev_mode.
 	 *
 	 * @dataProvider get_data_to_test_filtering_admin_bar_style_loader_tag_data
-	 * @covers \AMP_Theme_Support::filter_admin_bar_style_loader_tag()
+	 * @covers \AMP_Theme_Support::filter_style_loader_tag_for_dev_mode()
 	 * @covers \AMP_Theme_Support::is_exclusively_dependent()
 	 *
 	 * @param callable $setup_callback  Setup callback.
@@ -1409,7 +1415,7 @@ public function test_filter_admin_bar_style_loader_tag( $setup_callback, $assert
 		add_theme_support( 'amp' );
 		$this->go_to( '/' );
 		add_filter( 'amp_dev_mode_enabled', '__return_true' );
-		add_filter( 'style_loader_tag', [ 'AMP_Theme_Support', 'filter_admin_bar_style_loader_tag' ], 10, 2 );
+		add_filter( 'style_loader_tag', [ 'AMP_Theme_Support', 'filter_style_loader_tag_for_dev_mode' ], 10, 2 );
 		$setup_callback();
 		ob_start();
 		echo '<html><head>';
@@ -1424,7 +1430,7 @@ public function test_filter_admin_bar_style_loader_tag( $setup_callback, $assert
 	}
 
 	/**
-	 * Get data to test AMP_Theme_Support::filter_admin_bar_script_loader_tag().
+	 * Get data to test AMP_Theme_Support::filter_script_loader_tag_for_dev_mode().
 	 *
 	 * @return array
 	 */
@@ -1434,10 +1440,14 @@ public function get_data_to_test_filtering_admin_bar_script_loader_tag_data() {
 				static function () {
 					wp_enqueue_script( 'admin-bar' );
 					wp_enqueue_script( 'example-admin-bar', 'https://example.com/example-admin-bar.js', [ 'admin-bar' ], '0.1', false );
+					wp_add_inline_script( 'example-admin-bar', '/* inline-example-admin-bar */' );
+					wp_localize_script( 'example-admin-bar', 'exampleAdminBar', [ 'hello' => 'world' ] );
 				},
 				function ( DOMXPath $xpath ) {
 					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( @src, "/example-admin-bar" ) ]' );
 					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( @src, "/admin-bar" ) ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( text(), "inline-example-admin-bar" ) ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( text(), "exampleAdminBar" ) ]' );
 					if ( wp_script_is( 'hoverintent-js', 'registered' ) ) {
 						$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( @src, "/hoverintent-js" ) ]' );
 					}
@@ -1496,10 +1506,10 @@ function ( DOMXPath $xpath ) {
 	}
 
 	/**
-	 * Test filter_admin_bar_script_loader_tag.
+	 * Test filter_script_loader_tag_for_dev_mode.
 	 *
 	 * @dataProvider get_data_to_test_filtering_admin_bar_script_loader_tag_data
-	 * @covers \AMP_Theme_Support::filter_admin_bar_script_loader_tag()
+	 * @covers \AMP_Theme_Support::filter_script_loader_tag_for_dev_mode()
 	 * @covers \AMP_Theme_Support::is_exclusively_dependent()
 	 *
 	 * @param callable $setup_callback  Setup callback.
@@ -1509,7 +1519,7 @@ public function test_filter_admin_bar_script_loader_tag( $setup_callback, $asser
 		add_theme_support( 'amp' );
 		$this->go_to( '/' );
 		add_filter( 'amp_dev_mode_enabled', '__return_true' );
-		add_filter( 'script_loader_tag', [ 'AMP_Theme_Support', 'filter_admin_bar_script_loader_tag' ], 10, 2 );
+		add_filter( 'script_loader_tag', [ 'AMP_Theme_Support', 'filter_script_loader_tag_for_dev_mode' ], 10, 2 );
 		$setup_callback();
 		ob_start();
 		echo '<html><head>';
@@ -1519,8 +1529,14 @@ public function test_filter_admin_bar_script_loader_tag( $setup_callback, $asser
 		echo '</body></html>';
 		$output = ob_get_clean();
 
-		$dom = new DOMDocument();
-		@$dom->loadHTML( $output ); // phpcs:ignore WordPress.PHP.NoSilencedErrors.Discouraged
+		$dom = new Document();
+		$dom->loadHTML( $output );
+
+		AMP_Content_Sanitizer::sanitize_document(
+			$dom,
+			wp_array_slice_assoc( amp_get_content_sanitizers(), [ 'AMP_Dev_Mode_Sanitizer' ] ),
+			[ 'use_document_element' => true ]
+		);
 
 		$assert_callback( new DOMXPath( $dom ) );
 	}
@@ -1529,15 +1545,18 @@ public function test_filter_admin_bar_script_loader_tag( $setup_callback, $asser
 	 * Test init_admin_bar to ensure dashicons are not added to dev mode when directly enqueued.
 	 *
 	 * @covers \AMP_Theme_Support::init_admin_bar()
-	 * @covers \AMP_Theme_Support::filter_admin_bar_style_loader_tag()
+	 * @covers \AMP_Theme_Support::filter_style_loader_tag_for_dev_mode()
 	 */
 	public function test_init_admin_bar_for_directly_enqueued_dashicons() {
+		add_theme_support( 'amp' );
+		$this->go_to( '/' );
 		require_once ABSPATH . WPINC . '/class-wp-admin-bar.php';
 
 		global $wp_admin_bar;
 		$wp_admin_bar = new WP_Admin_Bar();
 		$wp_admin_bar->initialize();
 		AMP_Theme_Support::init_admin_bar();
+		AMP_Theme_Support::add_hooks();
 
 		// Enqueued directly.
 		wp_enqueue_style( 'dashicons' );
@@ -1555,15 +1574,18 @@ public function test_init_admin_bar_for_directly_enqueued_dashicons() {
 	 * Test init_admin_bar to ensure dashicons are not added to dev mode when indirectly enqueued.
 	 *
 	 * @covers \AMP_Theme_Support::init_admin_bar()
-	 * @covers \AMP_Theme_Support::filter_admin_bar_style_loader_tag()
+	 * @covers \AMP_Theme_Support::filter_style_loader_tag_for_dev_mode()
 	 */
 	public function test_init_admin_bar_for_indirectly_enqueued_dashicons() {
+		add_theme_support( 'amp' );
+		$this->go_to( '/' );
 		require_once ABSPATH . WPINC . '/class-wp-admin-bar.php';
 
 		global $wp_admin_bar;
 		$wp_admin_bar = new WP_Admin_Bar();
 		$wp_admin_bar->initialize();
 		AMP_Theme_Support::init_admin_bar();
+		AMP_Theme_Support::add_hooks();
 
 		// Enqueued indirectly.
 		wp_enqueue_style( 'my-font-pack', 'https://example.com/fonts', [ 'dashicons' ], '0.1' );

From 399f9b53aa0b4f69040dee3b79482c38dc66e1ed Mon Sep 17 00:00:00 2001
From: Weston Ruter <westonruter@google.com>
Date: Fri, 20 Dec 2019 13:45:18 -0800
Subject: [PATCH 2/3] Obtain sanitizer args after template rendering so done
 script/style handles are available

---
 includes/amp-helper-functions.php          |  4 +-
 includes/class-amp-theme-support.php       | 39 +++++++-----
 tests/php/test-class-amp-theme-support.php | 70 +++++++++++++++++++---
 3 files changed, 88 insertions(+), 25 deletions(-)

diff --git a/includes/amp-helper-functions.php b/includes/amp-helper-functions.php
index e7b3d9cc27a..f193247beb3 100644
--- a/includes/amp-helper-functions.php
+++ b/includes/amp-helper-functions.php
@@ -965,7 +965,7 @@ function amp_get_content_sanitizers( $post = null ) {
 		// Ensure script localization data gets flagged for dev mode. This only applies to wp_localize_script() as
 		// inline scripts added via wp_add_inline_script() get filtered by script_loader_tag and thus will have the
 		// data-ampdevmode attribute added via AMP_Theme_Support::filter_script_loader_tag_for_dev_mode().
-		foreach ( wp_scripts()->done as $script_handle ) { // @todo Fail! This is currently being called _before_ the template is rendered, so it will always be empty.
+		foreach ( wp_scripts()->done as $script_handle ) {
 			if ( ! AMP_Theme_Support::dependency_needs_dev_mode( wp_scripts(), $script_handle ) ) {
 				continue;
 			}
@@ -976,7 +976,7 @@ function amp_get_content_sanitizers( $post = null ) {
 		}
 
 		// Ensure all inline styles added via wp_add_inline_style() get the data-ampdevmode attribute.
-		foreach ( wp_styles()->done as $style_handle ) { // @todo Fail! This is currently being called _before_ the template is rendered, so it will always be empty.
+		foreach ( wp_styles()->done as $style_handle ) {
 			if ( AMP_Theme_Support::dependency_needs_dev_mode( wp_styles(), $style_handle ) ) {
 				$dev_mode_xpaths[] = sprintf( '//style[ @id = "%s" ]', "$style_handle-inline-css" );
 			}
diff --git a/includes/class-amp-theme-support.php b/includes/class-amp-theme-support.php
index a2bd4017e7d..ac8dcf1c733 100644
--- a/includes/class-amp-theme-support.php
+++ b/includes/class-amp-theme-support.php
@@ -99,13 +99,6 @@ class AMP_Theme_Support {
 
 	const PAIRED_BROWSING_QUERY_VAR = 'amp-paired-browsing';
 
-	/**
-	 * Sanitizer classes.
-	 *
-	 * @var array
-	 */
-	protected static $sanitizer_classes = [];
-
 	/**
 	 * Embed handlers.
 	 *
@@ -442,20 +435,32 @@ static function() {
 		}
 
 		self::add_hooks();
-		self::$sanitizer_classes = amp_get_content_sanitizers();
-		if ( ! $is_reader_mode ) {
-			self::$sanitizer_classes = AMP_Validation_Manager::filter_sanitizer_args( self::$sanitizer_classes );
-		}
 		self::$embed_handlers = self::register_content_embed_handlers();
-		self::$sanitizer_classes['AMP_Embed_Sanitizer']['embed_handlers'] = self::$embed_handlers;
 
-		foreach ( self::$sanitizer_classes as $sanitizer_class => $args ) {
+		// @todo It is not ideal that get_sanitizer_classes() is called here before the template is rendered and after the template is rendered.
+		foreach ( self::get_sanitizer_classes() as $sanitizer_class => $args ) {
 			if ( method_exists( $sanitizer_class, 'add_buffering_hooks' ) ) {
 				call_user_func( [ $sanitizer_class, 'add_buffering_hooks' ], $args );
 			}
 		}
 	}
 
+	/**
+	 * Get sanitizer classes.
+	 *
+	 * @see amp_get_content_sanitizers()
+	 *
+	 * @return array Mapping of sanitizer class name to constructor args array.
+	 */
+	protected static function get_sanitizer_classes() {
+		$sanitizer_classes = amp_get_content_sanitizers();
+		if ( self::READER_MODE_SLUG !== self::get_support_mode() ) {
+			$sanitizer_classes = AMP_Validation_Manager::filter_sanitizer_args( $sanitizer_classes );
+		}
+		$sanitizer_classes['AMP_Embed_Sanitizer']['embed_handlers'] = self::$embed_handlers;
+		return $sanitizer_classes;
+	}
+
 	/**
 	 * Ensure that the current AMP location is correct.
 	 *
@@ -1985,7 +1990,7 @@ public static function filter_customize_partial_render( $partial ) {
 				'allow_dirty_styles'   => true,
 				'allow_dirty_scripts'  => false,
 			];
-			AMP_Content_Sanitizer::sanitize_document( $dom, self::$sanitizer_classes, $args ); // @todo Include script assets in response?
+			AMP_Content_Sanitizer::sanitize_document( $dom, self::get_sanitizer_classes(), $args ); // @todo Include script assets in response?
 			$partial = AMP_DOM_Utils::get_content_from_dom( $dom );
 		}
 		return $partial;
@@ -2095,6 +2100,8 @@ public static function prepare_response( $response, $args = [] ) {
 		$current_url = amp_get_current_url();
 		$non_amp_url = amp_remove_endpoint( $current_url );
 
+		$sanitizers = self::get_sanitizer_classes();
+
 		/*
 		 * Set response cache hash, the data values dictates whether a new hash key should be generated or not.
 		 * This is also used as the ETag.
@@ -2104,7 +2111,7 @@ public static function prepare_response( $response, $args = [] ) {
 				[
 					$args,
 					$response,
-					self::$sanitizer_classes,
+					$sanitizers,
 					self::$embed_handlers,
 					AMP__VERSION,
 				]
@@ -2268,7 +2275,7 @@ public static function prepare_response( $response, $args = [] ) {
 			$dom->documentElement->setAttribute( 'amp', '' );
 		}
 
-		$assets = AMP_Content_Sanitizer::sanitize_document( $dom, self::$sanitizer_classes, $args );
+		$assets = AMP_Content_Sanitizer::sanitize_document( $dom, $sanitizers, $args );
 
 		// Respond early with results if performing a validate request.
 		if ( AMP_Validation_Manager::$is_validate_request ) {
diff --git a/tests/php/test-class-amp-theme-support.php b/tests/php/test-class-amp-theme-support.php
index 6550aad82cd..88dc449c8d9 100644
--- a/tests/php/test-class-amp-theme-support.php
+++ b/tests/php/test-class-amp-theme-support.php
@@ -1312,10 +1312,14 @@ public function assert_queried_element_exists( DOMXPath $xpath, $query ) {
 	 */
 	public function assert_dev_mode_is_on_queried_element( DOMXPath $xpath, $query ) {
 		$element = $xpath->query( $query )->item( 0 );
-		$this->assertInstanceOf( 'DOMElement', $element, 'Expected element for query: ' . $query );
+		$this->assertInstanceOf(
+			'DOMElement',
+			$element,
+			'Expected element for query: ' . $query . "\nDocument: " . $xpath->document->saveHTML()
+		);
 		$this->assertTrue(
 			$element->hasAttribute( AMP_Rule_Spec::DEV_MODE_ATTRIBUTE ),
-			'Expected dev mode to be enabled on element for query: ' . $query . "\nDocument: " . $element->ownerDocument->saveHTML()
+			'Expected dev mode to be enabled on element for query: ' . $query . "\nDocument: " . $xpath->document->saveHTML()
 		);
 	}
 
@@ -1327,8 +1331,15 @@ public function assert_dev_mode_is_on_queried_element( DOMXPath $xpath, $query )
 	 */
 	public function assert_dev_mode_is_not_on_queried_element( DOMXPath $xpath, $query ) {
 		$element = $xpath->query( $query )->item( 0 );
-		$this->assertInstanceOf( 'DOMElement', $element, 'Expected element for query: ' . $query );
-		$this->assertFalse( $element->hasAttribute( AMP_Rule_Spec::DEV_MODE_ATTRIBUTE ), 'Expected dev mode to not be enabled on element for query: ' . $query );
+		$this->assertInstanceOf(
+			'DOMElement',
+			$element,
+			'Expected element for query: ' . $query . "\nDocument: " . $xpath->document->saveHTML()
+		);
+		$this->assertFalse(
+			$element->hasAttribute( AMP_Rule_Spec::DEV_MODE_ATTRIBUTE ),
+			'Expected dev mode to not be enabled on element for query: ' . $query . "\nDocument: " . $xpath->document->saveHTML()
+		);
 	}
 
 	/**
@@ -1341,10 +1352,12 @@ public function get_data_to_test_filtering_admin_bar_style_loader_tag_data() {
 			'admin_bar_exclusively_dependent'             => [
 				static function () {
 					wp_enqueue_style( 'example-admin-bar', 'https://example.com/example-admin-bar.css', [ 'admin-bar' ], '0.1' );
+					wp_add_inline_style( 'example-admin-bar', '#wpadminbar:after { content: "Hey admin!" }' );
 				},
 				function ( DOMXPath $xpath ) {
 					$this->assert_dev_mode_is_on_queried_element( $xpath, '//link[ @id = "example-admin-bar-css" ]' );
 					$this->assert_dev_mode_is_on_queried_element( $xpath, '//link[ @id = "dashicons-css" ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//style[ @id = "example-admin-bar-inline-css" ]' );
 				},
 			],
 
@@ -1398,6 +1411,22 @@ function ( DOMXPath $xpath ) {
 					$this->assert_dev_mode_is_on_queried_element( $xpath, '//link[ @id = "dashicons-css" ]' );
 				},
 			],
+
+			'styles_have_dev_mode_when_flagged'           => [
+				static function () {
+					wp_enqueue_style( 'custom-style', 'https://example.com/custom-style.css', [], '1.0', false );
+					wp_add_inline_style( 'custom-style', '/* inline-custom-style */' );
+					wp_style_add_data( 'custom-style', 'ampdevmode', true );
+
+					wp_enqueue_style( 'excluded-style', 'https://example.com/excluded-style.js', [], '1.0', false );
+				},
+				function ( DOMXPath $xpath ) {
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//link[ contains( @href, "/custom-style" ) ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//style[ @id = "custom-style-inline-css" ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//style[ contains( text(), "inline-custom-style" ) ]' );
+					$this->assert_dev_mode_is_not_on_queried_element( $xpath, '//link[ contains( @href, "/excluded-style" ) ]' );
+				},
+			],
 		];
 	}
 
@@ -1423,9 +1452,15 @@ public function test_filter_admin_bar_style_loader_tag( $setup_callback, $assert
 		echo '</head><body></body></html>';
 		$output = ob_get_clean();
 
-		$dom = new DOMDocument();
+		$dom = new Document();
 		$dom->loadHTML( $output );
 
+		AMP_Content_Sanitizer::sanitize_document(
+			$dom,
+			wp_array_slice_assoc( amp_get_content_sanitizers(), [ 'AMP_Dev_Mode_Sanitizer' ] ),
+			[ 'use_document_element' => true ]
+		);
+
 		$assert_callback( new DOMXPath( $dom ) );
 	}
 
@@ -1440,13 +1475,15 @@ public function get_data_to_test_filtering_admin_bar_script_loader_tag_data() {
 				static function () {
 					wp_enqueue_script( 'admin-bar' );
 					wp_enqueue_script( 'example-admin-bar', 'https://example.com/example-admin-bar.js', [ 'admin-bar' ], '0.1', false );
-					wp_add_inline_script( 'example-admin-bar', '/* inline-example-admin-bar */' );
+					wp_add_inline_script( 'example-admin-bar', '/* inline-example-admin-bar-before */', 'before' );
+					wp_add_inline_script( 'example-admin-bar', '/* inline-example-admin-bar-after */', 'after' );
 					wp_localize_script( 'example-admin-bar', 'exampleAdminBar', [ 'hello' => 'world' ] );
 				},
 				function ( DOMXPath $xpath ) {
 					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( @src, "/example-admin-bar" ) ]' );
 					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( @src, "/admin-bar" ) ]' );
-					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( text(), "inline-example-admin-bar" ) ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( text(), "inline-example-admin-bar-before" ) ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( text(), "inline-example-admin-bar-after" ) ]' );
 					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( text(), "exampleAdminBar" ) ]' );
 					if ( wp_script_is( 'hoverintent-js', 'registered' ) ) {
 						$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( @src, "/hoverintent-js" ) ]' );
@@ -1502,6 +1539,25 @@ function ( DOMXPath $xpath ) {
 					$this->assert_dev_mode_is_not_on_queried_element( $xpath, '//script[ contains( @src, "/hoverintent-js" ) ]' );
 				},
 			],
+
+			'scripts_have_dev_mode_when_flagged'     => [
+				static function () {
+					wp_enqueue_script( 'custom-script', 'https://example.com/custom-script.js', [], '1.0', false );
+					wp_add_inline_script( 'custom-script', '/* inline-custom-script-before */', 'before' );
+					wp_add_inline_script( 'custom-script', '/* inline-custom-script-after */', 'after' );
+					wp_localize_script( 'custom-script', 'customScript', [ 'hello' => 'world' ] );
+					wp_script_add_data( 'custom-script', 'ampdevmode', true );
+
+					wp_enqueue_script( 'excluded-script', 'https://example.com/excluded-script.js', [], '1.0', false );
+				},
+				function ( DOMXPath $xpath ) {
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( @src, "/custom-script" ) ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( text(), "inline-custom-script-before" ) ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( text(), "inline-custom-script-after" ) ]' );
+					$this->assert_dev_mode_is_on_queried_element( $xpath, '//script[ contains( text(), "customScript" ) ]' );
+					$this->assert_dev_mode_is_not_on_queried_element( $xpath, '//script[ contains( @src, "/excluded-script" ) ]' );
+				},
+			],
 		];
 	}
 

From 9db65fc1c72f3c51d1ebefc57b8899246f1d265d Mon Sep 17 00:00:00 2001
From: Weston Ruter <westonruter@google.com>
Date: Fri, 20 Dec 2019 13:52:30 -0800
Subject: [PATCH 3/3] Add assertions to
 test_amp_get_content_sanitizers_with_dev_mode

---
 tests/php/test-amp-helper-functions.php | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tests/php/test-amp-helper-functions.php b/tests/php/test-amp-helper-functions.php
index 3d8f6167438..0863c1a70aa 100644
--- a/tests/php/test-amp-helper-functions.php
+++ b/tests/php/test-amp-helper-functions.php
@@ -784,7 +784,10 @@ function ( $xpaths ) use ( $element_xpaths ) {
 		// Check that AMP_Dev_Mode_Sanitizer is registered once in dev mode, and now also with admin bar showing.
 		add_filter( 'amp_dev_mode_enabled', '__return_true' );
 		add_filter( 'show_admin_bar', '__return_true' );
-		$sanitizers = amp_get_content_sanitizers();
+		wp_localize_script( 'admin-bar', 'myAdminBarStrings', [ 'hello' => 'world' ] );
+		wp_scripts()->done[] = 'admin-bar';
+		wp_styles()->done[]  = 'admin-bar';
+		$sanitizers          = amp_get_content_sanitizers();
 		$this->assertTrue( is_admin_bar_showing() );
 		$this->assertTrue( amp_is_dev_mode() );
 		$this->assertArrayHasKey( 'AMP_Dev_Mode_Sanitizer', $sanitizers );
@@ -794,6 +797,8 @@ function ( $xpaths ) use ( $element_xpaths ) {
 				[
 					'//*[ @id = "wpadminbar" ]',
 					'//*[ @id = "wpadminbar" ]//*',
+					'//script[ not( @src ) ][ contains( text(), "var myAdminBarStrings =" ) ]',
+					'//style[ @id = "admin-bar-inline-css" ]',
 				]
 			),
 			$sanitizers['AMP_Dev_Mode_Sanitizer']['element_xpaths']