Get String length after args in String#chomp!

Fixes RCE issue
Reported by @bouk

2 files changed, 16 insertions, 2 deletions

diff --git a/src/string.c b/src/string.c
index 5e490bf03..f47294291 100644
--- 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
index e67389b5c..80fcbe6fa 100644
--- a/test/t/string.rb
+++ b/test/t/string.rb
@@ -251,6 +251,19 @@ assert('String#chomp!', '15.2.10.5.10') do
   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('String#freeze') do
 
   assert_raise(RuntimeError) { str.upcase! }
 end
-