Add support for allowing comments or doctypes

wooorm · wooorm · commit dfe4feacd411 · 2018-06-11T11:10:00.000+02:00
Closes GH-11.
diff --git a/lib/index.js b/lib/index.js
@@ -12,6 +12,8 @@ var own = {}.hasOwnProperty
 /* Schema. */
 var NODES = {
   root: {children: all},
+  doctype: handleDoctype,
+  comment: handleComment,
   element: {
     tagName: handleTagName,
     properties: handleProperties,
@@ -57,25 +59,36 @@ function one(schema, node, stack) {
   var type = node && node.type
   var replacement = {type: node.type}
   var replace = true
+  var definition
   var allowed
   var result
   var key
 
   if (!own.call(NODES, type)) {
     replace = false
   } else {
-    allowed = xtend(NODES[type], NODES['*'])
+    definition = NODES[type]
 
-    for (key in allowed) {
-      result = allowed[key](schema, node[key], node, stack)
+    if (typeof definition === 'function') {
+      definition = definition(schema, node)
+    }
+
+    if (!definition) {
+      replace = false
+    } else {
+      allowed = xtend(definition, NODES['*'])
 
-      if (result === false) {
-        replace = false
+      for (key in allowed) {
+        result = allowed[key](schema, node[key], node, stack)
 
-        /* Set the non-safe value. */
-        replacement[key] = node[key]
-      } else if (result !== null && result !== undefined) {
-        replacement[key] = result
+        if (result === false) {
+          replace = false
+
+          /* Set the non-safe value. */
+          replacement[key] = node[key]
+        } else if (result !== null && result !== undefined) {
+          replacement[key] = result
+        }
       }
     }
   }
@@ -253,6 +266,11 @@ function handleProtocol(schema, value, prop) {
   return false
 }
 
+/* Always return a valid HTML5 doctype. */
+function handleDoctypeName() {
+  return 'html'
+}
+
 /* Sanitize `tagName`. */
 function handleTagName(schema, tagName, node, stack) {
   var name = typeof tagName === 'string' ? tagName : null
@@ -286,6 +304,14 @@ function handleTagName(schema, tagName, node, stack) {
   return name
 }
 
+function handleDoctype(schema) {
+  return schema.allowDoctypes ? {name: handleDoctypeName} : null
+}
+
+function handleComment(schema) {
+  return schema.allowComments ? {value: handleValue} : null
+}
+
 /* Sanitize `value`. */
 function handleValue(schema, value) {
   return typeof value === 'string' ? value : ''
diff --git a/readme.md b/readme.md
@@ -93,7 +93,7 @@ var gh = require('hast-util-sanitize/lib/github');
 
 var schema = merge(gh, {attributes: {'*': ['className']}});
 
-var tree = sanitize(h('div', {className: ['foo']}));
+var tree = sanitize(h('div', {className: ['foo']}), schema);
 // `tree` still has `className`.
 ```
 
@@ -213,6 +213,22 @@ should however be entirely stripped from the tree.
 ]
 ```
 
+###### `allowComments`
+
+Whether to allow comment nodes (`boolean`, default: `false`).
+
+```js
+"allowComments": true
+```
+
+###### `allowDoctypes`
+
+Whether to allow doctype nodes (`boolean`, default: `false`).
+
+```js
+"allowDoctypes": true
+```
+
 ## Contribute
 
 See [`contributing.md` in `syntax-tree/hast`][contributing] for ways to get
diff --git a/test.js b/test.js
@@ -53,10 +53,38 @@ test('sanitize()', function(t) {
       'should ignore `characterData`s'
     )
 
+    st.end()
+  })
+
+  t.test('`comment`', function(st) {
     st.equal(
       html(sanitize(u('comment', 'alpha'))),
       '',
-      'should ignore `comment`s'
+      'should ignore `comment`s by default'
+    )
+
+    st.equal(
+      html(sanitize(u('comment', 'alpha'), {allowComments: true})),
+      '<!--alpha-->',
+      'should allow `comment`s with `allowComments: true`'
+    )
+
+    st.end()
+  })
+
+  t.test('`doctype`', function(st) {
+    st.equal(
+      html(sanitize(u('doctype', {name: 'html'}, 'alpha'))),
+      '',
+      'should ignore `doctype`s by default'
+    )
+
+    st.equal(
+      html(
+        sanitize(u('doctype', {name: 'html'}, 'alpha'), {allowDoctypes: true})
+      ),
+      '<!DOCTYPE html>',
+      'should allow `doctype`s with `allowDoctypes: true`'
     )
 
     st.end()
