Merge pull request #668 from casperisfine/non-string-keys

JSON.generate: call to_json on String subclasses

Author: byroot

ext/json/ext/fbuffer/fbuffer.h

@@ -42,6 +42,10 @@ static VALUE fbuffer_to_s(FBuffer *fb);
 #define RB_UNLIKELY(expr) expr
 #endif
 
+#ifndef RB_LIKELY
+#define RB_LIKELY(expr) expr
+#endif
+
 static void fbuffer_stack_init(FBuffer *fb, unsigned long initial_length, char *stack_buffer, long stack_buffer_size)
 {
     fb->initial_length = (initial_length > 0) ? initial_length : FBUFFER_INITIAL_LENGTH_DEFAULT;

ext/json/ext/generator/generator.c

@@ -737,7 +737,11 @@ json_object_i(VALUE key, VALUE val, VALUE _arg)
             break;
     }
 
-    generate_json_string(buffer, data, state, key_to_s);
+    if (RB_LIKELY(RBASIC_CLASS(key_to_s) == rb_cString)) {
+        generate_json_string(buffer, data, state, key_to_s);
+    } else {
+        generate_json(buffer, data, state, key_to_s);
+    }
     if (RB_UNLIKELY(state->space_before)) fbuffer_append_str(buffer, state->space_before);
     fbuffer_append_char(buffer, ':');
     if (RB_UNLIKELY(state->space)) fbuffer_append_str(buffer, state->space);

java/src/json/ext/Generator.java

@@ -359,7 +359,14 @@ public void visit(IRubyObject key, IRubyObject value) {
                         }
                         if (objectNl.length() != 0) buffer.append(indent);
 
-                        STRING_HANDLER.generate(session, key.asString(), buffer);
+                        IRubyObject keyStr = key.callMethod(context, "to_s");
+                        if (keyStr.getMetaClass() == runtime.getString()) {
+                            STRING_HANDLER.generate(session, (RubyString)keyStr, buffer);
+                        } else {
+                            Utils.ensureString(keyStr);
+                            Handler<IRubyObject> keyHandler = (Handler<IRubyObject>) getHandlerFor(runtime, keyStr);
+                            keyHandler.generate(session, keyStr, buffer);
+                        }
                         session.infectBy(key);
 
                         buffer.append(spaceBefore);

lib/json/pure/generator.rb

@@ -307,19 +307,30 @@ def generate(obj)
 
         # Handles @allow_nan, @buffer_initial_length, other ivars must be the default value (see above)
         private def generate_json(obj, buf)
-          case obj
-          when Hash
+          klass = obj.class
+          if klass == Hash
             buf << '{'
             first = true
             obj.each_pair do |k,v|
               buf << ',' unless first
-              fast_serialize_string(k.to_s, buf)
+
+              key_str = k.to_s
+              if key_str.is_a?(::String)
+                if key_str.class == ::String
+                  fast_serialize_string(key_str, buf)
+                else
+                  generate_json(key_str, buf)
+                end
+              else
+                raise TypeError, "#{k.class}#to_s returns an instance of #{key_str.class}, expected a String"
+              end
+
               buf << ':'
               generate_json(v, buf)
               first = false
             end
             buf << '}'
-          when Array
+          elsif klass == Array
             buf << '['
             first = true
             obj.each do |e|
@@ -328,9 +339,9 @@ def generate(obj)
               first = false
             end
             buf << ']'
-          when String
+          elsif klass == String
             fast_serialize_string(obj, buf)
-          when Integer
+          elsif klass == Integer
             buf << obj.to_s
           else
             # Note: Float is handled this way since Float#to_s is slow anyway
@@ -419,7 +430,15 @@ def json_transform(state)
             each { |key, value|
               result << delim unless first
               result << state.indent * depth if indent
-              result = +"#{result}#{key.to_s.to_json(state)}#{state.space_before}:#{state.space}"
+
+              key_str = key.to_s
+              key_json = if key_str.is_a?(::String)
+                key_str = key_str.to_json(state)
+              else
+                raise TypeError, "#{key.class}#to_s returns an instance of #{key_str.class}, expected a String"
+              end
+
+              result = +"#{result}#{key_json}#{state.space_before}:#{state.space}"
               if state.strict? && !(false == value || true == value || nil == value || String === value || Array === value || Hash === value || Integer === value || Float === value)
                 raise GeneratorError, "#{value.class} not allowed in JSON"
               elsif value.respond_to?(:to_json)

test/json/json_generator_test.rb

@@ -486,6 +486,41 @@ def test_invalid_encoding_string
     end
   end
 
+  class MyCustomString < String
+    def to_json(_state = nil)
+      '"my_custom_key"'
+    end
+
+    def to_s
+      self
+    end
+  end
+
+  def test_string_subclass_as_keys
+    # Ref: https://github.com/ruby/json/issues/667
+    # if key.to_s doesn't return a bare string, we call `to_json` on it.
+    key = MyCustomString.new("won't be used")
+    assert_equal '{"my_custom_key":1}', JSON.generate(key => 1)
+  end
+
+  class FakeString
+    def to_json(_state = nil)
+      raise "Shouldn't be called"
+    end
+
+    def to_s
+      self
+    end
+  end
+
+  def test_custom_object_as_keys
+    key = FakeString.new
+    error = assert_raise(TypeError) do
+      JSON.generate(key => 1)
+    end
+    assert_match "FakeString", error.message
+  end
+
   def test_to_json_called_with_state_object
     object = Object.new
     called = false