From dc48407c1cdf4ad14e4384ee137b7251ccad7bdc Mon Sep 17 00:00:00 2001
From: NAITOH Jun <naitoh@gmail.com>
Date: Mon, 19 Aug 2024 16:24:23 +0900
Subject: [PATCH] Add support for XML entity expansion limitation in Stream
 parser

## Why?

See:
- https://github.com/ruby/rexml/pull/187
- https://github.com/ruby/rexml/pull/195

## Change
- Supported `REXML::Security.entity_expansion_limit=` in Stream parser
- Supported `REXML::Security.entity_expansion_text_limit=` in Stream parser
---
 lib/rexml/parsers/streamparser.rb |   4 ++
 test/test_stream.rb               | 101 ++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/lib/rexml/parsers/streamparser.rb b/lib/rexml/parsers/streamparser.rb
index f97c9124..7781fe44 100644
--- a/lib/rexml/parsers/streamparser.rb
+++ b/lib/rexml/parsers/streamparser.rb
@@ -14,6 +14,10 @@ def add_listener( listener )
         @parser.add_listener( listener )
       end
 
+      def entity_expansion_count
+        @parser.entity_expansion_count
+      end
+
       def parse
         # entity string
         while true
diff --git a/test/test_stream.rb b/test/test_stream.rb
index f08f18bf..4e769266 100644
--- a/test/test_stream.rb
+++ b/test/test_stream.rb
@@ -134,6 +134,81 @@ def teardown
       REXML::Security.entity_expansion_text_limit = @default_entity_expansion_text_limit
     end
 
+    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
+
+      assert_raise(RuntimeError.new("entity expansion has grown too large")) do
+        REXML::Document.parse_stream(source, MyListener.new)
+      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
+
+      listener = MyListener.new
+      REXML::Security.entity_expansion_limit = 100000
+      parser = REXML::Parsers::StreamParser.new( source, listener )
+      parser.parse
+      assert_equal(11111, parser.entity_expansion_count)
+
+      REXML::Security.entity_expansion_limit = @default_entity_expansion_limit
+      parser = REXML::Parsers::StreamParser.new( source, listener )
+      assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
+        parser.parse
+      end
+      assert do
+        parser.entity_expansion_count > @default_entity_expansion_limit
+      end
+    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
+
+      listener = MyListener.new
+      REXML::Security.entity_expansion_limit = 4
+      REXML::Document.parse_stream(source, listener)
+
+      REXML::Security.entity_expansion_limit = 3
+      assert_raise(RuntimeError.new("number of entity expansions exceeded, processing aborted.")) do
+        REXML::Document.parse_stream(source, listener)
+      end
+    end
+
     def test_with_only_default_entities
       member_value = "&lt;p&gt;#{'A' * @default_entity_expansion_text_limit}&lt;/p&gt;"
       source = <<-XML
@@ -159,6 +234,32 @@ def text(text)
         listener.text_value.bytesize > @default_entity_expansion_text_limit
       end
     end
+
+    def test_entity_expansion_text_limit
+      source = <<-XML
+<!DOCTYPE member [
+  <!ENTITY a "&b;&b;&b;">
+  <!ENTITY b "&c;&d;&e;">
+  <!ENTITY c "xxxxxxxxxx">
+  <!ENTITY d "yyyyyyyyyy">
+  <!ENTITY e "zzzzzzzzzz">
+]>
+<member>&a;</member>
+      XML
+
+      listener = MyListener.new
+      class << listener
+        attr_accessor :text_value
+        def text(text)
+          @text_value << text
+        end
+      end
+      listener.text_value = ""
+      REXML::Security.entity_expansion_text_limit = 90
+      REXML::Document.parse_stream(source, listener)
+
+      assert_equal(90, listener.text_value.size)
+    end
   end
 
   # For test_listener