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

class RubyCve_2014_8080Test < Test::Unit::TestCase

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

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

  def test__record_entity_expansion__with_small_num_text_expansions
    REXML::Security.entity_expansion_text_limit = 50

    xml = get_expandable_xml_that_expands_to(:less_than=>REXML::Security.entity_expansion_text_limit)

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

  def test__record_entity_expansion__with_too_many_text_expansions
    REXML::Security.entity_expansion_text_limit = 50

    xml = get_expandable_xml_that_expands_to(:more_than=>REXML::Security.entity_expansion_text_limit)

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

  private

  def get_expandable_xml_that_expands_to(opts = {})
    more_or_less = nil
    num_text_expansions_required = if opts.has_key?(:less_than)
      more_or_less = :less
      opts[:less_than].to_i
    elsif opts.has_key?(:more_than)
      more_or_less = :more
      opts[:more_than].to_i
    else
      raise ArgumentError.new("Argument must be either `:less_than=>x` or `:more_than=>x`, but was #{opts.inspect}")
    end

    expansion_keys = %w{a b}
    expansion_text = "x"*1

    expansions_per_key = Math.sqrt(num_text_expansions_required)
    expansions_per_key = more_or_less == :more ? expansions_per_key.ceil : expansions_per_key.floor

    <<-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}">
]>
<member>
&a;
</member>
    XML
  end

end