Fix Psych 4 compatibility and add caching for YAML.unsafe_load_file

.github/workflows/ci.yaml

@@ -36,6 +36,23 @@ jobs:
       - run: bundle install
       - run: bundle exec rake
 
+  psych4:
+    strategy:
+      matrix:
+        os: [ubuntu]
+        ruby: ['3.0']
+    runs-on: ${{ matrix.os }}-latest
+    env:
+      PSYCH_4: "1"
+    steps:
+      - uses: actions/checkout@v2
+      - uses: ruby/setup-ruby@v1
+        with:
+          ruby-version: ${{ matrix.ruby }}
+      - run: bundle install
+      - run: bundle exec rake
+
+
   minimal:
     strategy:
       matrix:

CHANGELOG.md

@@ -1,5 +1,7 @@
 # Unreleased
 
+* Improve support for Pysch 4.
+
 # 1.7.7
 
 * Fix `require_relative` in evaled code on latest ruby 3.1.0-dev. (#366)

Gemfile

@@ -4,6 +4,10 @@ source 'https://rubygems.org'
 # Specify your gem's dependencies in bootsnap.gemspec
 gemspec
 
+if ENV["PSYCH_4"]
+  gem "psych", ">= 4"
+end
+
 group :development do
   gem 'rubocop'
   gem 'byebug', platform: :ruby

lib/bootsnap/compile_cache/yaml.rb

@@ -23,7 +23,7 @@ def storage_to_output(data, kwargs)
         end
 
         def input_to_output(data, kwargs)
-          ::YAML.load(data, **(kwargs || {}))
+          ::YAML.unsafe_load(data, **(kwargs || {}))
         end
 
         def strict_load(payload, *args)
@@ -52,6 +52,13 @@ def init!
           require('msgpack')
           require('date')
 
+          if Patch.method_defined?(:unsafe_load_file) && !::YAML.respond_to?(:unsafe_load_file)
+            Patch.send(:remove_method, :unsafe_load_file)
+          end
+          if Patch.method_defined?(:load_file) && ::YAML::VERSION >= '4'
+            Patch.send(:remove_method, :load_file)
+          end
+
           # MessagePack serializes symbols as strings by default.
           # We want them to roundtrip cleanly, so we use a custom factory.
           # see: https://github.com/msgpack/msgpack-ruby/pull/122
@@ -126,6 +133,27 @@ def load_file(path, *args)
         end
 
         ruby2_keywords :load_file if respond_to?(:ruby2_keywords, true)
+
+        def unsafe_load_file(path, *args)
+          return super if args.size > 1
+          if kwargs = args.first
+            return super unless kwargs.is_a?(Hash)
+            return super unless (kwargs.keys - ::Bootsnap::CompileCache::YAML.supported_options).empty?
+          end
+
+          begin
+            ::Bootsnap::CompileCache::Native.fetch(
+              Bootsnap::CompileCache::YAML.cache_dir,
+              File.realpath(path),
+              ::Bootsnap::CompileCache::YAML,
+              kwargs,
+            )
+          rescue Errno::EACCES
+            ::Bootsnap::CompileCache.permission_error(path)
+          end
+        end
+
+        ruby2_keywords :unsafe_load_file if respond_to?(:ruby2_keywords, true)
       end
     end
   end

test/compile_cache/yaml_test.rb

@@ -6,17 +6,21 @@ class CompileCacheYAMLTest < Minitest::Test
 
   module FakeYaml
     Fallback = Class.new(StandardError)
-    extend self
-    singleton_class.prepend(Bootsnap::CompileCache::YAML::Patch)
+    class << self
+      def load_file(path, symbolize_names: false, freeze: false, fallback: nil)
+        raise Fallback
+      end
 
-    def load_file(path, symbolize_names: false, freeze: false, fallback: nil)
-      raise Fallback
+      def unsafe_load_file(path, symbolize_names: false, freeze: false, fallback: nil)
+        raise Fallback
+      end
     end
   end
 
   def setup
     super
     Bootsnap::CompileCache::YAML.init!
+    FakeYaml.singleton_class.prepend(Bootsnap::CompileCache::YAML::Patch)
   end
 
   def test_yaml_strict_load
@@ -44,49 +48,72 @@ def test_yaml_tags
     assert_equal "YAML tags are not supported: !ruby/object", error.message
   end
 
-  def test_load_file
-    Help.set_file('a.yml', "---\nfoo: bar", 100)
-    assert_equal({'foo' => 'bar'}, FakeYaml.load_file('a.yml'))
-  end
+  if YAML::VERSION >= '4'
+    def test_load_psych_4
+      # Until we figure out a proper strategy, only `YAML.unsafe_load_file`
+      # is cached with Psych >= 4
+      Help.set_file('a.yml', "foo: &foo\n  bar: 42\nplop:\n  <<: *foo", 100)
+      assert_raises FakeYaml::Fallback do
+        FakeYaml.load_file('a.yml')
+      end
+    end
+  else
+    def test_load_file
+      Help.set_file('a.yml', "---\nfoo: bar", 100)
+      assert_equal({'foo' => 'bar'}, FakeYaml.load_file('a.yml'))
+    end
 
-  def test_load_file_symbolize_names
-    Help.set_file('a.yml', "---\nfoo: bar", 100)
-    FakeYaml.load_file('a.yml')
+    def test_load_file_aliases
+      Help.set_file('a.yml', "foo: &foo\n  bar: 42\nplop:\n  <<: *foo", 100)
+      assert_equal({"foo" => { "bar" => 42 }, "plop" => { "bar" => 42} }, FakeYaml.load_file('a.yml'))
+    end
 
-    if ::Bootsnap::CompileCache::YAML.supported_options.include?(:symbolize_names)
-      2.times do
-        assert_equal({foo: 'bar'}, FakeYaml.load_file('a.yml', symbolize_names: true))
-      end
-    else
-      assert_raises(FakeYaml::Fallback) do # would call super
-        FakeYaml.load_file('a.yml', symbolize_names: true)
+    def test_load_file_symbolize_names
+      Help.set_file('a.yml', "---\nfoo: bar", 100)
+      FakeYaml.load_file('a.yml')
+
+      if ::Bootsnap::CompileCache::YAML.supported_options.include?(:symbolize_names)
+        2.times do
+          assert_equal({foo: 'bar'}, FakeYaml.load_file('a.yml', symbolize_names: true))
+        end
+      else
+        assert_raises(FakeYaml::Fallback) do # would call super
+          FakeYaml.load_file('a.yml', symbolize_names: true)
+        end
       end
     end
-  end
 
-  def test_load_file_freeze
-    Help.set_file('a.yml', "---\nfoo", 100)
-    FakeYaml.load_file('a.yml')
+    def test_load_file_freeze
+      Help.set_file('a.yml', "---\nfoo", 100)
+      FakeYaml.load_file('a.yml')
 
-    if ::Bootsnap::CompileCache::YAML.supported_options.include?(:freeze)
-      2.times do
-        string = FakeYaml.load_file('a.yml', freeze: true)
-        assert_equal("foo", string)
-        assert_predicate(string, :frozen?)
+      if ::Bootsnap::CompileCache::YAML.supported_options.include?(:freeze)
+        2.times do
+          string = FakeYaml.load_file('a.yml', freeze: true)
+          assert_equal("foo", string)
+          assert_predicate(string, :frozen?)
+        end
+      else
+        assert_raises(FakeYaml::Fallback) do # would call super
+          FakeYaml.load_file('a.yml', freeze: true)
+        end
       end
-    else
+    end
+
+    def test_load_file_unknown_option
+      Help.set_file('a.yml', "---\nfoo", 100)
+      FakeYaml.load_file('a.yml')
+
       assert_raises(FakeYaml::Fallback) do # would call super
-        FakeYaml.load_file('a.yml', freeze: true)
+        FakeYaml.load_file('a.yml', fallback: true)
       end
     end
   end
 
-  def test_load_file_unknown_option
-    Help.set_file('a.yml', "---\nfoo", 100)
-    FakeYaml.load_file('a.yml')
-
-    assert_raises(FakeYaml::Fallback) do # would call super
-      FakeYaml.load_file('a.yml', fallback: true)
+  if YAML.respond_to?(:unsafe_load_file)
+    def test_unsafe_load_file
+      Help.set_file('a.yml', "foo: &foo\n  bar: 42\nplop:\n  <<: *foo", 100)
+      assert_equal({"foo" => { "bar" => 42 }, "plop" => { "bar" => 42} }, FakeYaml.unsafe_load_file('a.yml'))
     end
   end
 end