Fix class/instance variable name validation

- `@@?` etc are invalid class variable name.
- `@1` etc are invalid instance variable name.

4 files changed, 44 insertions, 6 deletions

diff --git a/mrbgems/mruby-metaprog/src/metaprog.c b/mrbgems/mruby-metaprog/src/metaprog.c
index 75913dab5..0aafb4c34 100644
--- a/mrbgems/mruby-metaprog/src/metaprog.c
+++ b/mrbgems/mruby-metaprog/src/metaprog.c
@@ -403,7 +403,8 @@ mod_define_singleton_method(mrb_state *mrb, mrb_value self)
 static mrb_bool
 cv_name_p(mrb_state *mrb, const char *name, mrb_int len)
 {
-  return len > 2 && name[0] == '@' && name[1] == '@';
+  return len > 2 && name[0] == '@' && name[1] == '@' &&
+         !ISDIGIT(name[2]) && mrb_ident_p(name+2, len-2);
 }
 
 static void
diff --git a/mrbgems/mruby-metaprog/test/metaprog.rb b/mrbgems/mruby-metaprog/test/metaprog.rb
index 8bc3719b8..1262c9945 100644
--- a/mrbgems/mruby-metaprog/test/metaprog.rb
+++ b/mrbgems/mruby-metaprog/test/metaprog.rb
@@ -34,6 +34,30 @@ assert('Kernel#instance_variable_defined?', '15.3.1.3.20') do
   assert_false o.instance_variable_defined?("@b")
   assert_true o.instance_variable_defined?("@a"[0,2])
   assert_true o.instance_variable_defined?("@abc"[0,2])
+  assert_raise(NameError) { o.instance_variable_defined?("@0") }
+end
+
+assert('Kernel#instance_variable_get', '15.3.1.3.21') do
+  o = Class.new { attr_accessor :foo, :bar }.new
+  o.foo = "one"
+  o.bar = 2
+  assert_equal("one", o.instance_variable_get(:@foo))
+  assert_equal(2, o.instance_variable_get("@bar"))
+  assert_equal(nil, o.instance_variable_get(:@baz))
+  %w[foo @1].each do |n|
+    assert_raise(NameError) { o.instance_variable_get(n) }
+  end
+end
+
+assert('Kernel#instance_variable_set', '15.3.1.3.22') do
+  o = Class.new { attr_reader :foo, :_bar }.new
+  assert_equal("one", o.instance_variable_set(:@foo, "one"))
+  assert_equal("one", o.foo)
+  assert_equal(2, o.instance_variable_set("@_bar", 2))
+  assert_equal(2, o._bar)
+  %w[@6 @% @@a @ a].each do |n|
+    assert_raise(NameError) { o.instance_variable_set(n, 1) }
+  end
 end
 
 assert('Kernel#instance_variables', '15.3.1.3.23') do
@@ -125,6 +149,7 @@ assert('Module#class_variable_defined?', '15.2.2.4.16') do
 
   assert_true Test4ClassVariableDefined.class_variable_defined?(:@@cv)
   assert_false Test4ClassVariableDefined.class_variable_defined?(:@@noexisting)
+  assert_raise(NameError) { Test4ClassVariableDefined.class_variable_defined?("@@2") }
 end
 
 assert('Module#class_variable_get', '15.2.2.4.17') do
@@ -133,6 +158,10 @@ assert('Module#class_variable_get', '15.2.2.4.17') do
   end
 
   assert_equal 99, Test4ClassVariableGet.class_variable_get(:@@cv)
+  assert_raise(NameError) { Test4ClassVariableGet.class_variable_get(:@@a) }
+  %w[@@a? @@! @a a].each do |n|
+    assert_raise(NameError) { Test4ClassVariableGet.class_variable_get(n) }
+  end
 end
 
 assert('Module#class_variable_set', '15.2.2.4.18') do
@@ -148,6 +177,9 @@ assert('Module#class_variable_set', '15.2.2.4.18') do
   assert_true Test4ClassVariableSet.class_variables.include? :@@cv
   assert_equal 99, Test4ClassVariableSet.class_variable_get(:@@cv)
   assert_equal 101, Test4ClassVariableSet.new.foo
+  %w[@@ @@1 @@x= @x @ x 1].each do |n|
+    assert_raise(NameError) { Test4ClassVariableSet.class_variable_set(n, 1) }
+  end
 end
 
 assert('Module#class_variables', '15.2.2.4.19') do
@@ -223,6 +255,12 @@ assert('Module#remove_class_variable', '15.2.2.4.39') do
 
   assert_equal 99, Test4RemoveClassVariable.remove_class_variable(:@@cv)
   assert_false Test4RemoveClassVariable.class_variables.include? :@@cv
+  assert_raise(NameError) do
+    Test4RemoveClassVariable.remove_class_variable(:@@cv)
+  end
+  assert_raise(NameError) do
+    Test4RemoveClassVariable.remove_class_variable(:@v)
+  end
 end
 
 assert('Module#remove_method', '15.2.2.4.41') do
diff --git a/src/variable.c b/src/variable.c
index b712af261..90efe9e0e 100644
--- a/src/variable.c
+++ b/src/variable.c
@@ -437,7 +437,7 @@ mrb_iv_name_sym_p(mrb_state *mrb, mrb_sym iv_name)
   s = mrb_sym2name_len(mrb, iv_name, &len);
   if (len < 2) return FALSE;
   if (s[0] != '@') return FALSE;
-  if (s[1] == '@') return FALSE;
+  if (ISDIGIT(s[1])) return FALSE;
   return mrb_ident_p(s+1, len-1);
 }
 
diff --git a/test/t/kernel.rb b/test/t/kernel.rb
index 74176fbd0..d99358c0c 100644
--- a/test/t/kernel.rb
+++ b/test/t/kernel.rb
@@ -391,11 +391,10 @@ assert('Kernel#remove_instance_variable', '15.3.1.3.41') do
 
   tri = Test4RemoveInstanceVar.new
   assert_equal 99, tri.var
-  tri.remove
+  assert_equal 99, tri.remove
   assert_equal nil, tri.var
-  assert_raise NameError do
-    tri.remove
-  end
+  assert_raise(NameError) { tri.remove }
+  assert_raise(NameError) { tri.remove_instance_variable(:var) }
 end
 
 # Kernel#require is defined in mruby-require. '15.3.1.3.42'