Get String length after args in String#chomp!

Fixes RCE issue Reported by @bouk
mruby · Nov 24, 2016 · 76a1bdf · 76a1bdf
1 parent a630c4f
commit 76a1bdf
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/src/string.c b/src/string.c
@@ -1235,11 +1235,13 @@ mrb_str_chomp_bang(mrb_state *mrb, mrb_value str)
   char *p, *pp;
   mrb_int rslen;
   mrb_int len;
+  mrb_int argc;
   struct RString *s = mrb_str_ptr(str);
 
   mrb_str_modify(mrb, s);
+  argc = mrb_get_args(mrb, "|S", &rs);
   len = RSTR_LEN(s);
-  if (mrb_get_args(mrb, "|S", &rs) == 0) {
+  if (argc == 0) {
     if (len == 0) return mrb_nil_value();
   smart_chomp:
     if (RSTR_PTR(s)[len-1] == '\n') {

diff --git a/test/t/string.rb b/test/t/string.rb
@@ -251,6 +251,19 @@
   assert_equal 'abc', e
 end
 
+assert('String#chomp! uses the correct length') do
+  class A
+    def to_str
+      $s.replace("AA")
+      "A"
+    end
+  end
+
+  $s = "AAA"
+  $s.chomp!(A.new)
+  assert_equal $s, "A"
+end
+
 assert('String#chop', '15.2.10.5.11') do
   a = ''.chop
   b = 'abc'.chop
@@ -683,4 +696,3 @@
 
   assert_raise(RuntimeError) { str.upcase! }
 end
-