Fix JSON::Coder to call as_json proc for NaN and Infinity

Co-authored-by: Jean Boussier <jean.boussier@gmail.com>

Author: etiennebarrie

ext/json/ext/generator/generator.c

@@ -841,15 +841,19 @@ json_object_i(VALUE key, VALUE val, VALUE _arg)
     return ST_CONTINUE;
 }
 
-static void generate_json_object(FBuffer *buffer, struct generate_json_data *data, JSON_Generator_State *state, VALUE obj)
+static inline long increase_depth(JSON_Generator_State *state)
 {
-    long max_nesting = state->max_nesting;
     long depth = ++state->depth;
-    int j;
-
-    if (max_nesting != 0 && depth > max_nesting) {
+    if (RB_UNLIKELY(depth > state->max_nesting && state->max_nesting)) {
         rb_raise(eNestingError, "nesting of %ld is too deep", --state->depth);
     }
+    return depth;
+}
+
+static void generate_json_object(FBuffer *buffer, struct generate_json_data *data, JSON_Generator_State *state, VALUE obj)
+{
+    int j;
+    long depth = increase_depth(state);
 
     if (RHASH_SIZE(obj) == 0) {
         fbuffer_append(buffer, "{}", 2);
@@ -879,12 +883,8 @@ static void generate_json_object(FBuffer *buffer, struct generate_json_data *dat
 
 static void generate_json_array(FBuffer *buffer, struct generate_json_data *data, JSON_Generator_State *state, VALUE obj)
 {
-    long max_nesting = state->max_nesting;
-    long depth = ++state->depth;
     int i, j;
-    if (max_nesting != 0 && depth > max_nesting) {
-        rb_raise(eNestingError, "nesting of %ld is too deep", --state->depth);
-    }
+    long depth = increase_depth(state);
 
     if (RARRAY_LEN(obj) == 0) {
         fbuffer_append(buffer, "[]", 2);
@@ -1031,13 +1031,21 @@ static void generate_json_float(FBuffer *buffer, struct generate_json_data *data
 {
     double value = RFLOAT_VALUE(obj);
     char allow_nan = state->allow_nan;
-    VALUE tmp = rb_funcall(obj, i_to_s, 0);
     if (!allow_nan) {
         if (isinf(value) || isnan(value)) {
-            raise_generator_error(obj, "%"PRIsVALUE" not allowed in JSON", tmp);
+            if (state->strict && state->as_json) {
+                VALUE casted_obj = rb_proc_call_with_block(state->as_json, 1, &obj, Qnil);
+                if (casted_obj != obj) {
+                    increase_depth(state);
+                    generate_json(buffer, data, state, casted_obj);
+                    state->depth--;
+                    return;
+                }
+            }
+            raise_generator_error(obj, "%"PRIsVALUE" not allowed in JSON", rb_funcall(obj, i_to_s, 0));
         }
     }
-    fbuffer_append_str(buffer, tmp);
+    fbuffer_append_str(buffer, rb_funcall(obj, i_to_s, 0));
 }
 
 static void generate_json_fragment(FBuffer *buffer, struct generate_json_data *data, JSON_Generator_State *state, VALUE obj)

java/src/json/ext/Generator.java

@@ -261,7 +261,17 @@ void generate(ThreadContext context, Session session, RubyFloat object, OutputSt
                 double value = object.getValue();
 
                 if (Double.isInfinite(value) || Double.isNaN(value)) {
-                    if (!session.getState(context).allowNaN()) {
+                    GeneratorState state = session.getState(context);
+
+                    if (!state.allowNaN()) {
+                        if (state.strict() && state.getAsJSON() != null) {
+                            IRubyObject castedValue = state.getAsJSON().call(context, object);
+                            if (castedValue != object) {
+                                getHandlerFor(context.runtime, castedValue).generate(context, session, castedValue, buffer);
+                                return;
+                            }
+                        }
+
                         throw Utils.buildGeneratorError(context, object, object + " not allowed in JSON").toThrowable();
                     }
                 }

lib/json/truffle_ruby/generator.rb

@@ -570,18 +570,23 @@ def to_json(*) to_s end
 
         module Float
           # Returns a JSON string representation for this Float number.
-          def to_json(state = nil, *)
+          def to_json(state = nil, *args)
             state = State.from_state(state)
-            case
-            when infinite?
-              if state.allow_nan?
-                to_s
-              else
-                raise GeneratorError.new("#{self} not allowed in JSON", self)
-              end
-            when nan?
+            if infinite? || nan?
               if state.allow_nan?
                 to_s
+              elsif state.strict? && state.as_json
+                casted_value = state.as_json.call(self)
+
+                if casted_value.equal?(self)
+                  raise GeneratorError.new("#{self} not allowed in JSON", self)
+                end
+
+                state.check_max_nesting
+                state.depth += 1
+                result = casted_value.to_json(state, *args)
+                state.depth -= 1
+                result
               else
                 raise GeneratorError.new("#{self} not allowed in JSON", self)
               end

test/json/json_coder_test.rb

@@ -35,4 +35,19 @@ def test_json_coder_load_options
     coder = JSON::Coder.new(symbolize_names: true)
     assert_equal({a: 1}, coder.load('{"a":1}'))
   end
+
+  def test_json_coder_dump_NaN_or_Infinity
+    coder = JSON::Coder.new(&:inspect)
+    assert_equal "NaN", coder.load(coder.dump(Float::NAN))
+    assert_equal "Infinity", coder.load(coder.dump(Float::INFINITY))
+    assert_equal "-Infinity", coder.load(coder.dump(-Float::INFINITY))
+  end
+
+  def test_json_coder_dump_NaN_or_Infinity_loop
+    coder = JSON::Coder.new(&:itself)
+    error = assert_raise JSON::GeneratorError do
+      coder.dump(Float::NAN)
+    end
+    assert_include error.message, "NaN not allowed in JSON"
+  end
 end