From 28aae733d08a62a06003e1a390452c026c61a955 Mon Sep 17 00:00:00 2001
From: William Chou <willchou@google.com>
Date: Fri, 19 Oct 2018 12:11:09 -0400
Subject: [PATCH] amp-script: Use AMP purifier config and hooks (#18789)

* Hook up AMP purify config/hooks to amp-script.

* Upgrade to worker-dom 0.1.4.

* Fix types.

* Fix hello-world example.

* Fix dep-check.

* Use '.patched' suffix to fix SP issue. Such a long name.
---
 build-system/dep-check-config.js         |  3 +-
 build-system/tasks/compile.js            |  3 +-
 build-system/tasks/update-packages.js    | 13 +++---
 examples/amp-script/hello-world.amp.html | 12 ++++--
 examples/amp-script/hello-world.html     |  8 ++++
 examples/amp-script/hello-world.js       | 43 +++++++++++++-----
 extensions/amp-script/0.1/amp-script.js  | 27 ++++++++++--
 package.json                             |  2 +-
 src/purifier.js                          | 55 +++++++++++++++++-------
 yarn.lock                                |  8 ++--
 10 files changed, 127 insertions(+), 47 deletions(-)
 create mode 100644 examples/amp-script/hello-world.html

diff --git a/build-system/dep-check-config.js b/build-system/dep-check-config.js
index f9c9a8d67b8b..cb5081bf56ba 100644
--- a/build-system/dep-check-config.js
+++ b/build-system/dep-check-config.js
@@ -59,8 +59,9 @@ exports.rules = [
     mustNotDependOn: 'src/purifier.js',
     whitelist: [
       'src/sanitizer.js->src/purifier.js',
-      'extensions/amp-mustache/0.2/amp-mustache.js->src/purifier.js',
       'extensions/amp-bind/0.1/bind-impl.js->src/purifier.js',
+      'extensions/amp-mustache/0.2/amp-mustache.js->src/purifier.js',
+      'extensions/amp-script/0.1/amp-script.js->src/purifier.js',
     ],
   },
   {
diff --git a/build-system/tasks/compile.js b/build-system/tasks/compile.js
index b27b6d751b36..c4e9466e320d 100644
--- a/build-system/tasks/compile.js
+++ b/build-system/tasks/compile.js
@@ -256,7 +256,8 @@ function compile(entryModuleFilenames, outputDir, outputFilename, options) {
       'node_modules/set-dom/src/**/*.js',
       'node_modules/web-animations-js/web-animations.install.js',
       'node_modules/web-activities/activity-ports.js',
-      'node_modules/@ampproject/worker-dom/dist/**/*.js',
+      'node_modules/@ampproject/worker-dom/dist/' +
+          'unminified.index.safe.mjs.patched.js',
       'node_modules/document-register-element/build/' +
           'document-register-element.patched.js',
       // 'node_modules/core-js/modules/**.js',
diff --git a/build-system/tasks/update-packages.js b/build-system/tasks/update-packages.js
index dd3c85772f1e..248af63a0bb8 100644
--- a/build-system/tasks/update-packages.js
+++ b/build-system/tasks/update-packages.js
@@ -110,15 +110,14 @@ function patchRegisterElement() {
 }
 
 /**
- * Patch @ampproject/worker-dom/dist/index.safe.js with ES6 export.
+ * Closure Compiler doesn't recognize .mjs extension yet, so copy the file to
+ * have a .js file extension.
  */
 function patchWorkerDom() {
-  replaceInFile(
-      'node_modules/@ampproject/worker-dom/dist/index.safe.js',
-      'node_modules/@ampproject/worker-dom/dist/index.safe.patched.js',
-      // Replace local var with an ES6 export.
-      'var MainThread=function(R){',
-      'export const MainThread=function(R){');
+  const dir = 'node_modules/@ampproject/worker-dom/dist/';
+  fs.copyFileSync(
+      dir + 'unminified.index.safe.mjs',
+      dir + 'unminified.index.safe.mjs.patched.js');
 }
 
 /**
diff --git a/examples/amp-script/hello-world.amp.html b/examples/amp-script/hello-world.amp.html
index 40a120759c21..425ffd4ff18f 100644
--- a/examples/amp-script/hello-world.amp.html
+++ b/examples/amp-script/hello-world.amp.html
@@ -4,16 +4,20 @@
   <meta charset="utf-8">
   <link rel="canonical" href="self.html" />
   <meta name="viewport" content="width=device-width,minimum-scale=1">
+
   <!-- CSP -->
-  <!-- <meta http-equiv="Content-Security-Policy" content="default-src * data: blob:; script-src blob: http://localhost:8000/dist/amp.js http://localhost:8000/dist/v0/amp-script-0.1.max.js https://cdn.ampproject.org/v0.js https://cdn.ampproject.org/v0/ https://cdn.ampproject.org/viewer/ https://cdn.ampproject.org/rtv/; object-src 'none'; style-src 'unsafe-inline' https://cdn.ampproject.org/rtv/ https://cdn.materialdesignicons.com https://cloud.typography.com https://fast.fonts.net https://fonts.googleapis.com https://maxcdn.bootstrapcdn.com https://p.typekit.net https://use.fontawesome.com https://use.typekit.net; report-uri https://csp-collector.appspot.com/csp/amp"> -->
+  <!-- <meta http-equiv="Content-Security-Policy" content="default-src * data: blob:; script-src blob: http://localhost:8000/dist/v0.js http://localhost:8000/dist/v0/amp-script-0.1.js https://cdn.ampproject.org/v0.js https://cdn.ampproject.org/v0/ https://cdn.ampproject.org/viewer/ https://cdn.ampproject.org/rtv/; object-src 'none'; style-src 'unsafe-inline' https://cdn.ampproject.org/rtv/ https://cdn.materialdesignicons.com https://cloud.typography.com https://fast.fonts.net https://fonts.googleapis.com https://maxcdn.bootstrapcdn.com https://p.typekit.net https://use.fontawesome.com https://use.typekit.net; report-uri https://csp-collector.appspot.com/csp/amp"> -->
 
   <style amp-boilerplate>body{-webkit-animation:-amp-start 8s steps(1,end) 0s 1 normal both;-moz-animation:-amp-start 8s steps(1,end) 0s 1 normal both;-ms-animation:-amp-start 8s steps(1,end) 0s 1 normal both;animation:-amp-start 8s steps(1,end) 0s 1 normal both}@-webkit-keyframes -amp-start{from{visibility:hidden}to{visibility:visible}}@-moz-keyframes -amp-start{from{visibility:hidden}to{visibility:visible}}@-ms-keyframes -amp-start{from{visibility:hidden}to{visibility:visible}}@-o-keyframes -amp-start{from{visibility:hidden}to{visibility:visible}}@keyframes -amp-start{from{visibility:hidden}to{visibility:visible}}</style><noscript><style amp-boilerplate>body{-webkit-animation:none;-moz-animation:none;-ms-animation:none;animation:none}</style></noscript>
+  <style amp-custom>
+    button {
+      border: solid 1px black;
+    }
+  </style>
   <script async src="https://cdn.ampproject.org/v0.js"></script>
   <script async custom-template="amp-script" src="https://cdn.ampproject.org/v0/amp-script-0.1.js"></script>
 </head>
 <body>
-  <amp-script layout=container src="/examples/amp-script/hello-world.js">
-    <div class="root"><button>prompt</button></div>
-  </amp-script>
+  <amp-script layout=container src="/examples/amp-script/hello-world.js"><div class="root"><button>prompt</button><button>prompt</button><button>prompt</button><button>prompt</button></div></amp-script>
 </body>
 </html>
diff --git a/examples/amp-script/hello-world.html b/examples/amp-script/hello-world.html
new file mode 100644
index 000000000000..1247fda690bd
--- /dev/null
+++ b/examples/amp-script/hello-world.html
@@ -0,0 +1,8 @@
+<!doctype html>
+<html>
+<head>
+  <script async src="hello-world.js"></script>
+</head>
+<body>
+</body>
+</html>
diff --git a/examples/amp-script/hello-world.js b/examples/amp-script/hello-world.js
index b96f16075de1..a770bae44f8b 100644
--- a/examples/amp-script/hello-world.js
+++ b/examples/amp-script/hello-world.js
@@ -15,17 +15,40 @@
  */
 
 const root = document.createElement('div');
-const btn = document.createElement('button');
-const text = document.createTextNode('Insert Hello World!');
-
 root.className = "root";
-btn.appendChild(text);
-root.appendChild(btn);
+document.body.appendChild(root);
+
+function createButton(label, onClick) {
+  const btn = document.createElement('button');
+  const text = document.createTextNode(label);
+  btn.appendChild(text);
+  btn.addEventListener('click', onClick);
+  root.appendChild(btn);
+}
+
+createButton('Insert Hello World!', () => {
+  const el = document.createElement('h1');
+  el.textContent = 'Hello World!';
+  document.body.appendChild(el);
+});
 
-btn.addEventListener('click', () => {
-  const h1 = document.createElement('h1');
-  h1.textContent = 'Hello World!'
-  document.body.appendChild(h1);
+// <amp-img> should be allowed.
+createButton('Insert amp-img', () => {
+  const el = document.createElement('amp-img');
+  el.setAttribute('width', '300');
+  el.setAttribute('height', '200');
+  el.setAttribute('src', '/examples/img/hero@1x.jpg')
+  document.body.appendChild(el);
 });
 
-document.body.appendChild(root);
\ No newline at end of file
+// <script> should be sanitized.
+createButton('Insert <script>', () => {
+  const el = document.createElement('script');
+  document.body.appendChild(el);
+});
+
+// <img> should be sanitized.
+createButton('Insert <img>', () => {
+  const el = document.createElement('img');
+  document.body.appendChild(el);
+});
diff --git a/extensions/amp-script/0.1/amp-script.js b/extensions/amp-script/0.1/amp-script.js
index 386a48840662..9783c8db942f 100644
--- a/extensions/amp-script/0.1/amp-script.js
+++ b/extensions/amp-script/0.1/amp-script.js
@@ -15,16 +15,21 @@
  */
 
 import {Layout, isLayoutSizeDefined} from '../../../src/layout';
-import {MainThread} from '@ampproject/worker-dom/dist/index.safe.patched';
+import {addPurifyHooks, purifyConfig} from '../../../src/purifier';
 import {
   calculateExtensionScriptUrl,
 } from '../../../src/service/extension-location';
-import {dev} from '../../../src/log';
+import {dev, user} from '../../../src/log';
 import {getMode} from '../../../src/mode';
+import {
+  sanitizer,
+  upgradeElement,
+} from '@ampproject/worker-dom/dist/unminified.index.safe.mjs.patched';
 
 /** @const {string} */
 const TAG = 'amp-script';
 
+
 export class AmpScript extends AMP.BaseElement {
   /** @override */
   isLayoutSupported(layout) {
@@ -34,9 +39,25 @@ export class AmpScript extends AMP.BaseElement {
 
   /** @override */
   layoutCallback() {
+    // Configure worker-dom's sanitizer with AMP-specific config and hooks.
+    const config = purifyConfig();
+    sanitizer.configure(config, {
+      'beforeSanitize': purify => {
+        addPurifyHooks(purify, /* diffing */ false);
+      },
+      'afterSanitize': purify => {
+        purify.removeAllHooks();
+      },
+      'nodeWasRemoved': node => {
+        user().warn(TAG, 'Node was sanitized:', node);
+      },
+    });
+
     const url = this.workerThreadUrl_();
     dev().fine(TAG, 'Fetching amp-script-worker from:', url);
-    MainThread.upgradeElement(this.element, url);
+
+    upgradeElement(this.element, url);
+
     return Promise.resolve();
   }
 
diff --git a/package.json b/package.json
index 1b2e353f5391..70cb210acb4d 100644
--- a/package.json
+++ b/package.json
@@ -25,7 +25,7 @@
     "preinstall": "node build-system/check-package-manager.js"
   },
   "dependencies": {
-    "@ampproject/worker-dom": "0.1.2",
+    "@ampproject/worker-dom": "0.1.4",
     "document-register-element": "1.5.0",
     "dompurify": "1.0.8",
     "promise-pjs": "1.1.3",
diff --git a/src/purifier.js b/src/purifier.js
index bcce351aa28c..5f8c172fce88 100644
--- a/src/purifier.js
+++ b/src/purifier.js
@@ -29,8 +29,13 @@ import {urls} from './config';
 import {user} from './log';
 import purify from 'dompurify/dist/purify.es';
 
-/** @private @const {{addHook: !Function, removeAllHooks: !Function, sanitize: !Function}} */
-const DOMPurify = purify(self);
+/**
+ * @typedef {{addHook: !Function, removeAllHooks: !Function, sanitize: !Function}}
+ */
+let DomPurifyDef;
+
+/** @private @const {!DomPurifyDef} */
+const DomPurify = purify(self);
 
 /** @private @const {string} */
 const TAG = 'purifier';
@@ -219,14 +224,14 @@ const BLACKLISTED_TAG_SPECIFIC_ATTRS = dict({
 const INVALID_INLINE_STYLE_REGEX =
     /!important|position\s*:\s*fixed|position\s*:\s*sticky/i;
 
-/** @const {!Object} */
-const PURIFY_CONFIG = {
+/** @const {!JsonObject} */
+const PURIFY_CONFIG = dict({
   'USE_PROFILES': {
     'html': true,
     'svg': true,
     'svgFilters': true,
   },
-};
+});
 
 /**
  * Monotonically increasing counter used for keying nodes.
@@ -241,6 +246,19 @@ let KEY_COUNTER = 0;
  * @return {!Node}
  */
 export function purifyHtml(dirty, diffing = false) {
+  const config = purifyConfig();
+  addPurifyHooks(DomPurify, diffing);
+  const body = DomPurify.sanitize(dirty, config);
+  DomPurify.removeAllHooks();
+  return body;
+}
+
+/**
+ * Returns DOMPurify config for normal, escaped templates.
+ * Do not use for unescaped templates.
+ * @return {!JsonObject}
+ */
+export function purifyConfig() {
   const config = Object.assign({}, PURIFY_CONFIG, {
     'ADD_ATTR': WHITELISTED_ATTRS,
     'FORBID_TAGS': Object.keys(BLACKLISTED_TAGS),
@@ -249,7 +267,15 @@ export function purifyHtml(dirty, diffing = false) {
     // Avoid need for serializing to/from string by returning Node directly.
     'RETURN_DOM': true,
   });
+  return /** @type {!JsonObject} */ (config);
+}
 
+/**
+ * Adds AMP hooks to given DOMPurify object.
+ * @param {!DomPurifyDef} purifier
+ * @param {boolean} diffing
+ */
+export function addPurifyHooks(purifier, diffing) {
   // Reference to DOMPurify's `allowedTags` whitelist.
   let allowedTags;
   const allowedTagsChanges = [];
@@ -423,13 +449,10 @@ export function purifyHtml(dirty, diffing = false) {
     });
   };
 
-  DOMPurify.addHook('uponSanitizeElement', uponSanitizeElement);
-  DOMPurify.addHook('afterSanitizeElements', afterSanitizeElements);
-  DOMPurify.addHook('uponSanitizeAttribute', uponSanitizeAttribute);
-  DOMPurify.addHook('afterSanitizeAttributes', afterSanitizeAttributes);
-  const body = DOMPurify.sanitize(dirty, config);
-  DOMPurify.removeAllHooks();
-  return body;
+  purifier.addHook('uponSanitizeElement', uponSanitizeElement);
+  purifier.addHook('afterSanitizeElements', afterSanitizeElements);
+  purifier.addHook('uponSanitizeAttribute', uponSanitizeAttribute);
+  purifier.addHook('afterSanitizeAttributes', afterSanitizeAttributes);
 }
 
 /**
@@ -444,7 +467,7 @@ export function purifyTagsForTripleMustache(html, doc = self.document) {
   // Reference to DOMPurify's `allowedTags` whitelist.
   let allowedTags;
 
-  DOMPurify.addHook('uponSanitizeElement', (node, data) => {
+  DomPurify.addHook('uponSanitizeElement', (node, data) => {
     const {tagName} = data;
     allowedTags = data.allowedTags;
     if (tagName === 'template') {
@@ -454,7 +477,7 @@ export function purifyTagsForTripleMustache(html, doc = self.document) {
       }
     }
   });
-  DOMPurify.addHook('afterSanitizeElements', unusedNode => {
+  DomPurify.addHook('afterSanitizeElements', unusedNode => {
     // DOMPurify doesn't have an required-attribute tag whitelist API and
     // `allowedTags` has a per-invocation scope, so we need to remove
     // required-attribute tags after sanitizing each element.
@@ -464,12 +487,12 @@ export function purifyTagsForTripleMustache(html, doc = self.document) {
   // reparented to the head. So to support nested templates, we need
   // RETURN_DOM_FRAGMENT to keep the <template> and FORCE_BODY to prevent
   // reparenting. See https://github.com/cure53/DOMPurify/issues/285#issuecomment-397810671
-  const fragment = DOMPurify.sanitize(html, {
+  const fragment = DomPurify.sanitize(html, {
     'ALLOWED_TAGS': TRIPLE_MUSTACHE_WHITELISTED_TAGS,
     'FORCE_BODY': true,
     'RETURN_DOM_FRAGMENT': true,
   });
-  DOMPurify.removeAllHooks();
+  DomPurify.removeAllHooks();
   // Serialize DocumentFragment to HTML. XMLSerializer would also work, but adds
   // namespaces for all elements and attributes.
   const div = doc.createElement('div');
diff --git a/yarn.lock b/yarn.lock
index 80587b21bc44..93f6b9be213b 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2,10 +2,10 @@
 # yarn lockfile v1
 
 
-"@ampproject/worker-dom@0.1.2":
-  version "0.1.2"
-  resolved "https://registry.yarnpkg.com/@ampproject/worker-dom/-/worker-dom-0.1.2.tgz#aa365fef9d99f1af21e8d3d144195d2946145562"
-  integrity sha512-POCRe6350Gsc3Vfg5JRCCvio+TcUJSTYg1tJHUANXAKQpdisVeujk2mk6/XhCHVy3lVrQjCnAYomepJIoH7NEw==
+"@ampproject/worker-dom@0.1.4":
+  version "0.1.4"
+  resolved "https://registry.yarnpkg.com/@ampproject/worker-dom/-/worker-dom-0.1.4.tgz#db62bec216c99fa14e8f87138b1e7925a20a582b"
+  integrity sha512-0dgREzd48ae28j8J0aF3+pFGiLqLO0lI+84Ip000ZAjMUM2gFvd4KH1AzsqPkKx4/4IdJBtt3iQbTgsWjPdQ8Q==
   dependencies:
     dompurify "1.0.8"