require 'test/unit'
require 'rails-security-backports'

class RubyCve_2008_3790Test < Test::Unit::TestCase

  def setup
    @orig_limit = REXML::Security.entity_expansion_limit
  end

  def teardown
    REXML::Security.entity_expansion_limit = @orig_limit
  end


  def test__record_entity_expansion__with_small_num_expansions
    REXML::Security.entity_expansion_limit = 50

    xml = get_expandable_xml_that_expands_to(REXML::Security.entity_expansion_limit)

    assert_nothing_raised(RuntimeError, "Expected NO exception with xml:\n#{xml}") do
      REXML::Document.new(xml).root.text
    end
  end

  def test__record_entity_expansion__with_too_many_expansions
    REXML::Security.entity_expansion_limit = 50

    xml = get_expandable_xml_that_expands_to(REXML::Security.entity_expansion_limit + 1000)

    assert_raise(RuntimeError, "Expected exception with xml:\n#{xml}") do
      REXML::Document.new(xml).root.text
    end
  end

  private

  def get_expandable_xml_that_expands_to(num_expansions_required = 51)
    expansion_keys = %w{a b}
    expansion_text = "x"*10

    expansions_per_key = Math.sqrt(num_expansions_required).floor - 1
    num_expansions_expected = expansions_per_key ** 2
    num_expansions_expected += 1 + expansions_per_key

    non_exponential_expansions = num_expansions_required - 1 - num_expansions_expected

    <<-XML
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE member [
<!ENTITY a "#{'&b;'*expansions_per_key}">
<!ENTITY b "#{'&x;'*expansions_per_key}">
<!ENTITY x "#{expansion_text}">
<!ENTITY c "#{'&y;'*non_exponential_expansions}">
<!ENTITY y "#{expansion_text}">
]>
<member>
&a;
&c;
</member>
    XML
  end

end