Skip to content

Commit 61500c6

Browse files
committed
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]
1 parent 25c871f commit 61500c6

File tree

5 files changed

+68
-0
lines changed

5 files changed

+68
-0
lines changed

enumerator.c

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3332,6 +3332,24 @@ enumerator_plus(VALUE obj, VALUE eobj)
33323332
return new_enum_chain(rb_ary_new_from_args(2, obj, eobj));
33333333
}
33343334

3335+
/*
3336+
* call-seq:
3337+
* e.to_set -> set
3338+
*
3339+
* Returns a set generated from this enumerator.
3340+
*
3341+
* e = Enumerator.new { |y| y << 1 << 1 << 2 << 3 << 5 }
3342+
* e.to_set #=> #<Set: {1, 2, 3, 5}>
3343+
*/
3344+
static VALUE enumerator_to_set(int argc, VALUE *argv, VALUE obj)
3345+
{
3346+
VALUE size = rb_funcall(obj, id_size, 0);
3347+
if (RB_TYPE_P(size, T_FLOAT) && RFLOAT_VALUE(size) == INFINITY) {
3348+
rb_raise(rb_eArgError, "cannot convert an infinite enumerator to a set");
3349+
}
3350+
return rb_call_super(argc, argv);
3351+
}
3352+
33353353
/*
33363354
* Document-class: Enumerator::Product
33373355
*
@@ -4488,6 +4506,7 @@ InitVM_Enumerator(void)
44884506
rb_define_method(rb_cEnumerator, "rewind", enumerator_rewind, 0);
44894507
rb_define_method(rb_cEnumerator, "inspect", enumerator_inspect, 0);
44904508
rb_define_method(rb_cEnumerator, "size", enumerator_size, 0);
4509+
rb_define_method(rb_cEnumerator, "to_set", enumerator_to_set, -1);
44914510
rb_define_method(rb_cEnumerator, "+", enumerator_plus, 1);
44924511
rb_define_method(rb_mEnumerable, "chain", enum_chain, -1);
44934512

range.c

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1018,6 +1018,15 @@ range_to_a(VALUE range)
10181018
return rb_call_super(0, 0);
10191019
}
10201020

1021+
static VALUE
1022+
range_to_set(int argc, VALUE *argv, VALUE range)
1023+
{
1024+
if (NIL_P(RANGE_END(range))) {
1025+
rb_raise(rb_eRangeError, "cannot convert endless range to a set");
1026+
}
1027+
return rb_call_super(argc, argv);
1028+
}
1029+
10211030
static VALUE
10221031
range_enum_size(VALUE range, VALUE args, VALUE eobj)
10231032
{
@@ -2845,6 +2854,7 @@ Init_Range(void)
28452854
rb_define_method(rb_cRange, "minmax", range_minmax, 0);
28462855
rb_define_method(rb_cRange, "size", range_size, 0);
28472856
rb_define_method(rb_cRange, "to_a", range_to_a, 0);
2857+
rb_define_method(rb_cRange, "to_set", range_to_set, -1);
28482858
rb_define_method(rb_cRange, "entries", range_to_a, 0);
28492859
rb_define_method(rb_cRange, "to_s", range_to_s, 0);
28502860
rb_define_method(rb_cRange, "inspect", range_inspect, 0);

test/ruby/test_enumerator.rb

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1058,4 +1058,13 @@ def initialize(val)
10581058
enum = ary.each
10591059
assert_equal(35.0, enum.sum)
10601060
end
1061+
1062+
def test_to_set
1063+
e = Enumerator.new { it << 1 << 1 << 2 << 3 << 5 }
1064+
set = e.to_set
1065+
assert_equal(Set[1, 2, 3, 5], set)
1066+
1067+
ei = Enumerator.new(Float::INFINITY) { it << 1 << 1 << 2 << 3 << 5 }
1068+
assert_raise(ArgumentError) { ei.to_set }
1069+
end
10611070
end

test/ruby/test_range.rb

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1458,6 +1458,12 @@ def test_to_a
14581458
assert_raise(RangeError) { (1..).to_a }
14591459
end
14601460

1461+
def test_to_set
1462+
assert_equal(Set[1,2,3,4,5], (1..5).to_set)
1463+
assert_equal(Set[1,2,3,4], (1...5).to_set)
1464+
assert_raise(RangeError) { (1..).to_set }
1465+
end
1466+
14611467
def test_beginless_range_iteration
14621468
assert_raise(TypeError) { (..1).each { } }
14631469
end

test/ruby/test_set.rb

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -939,6 +939,30 @@ def test_to_set
939939
assert_same set, set.to_set
940940
assert_not_same set, set.to_set { |o| o }
941941
end
942+
943+
class MyEnum
944+
include Enumerable
945+
946+
def initialize(array)
947+
@array = array
948+
end
949+
950+
def each(&block)
951+
@array.each(&block)
952+
end
953+
954+
def size
955+
raise "should not be called"
956+
end
957+
end
958+
959+
def test_to_set_not_calling_size
960+
enum = MyEnum.new([1,2,3])
961+
962+
set = assert_nothing_raised { enum.to_set }
963+
assert(set.is_a?(Set))
964+
assert_equal(Set[1,2,3], set)
965+
end
942966
end
943967

944968
class TC_Set_Builtin < Test::Unit::TestCase

0 commit comments

Comments
 (0)