Added PermitScrubber which allows you to permit elements for sanitization.

2 files changed, 85 insertions, 13 deletions

diff --git a/actionview/lib/action_view/helpers/sanitize_helper/permit_scrubber.rb b/actionview/lib/action_view/helpers/sanitize_helper/permit_scrubber.rb
new file mode 100644
index 0000000000..749bef23fc
--- /dev/null
+++ b/actionview/lib/action_view/helpers/sanitize_helper/permit_scrubber.rb
@@ -0,0 +1,71 @@
+# === PermitScrubber
+#
+# PermitScrubber allows you to permit only your own tags and/or attributes.
+#
+# Supplied tags and attributes should be Enumerables
+#
+# +tags=+
+# If this value is set all other elements will be stripped (their inner elements will be kept).
+# If not set elements for which HTML5::Scrub.allowed_element? is false will be stripped.
+#
+# +attributes=+
+# Contain an elements allowed attributes.
+# If none is set HTML5::Scrub.scrub_attributes implementation will be used.
+class PermitScrubber < Loofah::Scrubber
+  attr_reader :tags, :attributes
+
+  def tags=(tags)
+    @tags = validate!(tags, :tags)
+  end
+
+  def attributes=(attributes)
+    @attributes = validate!(attributes, :attributes)
+  end
+
+  def scrub(node)
+    return CONTINUE if text_or_cdata_node?(node)
+
+    unless allowed_node?(node)
+      node.before node.children # strip
+      node.remove
+      return STOP
+    end
+
+    scrub_attributes(node)
+  end
+
+  protected
+
+  def allowed_node?(node)
+    if @tags
+      @tags.include?(node.name)
+    else
+      Loofah::HTML5::Scrub.allowed_element?(node.name)
+    end
+  end
+
+  def scrub_attributes(node)
+    if @attributes
+      node.attributes.each do |name, _|
+        node.remove_attribute(name) unless @attributes.include?(name)
+      end
+    else
+      Loofah::HTML5::Scrub.scrub_attributes(node)
+    end
+  end
+
+  def text_or_cdata_node?(node)
+    case node.type
+    when Nokogiri::XML::Node::TEXT_NODE, Nokogiri::XML::Node::CDATA_SECTION_NODE
+      return true
+    end
+    false
+  end
+
+  def validate!(var, name)
+    if var && !var.is_a?(Enumerable)
+      raise ArgumentError, "You should pass :#{name} as an Enumerable"
+    end
+    var
+  end
+end
diff --git a/actionview/lib/action_view/helpers/sanitize_helper/sanitizers.rb b/actionview/lib/action_view/helpers/sanitize_helper/sanitizers.rb
index 91de4c8ba1..cbddf3481c 100644
--- a/actionview/lib/action_view/helpers/sanitize_helper/sanitizers.rb
+++ b/actionview/lib/action_view/helpers/sanitize_helper/sanitizers.rb
@@ -1,5 +1,6 @@
 require 'active_support/core_ext/class/attribute'
 require 'active_support/deprecation'
+require 'action_view/helpers/sanitize_helper/permit_scrubber'
 require 'loofah'
 
 module ActionView
@@ -25,13 +26,23 @@ module ActionView
   end
 
   class WhiteListSanitizer
+
+    def initialize
+      @permit_scrubber = PermitScrubber.new
+    end
+
     def sanitize(html, options = {})
       return nil unless html
-      validate_options(options)
 
       loofah_fragment = Loofah.fragment(html)
-      loofah_fragment.scrub!(:strip)
-      loofah_fragment.xpath("./form").each { |form| form.remove }
+      if options[:tags] || options[:attributes]
+        @permit_scrubber.tags = options[:tags]
+        @permit_scrubber.attributes = options[:attributes]
+        loofah_fragment.scrub!(@permit_scrubber)
+      else
+        loofah_fragment.scrub!(:strip)
+        loofah_fragment.xpath("./form").each { |form| form.remove }
+      end
       loofah_fragment.to_s
     end
 
@@ -97,16 +108,6 @@ module ActionView
     self.allowed_protocols = Loofah::HTML5::WhiteList::ALLOWED_PROTOCOLS
 
     protected
-      def validate_options(options)
-        if options[:tags] && !options[:tags].is_a?(Enumerable)
-          raise ArgumentError, "You should pass :tags as an Enumerable"
-        end
-
-        if options[:attributes] && !options[:attributes].is_a?(Enumerable)
-          raise ArgumentError, "You should pass :attributes as an Enumerable"
-        end
-      end
-
       def contains_bad_protocols?(attr_name, value)
         protocol_separator = ':'
         self.uri_attributes.include?(attr_name) &&