Add support for XML entity expansion limitation in SAX and pull parsers

lib/rexml/parsers/baseparser.rb

@@ -154,13 +154,15 @@ def initialize( source )
         self.stream = source
         @listeners = []
         @prefixes = Set.new
+        @entity_expansion_count = 0
       end
 
       def add_listener( listener )
         @listeners << listener
       end
 
       attr_reader :source
+      attr_reader :entity_expansion_count
 
       def stream=( source )
         @source = SourceFactory.create_from( source )
@@ -513,7 +515,9 @@ def pull_event
       def entity( reference, entities )
         value = nil
         value = entities[ reference ] if entities
-        if not value
+        if value
+          record_entity_expansion
+        else
           value = DEFAULT_ENTITIES[ reference ]
           value = value[2] if value
         end
@@ -552,12 +556,17 @@ def unnormalize( string, entities=nil, filter=nil )
         }
         matches.collect!{|x|x[0]}.compact!
         if matches.size > 0
+          sum = 0
           matches.each do |entity_reference|
             unless filter and filter.include?(entity_reference)
               entity_value = entity( entity_reference, entities )
               if entity_value
                 re = Private::DEFAULT_ENTITIES_PATTERNS[entity_reference] || /&#{entity_reference};/
                 rv.gsub!( re, entity_value )
+                sum += rv.bytesize
+                if sum > Security.entity_expansion_text_limit
+                  raise "entity expansion has grown too large"
+                end
               else
                 er = DEFAULT_ENTITIES[entity_reference]
                 rv.gsub!( er[0], er[2] ) if er
@@ -570,6 +579,14 @@ def unnormalize( string, entities=nil, filter=nil )
       end
 
       private
+
+      def record_entity_expansion
+        @entity_expansion_count += 1
+        if @entity_expansion_count > Security.entity_expansion_limit
+          raise "number of entity expansions exceeded, processing aborted."
+        end
+      end
+
       def need_source_encoding_update?(xml_declaration_encoding)
         return false if xml_declaration_encoding.nil?
         return false if /\AUTF-16\z/i =~ xml_declaration_encoding

lib/rexml/parsers/pullparser.rb

@@ -47,6 +47,10 @@ def add_listener( listener )
         @listeners << listener
       end
 
+      def entity_expansion_count
+        @parser.entity_expansion_count
+      end
+
       def each
         while has_next?
           yield self.pull

lib/rexml/parsers/sax2parser.rb

@@ -22,6 +22,10 @@ def source
         @parser.source
       end
 
+      def entity_expansion_count
+        @parser.entity_expansion_count
+      end
+
       def add_listener( listener )
         @parser.add_listener( listener )
       end

test/test_document.rb

@@ -41,7 +41,7 @@ def teardown
 
       class GeneralEntityTest < self
         def test_have_value
-          xml = <<EOF
+          xml = <<XML
 <?xml version="1.0" encoding="UTF-8"?>
 <!DOCTYPE member [
   <!ENTITY a "&b;&b;&b;&b;&b;&b;&b;&b;&b;&b;">
@@ -55,23 +55,24 @@ def test_have_value
 <member>
 &a;
 </member>
-EOF
+XML
 
           doc = REXML::Document.new(xml)
-          assert_raise(RuntimeError) do
+          assert_raise(RuntimeError.new("entity expansion has grown too large")) do
             doc.root.children.first.value
           end
+
           REXML::Security.entity_expansion_limit = 100
           assert_equal(100, REXML::Security.entity_expansion_limit)
           doc = REXML::Document.new(xml)
-          assert_raise(RuntimeError) do
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
             doc.root.children.first.value
           end
           assert_equal(101, doc.entity_expansion_count)
         end
 
         def test_empty_value
-          xml = <<EOF
+          xml = <<XML
 <?xml version="1.0" encoding="UTF-8"?>
 <!DOCTYPE member [
   <!ENTITY a "&b;&b;&b;&b;&b;&b;&b;&b;&b;&b;">
@@ -85,23 +86,24 @@ def test_empty_value
 <member>
 &a;
 </member>
-EOF
+XML
 
           doc = REXML::Document.new(xml)
-          assert_raise(RuntimeError) do
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
             doc.root.children.first.value
           end
+
           REXML::Security.entity_expansion_limit = 100
           assert_equal(100, REXML::Security.entity_expansion_limit)
           doc = REXML::Document.new(xml)
-          assert_raise(RuntimeError) do
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
             doc.root.children.first.value
           end
           assert_equal(101, doc.entity_expansion_count)
         end
 
         def test_with_default_entity
-          xml = <<EOF
+          xml = <<XML
 <?xml version="1.0" encoding="UTF-8"?>
 <!DOCTYPE member [
   <!ENTITY a "a">
@@ -112,14 +114,15 @@ def test_with_default_entity
 &a2;
 &lt;
 </member>
-EOF
+XML
 
           REXML::Security.entity_expansion_limit = 4
           doc = REXML::Document.new(xml)
           assert_equal("\na\na a\n<\n", doc.root.children.first.value)
+
           REXML::Security.entity_expansion_limit = 3
           doc = REXML::Document.new(xml)
-          assert_raise(RuntimeError) do
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
             doc.root.children.first.value
           end
         end

test/test_pullparser.rb

@@ -155,5 +155,101 @@ def test_peek
       end
       assert_equal( 0, names.length )
     end
+
+    class EntityExpansionLimitTest < Test::Unit::TestCase
+      def setup
+        @default_entity_expansion_limit = REXML::Security.entity_expansion_limit
+      end
+
+      def teardown
+        REXML::Security.entity_expansion_limit = @default_entity_expansion_limit
+      end
+
+      class GeneralEntityTest < self
+        def test_have_value
+          source = <<-XML
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE member [
+  <!ENTITY a "&b;&b;&b;&b;&b;&b;&b;&b;&b;&b;">
+  <!ENTITY b "&c;&c;&c;&c;&c;&c;&c;&c;&c;&c;">
+  <!ENTITY c "&d;&d;&d;&d;&d;&d;&d;&d;&d;&d;">
+  <!ENTITY d "&e;&e;&e;&e;&e;&e;&e;&e;&e;&e;">
+  <!ENTITY e "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx">
+]>
+<member>
+&a;
+</member>
+          XML
+
+          parser = REXML::Parsers::PullParser.new(source)
+          assert_raise(RuntimeError.new("entity expansion has grown too large")) do
+            while parser.has_next?
+              parser.pull
+            end
+          end
+        end
+
+        def test_empty_value
+          source = <<-XML
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE member [
+  <!ENTITY a "&b;&b;&b;&b;&b;&b;&b;&b;&b;&b;">
+  <!ENTITY b "&c;&c;&c;&c;&c;&c;&c;&c;&c;&c;">
+  <!ENTITY c "&d;&d;&d;&d;&d;&d;&d;&d;&d;&d;">
+  <!ENTITY d "&e;&e;&e;&e;&e;&e;&e;&e;&e;&e;">
+  <!ENTITY e "">
+]>
+<member>
+&a;
+</member>
+          XML
+
+          parser = REXML::Parsers::PullParser.new(source)
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
+            while parser.has_next?
+              parser.pull
+            end
+          end
+
+          REXML::Security.entity_expansion_limit = 100
+          parser = REXML::Parsers::PullParser.new(source)
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
+            while parser.has_next?
+              parser.pull
+            end
+          end
+          assert_equal(101, parser.entity_expansion_count)
+        end
+
+        def test_with_default_entity
+          source = <<-XML
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE member [
+  <!ENTITY a "a">
+  <!ENTITY a2 "&a; &a;">
+]>
+<member>
+&a;
+&a2;
+&lt;
+</member>
+          XML
+
+          REXML::Security.entity_expansion_limit = 4
+          parser = REXML::Parsers::PullParser.new(source)
+          while parser.has_next?
+            parser.pull
+          end
+
+          REXML::Security.entity_expansion_limit = 3
+          parser = REXML::Parsers::PullParser.new(source)
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
+            while parser.has_next?
+              parser.pull
+            end
+          end
+        end
+      end
+    end
   end
 end

test/test_sax.rb

@@ -99,6 +99,92 @@ def test_sax2
       end
     end
 
+    class EntityExpansionLimitTest < Test::Unit::TestCase
+      def setup
+        @default_entity_expansion_limit = REXML::Security.entity_expansion_limit
+      end
+
+      def teardown
+        REXML::Security.entity_expansion_limit = @default_entity_expansion_limit
+      end
+
+      class GeneralEntityTest < self
+        def test_have_value
+          source = <<-XML
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE member [
+  <!ENTITY a "&b;&b;&b;&b;&b;&b;&b;&b;&b;&b;">
+  <!ENTITY b "&c;&c;&c;&c;&c;&c;&c;&c;&c;&c;">
+  <!ENTITY c "&d;&d;&d;&d;&d;&d;&d;&d;&d;&d;">
+  <!ENTITY d "&e;&e;&e;&e;&e;&e;&e;&e;&e;&e;">
+  <!ENTITY e "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx">
+]>
+<member>
+&a;
+</member>
+          XML
+
+          sax = REXML::Parsers::SAX2Parser.new(source)
+          assert_raise(RuntimeError.new("entity expansion has grown too large")) do
+            sax.parse
+          end
+        end
+
+        def test_empty_value
+          source = <<-XML
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE member [
+  <!ENTITY a "&b;&b;&b;&b;&b;&b;&b;&b;&b;&b;">
+  <!ENTITY b "&c;&c;&c;&c;&c;&c;&c;&c;&c;&c;">
+  <!ENTITY c "&d;&d;&d;&d;&d;&d;&d;&d;&d;&d;">
+  <!ENTITY d "&e;&e;&e;&e;&e;&e;&e;&e;&e;&e;">
+  <!ENTITY e "">
+]>
+<member>
+&a;
+</member>
+          XML
+
+          sax = REXML::Parsers::SAX2Parser.new(source)
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
+            sax.parse
+          end
+
+          REXML::Security.entity_expansion_limit = 100
+          sax = REXML::Parsers::SAX2Parser.new(source)
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
+            sax.parse
+          end
+          assert_equal(101, sax.entity_expansion_count)
+        end
+
+        def test_with_default_entity
+          source = <<-XML
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE member [
+  <!ENTITY a "a">
+  <!ENTITY a2 "&a; &a;">
+]>
+<member>
+&a;
+&a2;
+&lt;
+</member>
+          XML
+
+          REXML::Security.entity_expansion_limit = 4
+          sax = REXML::Parsers::SAX2Parser.new(source)
+          sax.parse
+
+          REXML::Security.entity_expansion_limit = 3
+          sax = REXML::Parsers::SAX2Parser.new(source)
+          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
+            sax.parse
+          end
+        end
+      end
+    end
+
     # used by test_simple_doctype_listener
     # submitted by Jeff Barczewski
     class SimpleDoctypeListener