Add size checks to Range#to_set and Enumerator#to_set [Bug #21654]

These two class are most common sources of infinite sequences.  This change should effectively prevent accidental infinite loops when calling to_set on them. [Bug #21513]

enumerator.c

@@ -3332,6 +3332,24 @@ enumerator_plus(VALUE obj, VALUE eobj)
     return new_enum_chain(rb_ary_new_from_args(2, obj, eobj));
 }
 
+/*
+ * call-seq:
+ *   e.to_set -> set
+ *
+ * Returns a set generated from this enumerator.
+ *
+ *   e = Enumerator.new { |y| y << 1 << 1 << 2 << 3 << 5 }
+ *   e.to_set #=> #<Set: {1, 2, 3, 5}>
+ */
+static VALUE enumerator_to_set(int argc, VALUE *argv, VALUE obj)
+{
+    VALUE size = rb_funcall(obj, id_size, 0);
+    if (RB_TYPE_P(size, T_FLOAT) && RFLOAT_VALUE(size) == INFINITY) {
+        rb_raise(rb_eArgError, "cannot convert an infinite enumerator to a set");
+    }
+    return rb_call_super(argc, argv);
+}
+
 /*
  * Document-class: Enumerator::Product
  *
@@ -4488,6 +4506,7 @@ InitVM_Enumerator(void)
     rb_define_method(rb_cEnumerator, "rewind", enumerator_rewind, 0);
     rb_define_method(rb_cEnumerator, "inspect", enumerator_inspect, 0);
     rb_define_method(rb_cEnumerator, "size", enumerator_size, 0);
+    rb_define_method(rb_cEnumerator, "to_set", enumerator_to_set, -1);
     rb_define_method(rb_cEnumerator, "+", enumerator_plus, 1);
     rb_define_method(rb_mEnumerable, "chain", enum_chain, -1);
 

range.c

@@ -1018,6 +1018,15 @@ range_to_a(VALUE range)
     return rb_call_super(0, 0);
 }
 
+static VALUE
+range_to_set(int argc, VALUE *argv, VALUE range)
+{
+    if (NIL_P(RANGE_END(range))) {
+        rb_raise(rb_eRangeError, "cannot convert endless range to a set");
+    }
+    return rb_call_super(argc, argv);
+}
+
 static VALUE
 range_enum_size(VALUE range, VALUE args, VALUE eobj)
 {
@@ -2845,6 +2854,7 @@ Init_Range(void)
     rb_define_method(rb_cRange, "minmax", range_minmax, 0);
     rb_define_method(rb_cRange, "size", range_size, 0);
     rb_define_method(rb_cRange, "to_a", range_to_a, 0);
+    rb_define_method(rb_cRange, "to_set", range_to_set, -1);
     rb_define_method(rb_cRange, "entries", range_to_a, 0);
     rb_define_method(rb_cRange, "to_s", range_to_s, 0);
     rb_define_method(rb_cRange, "inspect", range_inspect, 0);

test/ruby/test_enumerator.rb

@@ -1058,4 +1058,13 @@ class TestEnumerator < Test::Unit::TestCase
     enum = ary.each
     assert_equal(35.0, enum.sum)
   end
+
+  def test_to_set
+    e = Enumerator.new { it << 1 << 1 << 2 << 3 << 5 }
+    set = e.to_set
+    assert_equal(Set[1, 2, 3, 5], set)
+
+    ei = Enumerator.new(Float::INFINITY) { it << 1 << 1 << 2 << 3 << 5 }
+    assert_raise(ArgumentError) { ei.to_set }
+  end
 end

test/ruby/test_range.rb

@@ -1458,6 +1458,12 @@ class TestRange < Test::Unit::TestCase
     assert_raise(RangeError) { (1..).to_a }
   end
 
+  def test_to_set
+    assert_equal(Set[1,2,3,4,5], (1..5).to_set)
+    assert_equal(Set[1,2,3,4], (1...5).to_set)
+    assert_raise(RangeError) { (1..).to_set }
+  end
+
   def test_beginless_range_iteration
     assert_raise(TypeError) { (..1).each { } }
   end

test/ruby/test_set.rb

@@ -939,6 +939,30 @@ class TC_Enumerable < Test::Unit::TestCase
     assert_same set, set.to_set
     assert_not_same set, set.to_set { |o| o }
   end
+
+  class MyEnum
+    include Enumerable
+
+    def initialize(array)
+      @array = array
+    end
+
+    def each(&block)
+      @array.each(&block)
+    end
+
+    def size
+      raise "should not be called"
+    end
+  end
+
+  def test_to_set_not_calling_size
+    enum = MyEnum.new([1,2,3])
+
+    set = assert_nothing_raised { enum.to_set }
+    assert(set.is_a?(Set))
+    assert_equal(Set[1,2,3], set)
+  end
 end
 
 class TC_Set_Builtin < Test::Unit::TestCase