Fix JSON::Coder to cast non-string keys.

Author: byroot

ext/json/ext/generator/generator.c

@@ -1029,6 +1029,9 @@ json_object_i(VALUE key, VALUE val, VALUE _arg)
     }
 
     VALUE key_to_s;
+    bool as_json_called = false;
+
+  start:
     switch (rb_type(key)) {
         case T_STRING:
             if (RB_LIKELY(RBASIC_CLASS(key) == rb_cString)) {
@@ -1042,7 +1045,13 @@ json_object_i(VALUE key, VALUE val, VALUE _arg)
             break;
         default:
             if (data->state->strict) {
-                raise_generator_error(key, "%"PRIsVALUE" not allowed in JSON", rb_funcall(key, i_to_s, 0));
+                if (RTEST(data->state->as_json) && !as_json_called) {
+                    key = rb_proc_call_with_block(data->state->as_json, 1, &key, Qnil);
+                    as_json_called = true;
+                    goto start;
+                } else {
+                    raise_generator_error(key, "%"PRIsVALUE" not allowed as object key in JSON", CLASS_OF(key));
+                }
             }
             key_to_s = rb_convert_type(key, T_STRING, "String", "to_s");
             break;

java/src/json/ext/Generator.java

@@ -520,6 +520,23 @@ static void generateHash(ThreadContext context, Session session, RubyHash object
         buffer.write('}');
     }
 
+    private static IRubyObject castKey(ThreadContext context, IRubyObject key) {
+        RubyClass keyClass = key.getType();
+        Ruby runtime = context.runtime;
+
+        if (key instanceof RubyString) {
+            if (keyClass == runtime.getString()) {
+                return key;
+            } else {
+                return key.callMethod(context, "to_s");
+            }
+        } else if (keyClass == runtime.getSymbol()) {
+            return ((RubySymbol) key).id2name(context);
+        } else {
+            return null;
+        }
+    }
+
     private static void processEntry(ThreadContext context, Session session, OutputStream buffer, RubyHash.RubyHashEntry entry, boolean firstPair, ByteList objectNl, byte[] indent, ByteList spaceBefore, ByteList space) {
         IRubyObject key = (IRubyObject) entry.getKey();
         IRubyObject value = (IRubyObject) entry.getValue();
@@ -533,21 +550,22 @@ private static void processEntry(ThreadContext context, Session session, OutputS
 
             Ruby runtime = context.runtime;
 
-            IRubyObject keyStr;
-            RubyClass keyClass = key.getType();
-            if (key instanceof RubyString) {
-                if (keyClass == runtime.getString()) {
-                    keyStr = key;
-                } else {
-                    keyStr = key.callMethod(context, "to_s");
+            IRubyObject keyStr = castKey(context, key);
+            if (keyStr == null || !(keyStr instanceof RubyString)) {
+                GeneratorState state = session.getState(context);
+                if (state.strict()) {
+                    if (state.getAsJSON() != null) {
+                        key = state.getAsJSON().call(context, key);
+                        keyStr = castKey(context, key);
+                    }
+
+                    if (keyStr == null) {
+                        throw Utils.buildGeneratorError(context, key, key.getType().name(context) + " not allowed as object key in JSON").toThrowable();
+                    }
                 }
-            } else if (keyClass == runtime.getSymbol()) {
-                keyStr = ((RubySymbol) key).id2name(context);
-            } else {
-                if (session.getState(context).strict()) {
-                    throw Utils.buildGeneratorError(context, key, key + " not allowed in JSON").toThrowable();
+                else {
+                    keyStr = TypeConverter.convertToType(key, runtime.getString(), "to_s");
                 }
-                keyStr = TypeConverter.convertToType(key, runtime.getString(), "to_s");
             }
 
             if (keyStr.getMetaClass() == runtime.getString()) {
@@ -673,7 +691,7 @@ void generate(ThreadContext context, Session session, IRubyObject object, Output
     static RubyString generateGenericNew(ThreadContext context, Session session, IRubyObject object) {
         GeneratorState state = session.getState(context);
         if (state.strict()) {
-            if (state.getAsJSON() != null ) {
+            if (state.getAsJSON() != null) {
                 IRubyObject value = state.getAsJSON().call(context, object);
                 Handler handler = getHandlerFor(context.runtime, value);
                 if (handler == GENERIC_HANDLER) {

lib/json/truffle_ruby/generator.rb

@@ -477,7 +477,13 @@ def json_transform(state)
               result << state.indent * depth if indent
 
               if state.strict? && !(Symbol === key || String === key)
-                raise GeneratorError.new("#{key.class} not allowed in JSON", value)
+                if state.as_json
+                  key = state.as_json.call(key)
+                end
+
+                unless Symbol === key || String === key
+                  raise GeneratorError.new("#{key.class} not allowed as object key in JSON", value)
+                end
               end
 
               key_str = key.to_s

test/json/json_coder_test.rb

@@ -18,6 +18,18 @@ def test_json_coder_with_proc_with_unsupported_value
     assert_raise(JSON::GeneratorError) { coder.dump([Object.new]) }
   end
 
+  def test_json_coder_hash_key
+    obj = Object.new
+    coder = JSON::Coder.new(&:to_s)
+    assert_equal %({#{obj.to_s.inspect}:1}), coder.dump({ obj => 1 })
+
+    coder = JSON::Coder.new { 42 }
+    error = assert_raise JSON::GeneratorError do
+      coder.dump({ obj => 1 })
+    end
+    assert_equal "Integer not allowed as object key in JSON", error.message
+  end
+
   def test_json_coder_options
     coder = JSON::Coder.new(array_nl: "\n") do |object|
       42