[ruby/json] Sort RFC 8785 keys after as_json conversion

https://github.com/ruby/json/commit/88b358d8b2

ext/json/generator/generator.c

@@ -1051,7 +1051,8 @@ static inline long increase_depth(struct generate_json_data *data)
 static void generate_json_object(FBuffer *buffer, struct generate_json_data *data, VALUE obj)
 {
     if (RB_UNLIKELY(data->state->rfc8785)) {
-        obj = rb_proc_call_with_block(rfc8785_sort_keys_proc, 1, &obj, Qnil);
+        VALUE args[2] = {obj, data->state->strict ? data->state->as_json : Qnil};
+        obj = rb_proc_call_with_block(rfc8785_sort_keys_proc, 2, args, Qnil);
         Check_Type(obj, T_HASH);
     } else if (RB_UNLIKELY(data->state->sort_keys)) {
         obj = rb_proc_call_with_block(data->state->sort_keys, 1, &obj, Qnil);

ext/json/lib/json/ext.rb

@@ -100,7 +100,33 @@ module JSON
       end
     })
 
-    generator::State.rfc8785_sort_keys_proc = shareable_lambda(->(hash) {
+    generator::State.rfc8785_sort_keys_proc = shareable_lambda(->(hash, as_json) {
+      if as_json
+        converted = {}
+        hash.each do |key, value|
+          string = (String === key || Symbol === key) ? key.to_s : nil
+          unless string && (string.ascii_only? || (string.encoding == Encoding::UTF_8 && string.valid_encoding?))
+            key = as_json.call(key, true)
+            unless String === key || Symbol === key
+              raise GeneratorError.new("#{key.class} not allowed as object key in JSON", key)
+            end
+            begin
+              string = key.to_s.encode(Encoding::UTF_8)
+            rescue EncodingError => error
+              raise GeneratorError.new(error.message, key)
+            end
+            unless string.valid_encoding?
+              raise GeneratorError.new("source sequence is illegal/malformed utf-8", key)
+            end
+            key = string
+          end
+          if converted.key?(key)
+            raise GeneratorError.new("detected duplicate key #{key.inspect}", key)
+          end
+          converted[key] = value
+        end
+        hash = converted
+      end
       hash.sort_by { |k,| k.to_s.encode(Encoding::UTF_16) }.to_h
     })
 

test/json/json_coder_test.rb

@@ -64,6 +64,54 @@ class JSONCoderTest < Test::Unit::TestCase
     assert_equal "Integer not allowed as object key in JSON", error.message
   end
 
+  def test_json_coder_rfc8785_converted_keys
+    time = Time.utc(1999)
+    calls = []
+    coder = JSON::Coder.new(rfc8785: true) do |object, is_key|
+      calls << [object, is_key]
+      object.year.to_s
+    end
+    input = { time => 1, '1999!' => 2, Time.utc(2000) => 3 }
+    assert_equal '{"1999":1,"1999!":2,"2000":3}', coder.dump(input)
+    assert_equal [[time, true], [Time.utc(2000), true]], calls
+    assert_equal [time, '1999!', Time.utc(2000)], input.keys
+  end
+
+  def test_json_coder_rfc8785_converted_key_validation
+    object = Object.new
+    coder = JSON::Coder.new(rfc8785: true) { 42 }
+    error = assert_raise(JSON::GeneratorError) { coder.dump(object => 1) }
+    assert_equal 'Integer not allowed as object key in JSON', error.message
+    assert_equal 42, error.invalid_object
+
+    coder = JSON::Coder.new(rfc8785: true) { :a }
+    assert_equal '{"a":1,"b":2}', coder.dump('b' => 2, object => 1)
+
+    invalid = "\xFF"
+    calls = 0
+    coder = JSON::Coder.new(rfc8785: true) { calls += 1; invalid }
+    error = assert_raise(JSON::GeneratorError) { coder.dump(object => 1) }
+    assert_same invalid, error.invalid_object
+    assert_equal 1, calls
+  end
+
+  def test_json_coder_rfc8785_converted_key_collision
+    coder = JSON::Coder.new(rfc8785: true) { 'a' }
+    assert_raise(JSON::GeneratorError) { coder.dump(Object.new => 1, 'a' => 2) }
+    assert_raise(JSON::GeneratorError) { coder.dump(Object.new => 1, Object.new => 2) }
+    assert_raise(JSON::GeneratorError) { coder.dump(Object.new => 1, :a => 2) }
+  end
+
+  def test_json_coder_rfc8785_invalid_encoding_key
+    calls = []
+    coder = JSON::Coder.new(rfc8785: true) do |object, is_key|
+      calls << [object, is_key]
+      'a'
+    end
+    assert_equal '{"a":1,"b":2}', coder.dump('b' => 2, "\xFF" => 1)
+    assert_equal [["\xFF", true]], calls
+  end
+
   def test_json_coder_options
     coder = JSON::Coder.new(array_nl: "\n") do |object|
       42

test/json/json_generator_test.rb

@@ -1297,6 +1297,25 @@ class JSONGeneratorTest < Test::Unit::TestCase
     assert_rfc8785 '{"a":1,"b":2}', { a: 1, 'b' => 2 }
   end
 
+  def test_rfc8785_key_order_with_as_json
+    key = Object.new
+    def key.to_s
+      raise 'must not sort the original key'
+    end
+    as_json = ->(object, is_key) { is_key ? 'a' : object }
+    options = { rfc8785: true, strict: true, as_json: as_json }
+    object = [{ 'b' => 2, key => 1 }]
+    expected = '[{"a":1,"b":2}]'
+    assert_equal expected, JSON.generate(object, options)
+    assert_equal expected, object.to_json(options)
+    state = JSON::State.new(options)
+    io = StringIO.new
+    state.generate(object, io)
+    assert_equal expected, io.string
+
+    assert_equal '{"1":1,"b":2}', JSON.generate({ 'b' => 2, 1 => 1 }, options.merge(strict: false))
+  end
+
   def test_rfc8785_key_order_with_default_external_encoding
     # https://www.rfc-editor.org/rfc/rfc8785.html#section-3.2.3
     hash = { "\u20ac" => 1, "\r" => 2, "\ufb33" => 3, '1' => 4,