From d3e4ba6dd53eee938416026b90ad30f4bc3bf369 Mon Sep 17 00:00:00 2001
From: Patrick Simonian <Patrick.Simonian@gov.bc.ca>
Date: Mon, 16 Dec 2019 10:20:20 -0800
Subject: [PATCH] add feature to provide additional html escape sequences

address test failures
---
 packages/gatsby-remark-prismjs/README.md         |  4 ++++
 .../src/__tests__/escape-html.js                 | 16 ++++++++++++++++
 .../src/__tests__/highlight-code.js              | 12 +++++++-----
 .../gatsby-remark-prismjs/src/escape-html.js     | 15 +++++++++++++++
 .../gatsby-remark-prismjs/src/highlight-code.js  |  4 +++-
 packages/gatsby-remark-prismjs/src/index.js      |  6 ++++--
 6 files changed, 49 insertions(+), 8 deletions(-)
 create mode 100644 packages/gatsby-remark-prismjs/src/__tests__/escape-html.js
 create mode 100644 packages/gatsby-remark-prismjs/src/escape-html.js

diff --git a/packages/gatsby-remark-prismjs/README.md b/packages/gatsby-remark-prismjs/README.md
index 600191a1d1920..4a3e5f0851e16 100644
--- a/packages/gatsby-remark-prismjs/README.md
+++ b/packages/gatsby-remark-prismjs/README.md
@@ -75,6 +75,10 @@ plugins: [
               host: "localhost",
               global: false,
             },
+            // By default the HTML entities <>&'" are escaped.
+            // Add additional HTML escapes by providing a mapping
+            // of HTML entities and their escape value IE: { '}': '&#123;' }
+            escapeEntities: {},
           },
         },
       ],
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/escape-html.js b/packages/gatsby-remark-prismjs/src/__tests__/escape-html.js
new file mode 100644
index 0000000000000..d88dfb41a76d9
--- /dev/null
+++ b/packages/gatsby-remark-prismjs/src/__tests__/escape-html.js
@@ -0,0 +1,16 @@
+const escapeHtml = require(`../escape-html`)
+
+describe(`escaping html entities`, () => {
+  it(`returns code unchanged when no escapes are provided`, () => {
+    const code = `hello world&`
+    expect(escapeHtml(code)).toBe(code)
+  })
+
+  it(`escapes html`, () => {
+    const code = `&hello world&`
+    const mapping = {
+      "&": `&amp;`,
+    }
+    expect(escapeHtml(code, mapping)).toBe(`&amp;hello world&amp;`)
+  })
+})
diff --git a/packages/gatsby-remark-prismjs/src/__tests__/highlight-code.js b/packages/gatsby-remark-prismjs/src/__tests__/highlight-code.js
index 4928b422ea4b0..12758649854fc 100644
--- a/packages/gatsby-remark-prismjs/src/__tests__/highlight-code.js
+++ b/packages/gatsby-remark-prismjs/src/__tests__/highlight-code.js
@@ -13,7 +13,7 @@ int sum(a, b) {
 }
 `
     expect(
-      highlightCode(language, code, lineNumbersHighlight)
+      highlightCode(language, code, {}, lineNumbersHighlight)
     ).toMatchSnapshot()
   })
 
@@ -48,7 +48,7 @@ class Counter extends React.Component {
 
 export default Counter
 `
-    const processed = highlightCode(language, code, lineNumbersHighlight)
+    const processed = highlightCode(language, code, {}, lineNumbersHighlight)
 
     expect(processed).toMatchSnapshot()
     // expect spans to not contain \n as it would break line highlighting
@@ -62,7 +62,7 @@ export default Counter
       const highlightCode = require(`../highlight-code`)
       const language = `text`
       const code = `<button />`
-      expect(highlightCode(language, code, [], true)).toMatch(
+      expect(highlightCode(language, code, {}, [], true)).toMatch(
         `&lt;button /&gt;`
       )
       expect(console.warn).toHaveBeenCalledWith(
@@ -141,7 +141,9 @@ export default Counter
       const language = `javascript`
       const linesToHighlight = [1]
       const code = `const a = 1\nconst b = 2`
-      expect(highlightCode(language, code, linesToHighlight)).not.toMatch(/\n$/)
+      expect(highlightCode(language, code, {}, linesToHighlight)).not.toMatch(
+        /\n$/
+      )
     })
 
     it(`a trailing newline is preserved`, () => {
@@ -149,7 +151,7 @@ export default Counter
       const language = `javascript`
       const linesToHighlight = [1]
       const code = `const a = 1\nconst b = 2\n`
-      expect(highlightCode(language, code, linesToHighlight)).toMatch(
+      expect(highlightCode(language, code, {}, linesToHighlight)).toMatch(
         /[^\n]\n$/
       )
     })
diff --git a/packages/gatsby-remark-prismjs/src/escape-html.js b/packages/gatsby-remark-prismjs/src/escape-html.js
new file mode 100644
index 0000000000000..68173b26b4461
--- /dev/null
+++ b/packages/gatsby-remark-prismjs/src/escape-html.js
@@ -0,0 +1,15 @@
+const isEmpty = require(`lodash/isEmpty`)
+
+module.exports = (code, additionalEscapeCharacters = {}) => {
+  if (isEmpty(additionalEscapeCharacters)) return code
+  const escapedChars = char => additionalEscapeCharacters[char]
+  const chars = Object.keys(additionalEscapeCharacters)
+
+  const charsRe = new RegExp(`[${chars.join()}]`, `g`)
+
+  const rehasUnescapedChars = new RegExp(charsRe.source)
+
+  return code && rehasUnescapedChars.test(code)
+    ? code.replace(charsRe, escapedChars)
+    : code
+}
diff --git a/packages/gatsby-remark-prismjs/src/highlight-code.js b/packages/gatsby-remark-prismjs/src/highlight-code.js
index e5ddaaef0d732..19b23a17c9e14 100644
--- a/packages/gatsby-remark-prismjs/src/highlight-code.js
+++ b/packages/gatsby-remark-prismjs/src/highlight-code.js
@@ -2,11 +2,13 @@ const Prism = require(`prismjs`)
 const _ = require(`lodash`)
 const loadPrismLanguage = require(`./load-prism-language`)
 const handleDirectives = require(`./directives`)
+const escapeHTML = require(`./escape-html`)
 const unsupportedLanguages = new Set()
 
 module.exports = (
   language,
   code,
+  additionalEscapeCharacters = {},
   lineNumbersHighlight = [],
   noInlineHighlight = false
 ) => {
@@ -34,7 +36,7 @@ module.exports = (
         console.warn(message, `applying generic code block`)
         unsupportedLanguages.add(lang)
       }
-      return _.escape(code)
+      return escapeHTML(_.escape(code), additionalEscapeCharacters)
     }
   }
 
diff --git a/packages/gatsby-remark-prismjs/src/index.js b/packages/gatsby-remark-prismjs/src/index.js
index 7dae4be862efd..11ff27fbae430 100644
--- a/packages/gatsby-remark-prismjs/src/index.js
+++ b/packages/gatsby-remark-prismjs/src/index.js
@@ -20,6 +20,7 @@ module.exports = (
       host: `localhost`,
       global: false,
     },
+    escapeEntities = {},
   } = {}
 ) => {
   const normalizeLanguage = lang => {
@@ -94,7 +95,7 @@ module.exports = (
     +   `<pre${numLinesStyle} class="${className}${numLinesClass}">`
     +     `<code class="${className}">`
     +       `${useCommandLine ? commandLine(node.value, outputLines, promptUser, promptHost) : ``}`
-    +       `${highlightCode(languageName, node.value, highlightLines, noInlineHighlight)}`
+    +       `${highlightCode(languageName, node.value, escapeEntities, highlightLines, noInlineHighlight)}`
     +     `</code>`
     +     `${numLinesNumber}`
     +   `</pre>`
@@ -118,7 +119,8 @@ module.exports = (
       node.type = `html`
       node.value = `<code class="${className}">${highlightCode(
         languageName,
-        node.value
+        node.value,
+        escapeEntities
       )}</code>`
     })
   }