From a659c63e37414506dfb0d4655e031bb7a2e73fc8 Mon Sep 17 00:00:00 2001
From: Sutou Kouhei <kou@clear-code.com>
Date: Sat, 20 Feb 2021 07:22:57 +0900
Subject: [PATCH] Fix a bug that invalid notation declaration may be generated

HackerOne: HO-1104077

It's caused by quote character.

Reported by Juho Nurminen. Thanks!!!
---
 lib/rexml/doctype.rb | 24 +++++++++--
 test/test_doctype.rb | 99 +++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 118 insertions(+), 5 deletions(-)

diff --git a/lib/rexml/doctype.rb b/lib/rexml/doctype.rb
index dcfa0cfc..3e86cccc 100644
--- a/lib/rexml/doctype.rb
+++ b/lib/rexml/doctype.rb
@@ -255,13 +255,29 @@ def to_s
       c = nil
       c = parent.context if parent
       if c and c[:prologue_quote] == :apostrophe
-        quote = "'"
+        default_quote = "'"
       else
-        quote = "\""
+        default_quote = "\""
       end
       notation = "<!NOTATION #{@name} #{@middle}"
-      notation << " #{quote}#{@public}#{quote}" if @public
-      notation << " #{quote}#{@system}#{quote}" if @system
+      if @public
+        if @public.include?("'")
+          quote = "\""
+        else
+          quote = default_quote
+        end
+        notation << " #{quote}#{@public}#{quote}"
+      end
+      if @system
+        if @system.include?("'")
+          quote = "\""
+        elsif @system.include?("\"")
+          quote = "'"
+        else
+          quote = default_quote
+        end
+        notation << " #{quote}#{@system}#{quote}"
+      end
       notation << ">"
       notation
     end
diff --git a/test/test_doctype.rb b/test/test_doctype.rb
index a00c5d00..14284c94 100644
--- a/test/test_doctype.rb
+++ b/test/test_doctype.rb
@@ -89,11 +89,26 @@ def test_to_s
                    decl(@id, nil).to_s)
     end
 
+    def test_to_s_pubid_literal_include_apostrophe
+      assert_equal("<!NOTATION #{@name} PUBLIC \"#{@id}'\">",
+                   decl("#{@id}'", nil).to_s)
+    end
+
     def test_to_s_with_uri
       assert_equal("<!NOTATION #{@name} PUBLIC \"#{@id}\" \"#{@uri}\">",
                    decl(@id, @uri).to_s)
     end
 
+    def test_to_s_system_literal_include_apostrophe
+      assert_equal("<!NOTATION #{@name} PUBLIC \"#{@id}\" \"system'literal\">",
+                   decl(@id, "system'literal").to_s)
+    end
+
+    def test_to_s_system_literal_include_double_quote
+      assert_equal("<!NOTATION #{@name} PUBLIC \"#{@id}\" 'system\"literal'>",
+                   decl(@id, "system\"literal").to_s)
+    end
+
     def test_to_s_apostrophe
       document = REXML::Document.new(<<-XML)
       <!DOCTYPE root SYSTEM "urn:x-test:sysid" [
@@ -107,6 +122,49 @@ def test_to_s_apostrophe
                    notation.to_s)
     end
 
+    def test_to_s_apostrophe_pubid_literal_include_apostrophe
+      document = REXML::Document.new(<<-XML)
+      <!DOCTYPE root SYSTEM "urn:x-test:sysid" [
+        #{decl("#{@id}'", @uri).to_s}
+      ]>
+      <root/>
+      XML
+      # This isn't used for PubidLiteral because PubidChar includes '.
+      document.context[:prologue_quote] = :apostrophe
+      notation = document.doctype.notations[0]
+      assert_equal("<!NOTATION #{@name} PUBLIC \"#{@id}'\" '#{@uri}'>",
+                   notation.to_s)
+    end
+
+    def test_to_s_apostrophe_system_literal_include_apostrophe
+      document = REXML::Document.new(<<-XML)
+      <!DOCTYPE root SYSTEM "urn:x-test:sysid" [
+        #{decl(@id, "system'literal").to_s}
+      ]>
+      <root/>
+      XML
+      # This isn't used for SystemLiteral because SystemLiteral includes '.
+      document.context[:prologue_quote] = :apostrophe
+      notation = document.doctype.notations[0]
+      assert_equal("<!NOTATION #{@name} PUBLIC '#{@id}' \"system'literal\">",
+                   notation.to_s)
+    end
+
+    def test_to_s_apostrophe_system_literal_include_double_quote
+      document = REXML::Document.new(<<-XML)
+      <!DOCTYPE root SYSTEM "urn:x-test:sysid" [
+        #{decl(@id, "system\"literal").to_s}
+      ]>
+      <root/>
+      XML
+      # This isn't used for SystemLiteral because SystemLiteral includes ".
+      # But quoted by ' because SystemLiteral includes ".
+      document.context[:prologue_quote] = :apostrophe
+      notation = document.doctype.notations[0]
+      assert_equal("<!NOTATION #{@name} PUBLIC '#{@id}' 'system\"literal'>",
+                   notation.to_s)
+    end
+
     private
     def decl(id, uri)
       REXML::NotationDecl.new(@name, "PUBLIC", id, uri)
@@ -124,6 +182,16 @@ def test_to_s
                    decl(@id).to_s)
     end
 
+    def test_to_s_include_apostrophe
+      assert_equal("<!NOTATION #{@name} SYSTEM \"#{@id}'\">",
+                   decl("#{@id}'").to_s)
+    end
+
+    def test_to_s_include_double_quote
+      assert_equal("<!NOTATION #{@name} SYSTEM '#{@id}\"'>",
+                   decl("#{@id}\"").to_s)
+    end
+
     def test_to_s_apostrophe
       document = REXML::Document.new(<<-XML)
       <!DOCTYPE root SYSTEM "urn:x-test:sysid" [
@@ -137,9 +205,38 @@ def test_to_s_apostrophe
                    notation.to_s)
     end
 
+    def test_to_s_apostrophe_include_apostrophe
+      document = REXML::Document.new(<<-XML)
+      <!DOCTYPE root SYSTEM "urn:x-test:sysid" [
+        #{decl("#{@id}'").to_s}
+      ]>
+      <root/>
+      XML
+      # This isn't used for SystemLiteral because SystemLiteral includes '.
+      document.context[:prologue_quote] = :apostrophe
+      notation = document.doctype.notations[0]
+      assert_equal("<!NOTATION #{@name} SYSTEM \"#{@id}'\">",
+                   notation.to_s)
+    end
+
+    def test_to_s_apostrophe_include_double_quote
+      document = REXML::Document.new(<<-XML)
+      <!DOCTYPE root SYSTEM "urn:x-test:sysid" [
+        #{decl("#{@id}\"").to_s}
+      ]>
+      <root/>
+      XML
+      # This isn't used for SystemLiteral because SystemLiteral includes ".
+      # But quoted by ' because SystemLiteral includes ".
+      document.context[:prologue_quote] = :apostrophe
+      notation = document.doctype.notations[0]
+      assert_equal("<!NOTATION #{@name} SYSTEM '#{@id}\"'>",
+                   notation.to_s)
+    end
+
     private
     def decl(id)
-      REXML::NotationDecl.new(@name, "SYSTEM", id, nil)
+      REXML::NotationDecl.new(@name, "SYSTEM", nil, id)
     end
   end
 end