Fix segfault in Array#sample

2 files changed, 14 insertions, 1 deletions

diff --git a/mrbgems/mruby-random/src/random.c b/mrbgems/mruby-random/src/random.c
index 3b8d0dd2f..587e6c01a 100644
--- a/mrbgems/mruby-random/src/random.c
+++ b/mrbgems/mruby-random/src/random.c
@@ -266,7 +266,7 @@ mrb_ary_sample(mrb_state *mrb, mrb_value ary)
   mrb_int n = 0;
   mrb_bool given;
   mt_state *random = NULL;
-  mrb_int len = RARRAY_LEN(ary);
+  mrb_int len;
 
   mrb_get_args(mrb, "|i?d", &n, &given, &random, &mt_state_type);
   if (random == NULL) {
@@ -274,6 +274,7 @@ mrb_ary_sample(mrb_state *mrb, mrb_value ary)
   }
   mrb_random_rand_seed(mrb, random);
   mt_rand(random);
+  len = RARRAY_LEN(ary);
   if (!given) {                 /* pick one element */
     switch (len) {
     case 0:
diff --git a/mrbgems/mruby-random/test/random.rb b/mrbgems/mruby-random/test/random.rb
index 1653ae4a6..1c59be3a6 100644
--- a/mrbgems/mruby-random/test/random.rb
+++ b/mrbgems/mruby-random/test/random.rb
@@ -74,3 +74,15 @@ assert('Array#shuffle!(random)') do
 
   ary1 != [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] and 10.times { |x| ary1.include? x } and ary1 == ary2
 end
+
+assert('Array#sample checks input length after reading arguments') do
+  $ary = [1, 2, 3]
+  class ArrayChange
+    def to_i
+      $ary << 4
+      4
+    end
+  end
+
+  assert_equal [1, 2, 3, 4], $ary.sample(ArrayChange.new).sort
+end