From 9ec1f12b33948df56c16f7900a07ddeeb8c9a334 Mon Sep 17 00:00:00 2001
From: "Joshua C. Jackson" <jjackson@cca.edu>
Date: Fri, 1 Dec 2017 13:44:39 -0800
Subject: [PATCH] Remove image, embed, document, and link entities on paste if
 not enabled (#103)

---
 lib/api/DraftUtils.js | 65 ++++++++++++++++++++++++++++++++++++++++++-
 lib/api/behavior.js   |  9 ++++--
 2 files changed, 71 insertions(+), 3 deletions(-)

diff --git a/lib/api/DraftUtils.js b/lib/api/DraftUtils.js
index 0c1119f0..80001178 100644
--- a/lib/api/DraftUtils.js
+++ b/lib/api/DraftUtils.js
@@ -222,8 +222,8 @@ export default {
         const shouldNormaliseType = block => {
             const type = block.getType();
             return (
-                // TODO Should handle atomic blocks separately.
                 type !== BLOCK_TYPE.UNSTYLED &&
+                type !== BLOCK_TYPE.ATOMIC &&
                 enabledBlockTypes.indexOf(type) === -1
             );
         };
@@ -242,6 +242,69 @@ export default {
         });
     },
 
+    /**
+     * Reset all entity types (images, links, documents, embeds) that are unavailable.
+     * Meant to be used after a paste of unconstrained content.
+     */
+    normalizeEntityType(editorState, enabledTypes) {
+        let contentState = editorState.getCurrentContent();
+        let blockMap = contentState.getBlockMap();
+
+        const isValidEntity = char => {
+            let isValid = true;
+            const entityKey = char.getEntity();
+            if (entityKey !== null) {
+                const entityType = contentState.getEntity(entityKey).getType();
+                if (enabledTypes.indexOf(entityType) === -1) {
+                    isValid = false;
+                }
+            }
+            return isValid;
+        };
+
+        const isValidAtomicBlock = block => {
+            // Remove invalid image and document blocks if not enabled
+            let isValidBlock = true;
+            block.findEntityRanges(
+                char => {
+                    isValidBlock = isValidEntity(char);
+                },
+                () => {},
+            );
+            return isValidBlock;
+        };
+
+        const isValidInline = block => {
+            let altered = false;
+
+            const chars = block.getCharacterList().filter(char => {
+                let newChar = char;
+                if (!isValidEntity(char)) {
+                    altered = true;
+                    newChar = CharacterMetadata.applyEntity(newChar, null);
+                }
+                return newChar;
+            });
+
+            return altered ? block.set('characterList', chars) : block;
+        };
+
+        blockMap = blockMap.filter(isValidAtomicBlock);
+        const blocks = blockMap.map(isValidInline);
+        blockMap = blockMap.merge(blocks);
+        contentState = contentState.merge({ blockMap });
+
+        let newEditorState = EditorState.set(editorState, {
+            currentContent: contentState,
+        });
+
+        // TODO: instead of moving selection and focus to the end, move it to previous block,
+        // if the last block in the paste selection was removed.
+        newEditorState = EditorState.moveSelectionToEnd(newEditorState);
+        newEditorState = EditorState.moveFocusToEnd(newEditorState);
+        return newEditorState;
+    },
+
     /**
      * Remove all styles that use unavailable types.
      * Meant to be used after a paste of unconstrained content.
diff --git a/lib/api/behavior.js b/lib/api/behavior.js
index fb1df62b..e4282797 100644
--- a/lib/api/behavior.js
+++ b/lib/api/behavior.js
@@ -281,11 +281,11 @@ export default {
         enableLineBreak,
         blockTypes = [],
         inlineStyles = [],
-        // TODO Implement.
-        // entityTypes = [],
+        entityTypes = [],
     ) {
         let nextEditorState = editorState;
         const enabledBlockTypes = blockTypes.map(type => type.type);
+        const enabledEntityTypes = entityTypes.map(type => type.type);
         const enabledInlineStyles = inlineStyles.map(type => type.type);
 
         nextEditorState = DraftUtils.normaliseBlockDepth(
@@ -298,6 +298,11 @@ export default {
             enabledBlockTypes,
         );
 
+        nextEditorState = DraftUtils.normalizeEntityType(
+            nextEditorState,
+            enabledEntityTypes,
+        );
+
         // TODO Re-test what happens when pasting a styled element, eg. h2 with bold, or blockquote with italic.
         nextEditorState = DraftUtils.normaliseInlineStyle(
             nextEditorState,