From 4fd8b6b1e315c8e4adcf6ec623987943ee254169 Mon Sep 17 00:00:00 2001
From: Viktor Ivarsson <viktor.ivarsson@twingly.com>
Date: Wed, 7 Aug 2024 10:25:19 +0200
Subject: [PATCH] Improve `#unnormalize` and fix `sum` calculation

* Improve `#unnormalize` by only iterating over unique matches
* Fix bug where `sum` for `#unnormalize` is calculated multiple times over causing a runtime error "entity expansion has grown too large"
* Adjust tests to the reflect the changes to the `entity_expansion_count`

See #193
---
 lib/rexml/parsers/baseparser.rb |  2 +-
 test/test_pullparser.rb         | 17 ++++++++---------
 test/test_sax.rb                | 13 ++++++-------
 3 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/lib/rexml/parsers/baseparser.rb b/lib/rexml/parsers/baseparser.rb
index 28810bfa..0a5a8ac0 100644
--- a/lib/rexml/parsers/baseparser.rb
+++ b/lib/rexml/parsers/baseparser.rb
@@ -549,7 +549,7 @@ def unnormalize( string, entities=nil, filter=nil )
         matches.collect!{|x|x[0]}.compact!
         if matches.size > 0
           sum = 0
-          matches.each do |entity_reference|
+          matches.uniq.each do |entity_reference|
             unless filter and filter.include?(entity_reference)
               entity_value = entity( entity_reference, entities )
               if entity_value
diff --git a/test/test_pullparser.rb b/test/test_pullparser.rb
index 55205af8..f9808bab 100644
--- a/test/test_pullparser.rb
+++ b/test/test_pullparser.rb
@@ -204,21 +204,20 @@ def test_empty_value
 </member>
           XML
 
+          REXML::Security.entity_expansion_limit = 5
           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
+          while parser.has_next?
+            parser.pull
           end
 
-          REXML::Security.entity_expansion_limit = 100
+          REXML::Security.entity_expansion_limit = 4
           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)
+          assert_equal(5, parser.entity_expansion_count)
         end
 
         def test_with_default_entity
@@ -235,15 +234,15 @@ def test_with_default_entity
 </member>
           XML
 
-          REXML::Security.entity_expansion_limit = 4
+          REXML::Security.entity_expansion_limit = 3
           parser = REXML::Parsers::PullParser.new(source)
           while parser.has_next?
             parser.pull
           end
 
-          REXML::Security.entity_expansion_limit = 3
+          REXML::Security.entity_expansion_limit = 2
           parser = REXML::Parsers::PullParser.new(source)
-          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
+          assert_raise(RuntimeError) do
             while parser.has_next?
               parser.pull
             end
diff --git a/test/test_sax.rb b/test/test_sax.rb
index 5e3ad75b..bb7bbd96 100644
--- a/test/test_sax.rb
+++ b/test/test_sax.rb
@@ -145,17 +145,16 @@ def test_empty_value
 </member>
           XML
 
+          REXML::Security.entity_expansion_limit = 5
           sax = REXML::Parsers::SAX2Parser.new(source)
-          assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
-            sax.parse
-          end
+          sax.parse
 
-          REXML::Security.entity_expansion_limit = 100
+          REXML::Security.entity_expansion_limit = 4
           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)
+          assert_equal(5, sax.entity_expansion_count)
         end
 
         def test_with_default_entity
@@ -172,11 +171,11 @@ def test_with_default_entity
 </member>
           XML
 
-          REXML::Security.entity_expansion_limit = 4
+          REXML::Security.entity_expansion_limit = 3
           sax = REXML::Parsers::SAX2Parser.new(source)
           sax.parse
 
-          REXML::Security.entity_expansion_limit = 3
+          REXML::Security.entity_expansion_limit = 2
           sax = REXML::Parsers::SAX2Parser.new(source)
           assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
             sax.parse