Handle YAML.load_file on Psych 4

.github/workflows/ci.yaml

@@ -13,6 +13,7 @@ on:
 jobs:
   platforms:
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu, macos, windows]
         ruby: ['2.5']
@@ -40,6 +41,7 @@ jobs:
 
   psych4:
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu]
         ruby: ['3.0']
@@ -57,6 +59,7 @@ jobs:
 
   minimal:
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu]
         ruby: ['jruby', 'truffleruby']

CHANGELOG.md

@@ -1,5 +1,13 @@
 # Unreleased
 
+* Improve support of Psych 4. (#392)
+  Since `1.8.0`, `YAML.load_file` was no longer cached when Psych 4 was used. This is because `load_file` loads
+  in safe mode by default, so the Bootsnap cache could defeat that safety.
+  Now when precompiling YAML files, Bootsnap first try to parse them in safe mode, and if it can't fallback to unsafe mode,
+  and the cache contains a flag that records wether it was generated in safe mode or not.
+  `YAML.unsafe_load_file` will use safe caches just fine, but `YAML.load_file` will fallback to uncached YAML parsing
+  if the cache was generated using unsafe parsing.
+
 * Minimize the Kernel.require extra stack frames. (#393)
   This should reduce the noise generated by bootsnap on `LoadError`.
 

ext/bootsnap/bootsnap.c

@@ -75,7 +75,7 @@ struct bs_cache_key {
 STATIC_ASSERT(sizeof(struct bs_cache_key) == KEY_SIZE);
 
 /* Effectively a schema version. Bumping invalidates all previous caches */
-static const uint32_t current_version = 3;
+static const uint32_t current_version = 4;
 
 /* hash of e.g. "x86_64-darwin17", invalidating when ruby is recompiled on a
  * new OS ABI, etc. */
@@ -426,6 +426,7 @@ open_current_file(char * path, struct bs_cache_key * key, const char ** errno_pr
 #define ERROR_WITH_ERRNO -1
 #define CACHE_MISS -2
 #define CACHE_STALE -3
+#define CACHE_UNCOMPILABLE -4
 
 /*
  * Read the cache key from the given fd, which must have position 0 (e.g.
@@ -507,14 +508,14 @@ fetch_cached_data(int fd, ssize_t data_size, VALUE handler, VALUE args, VALUE *
   if (data_size > 100000000000) {
     *errno_provenance = "bs_fetch:fetch_cached_data:datasize";
     errno = EINVAL; /* because wtf? */
-    ret = -1;
+    ret = ERROR_WITH_ERRNO;
     goto done;
   }
   data = ALLOC_N(char, data_size);
   nread = read(fd, data, data_size);
   if (nread < 0) {
     *errno_provenance = "bs_fetch:fetch_cached_data:read";
-    ret = -1;
+    ret = ERROR_WITH_ERRNO;
     goto done;
   }
   if (nread != data_size) {
@@ -525,6 +526,10 @@ fetch_cached_data(int fd, ssize_t data_size, VALUE handler, VALUE args, VALUE *
   storage_data = rb_str_new(data, data_size);
 
   *exception_tag = bs_storage_to_output(handler, args, storage_data, output_data);
+  if (*output_data == uncompilable) {
+    ret = CACHE_UNCOMPILABLE;
+    goto done;
+  }
   ret = 0;
 done:
   if (data != NULL) xfree(data);
@@ -737,7 +742,15 @@ bs_fetch(char * path, VALUE path_v, char * cache_path, VALUE handler, VALUE args
       &output_data, &exception_tag, &errno_provenance
     );
     if (exception_tag != 0) goto raise;
-    else if (res == CACHE_MISS || res == CACHE_STALE) valid_cache = 0;
+    else if (res == CACHE_UNCOMPILABLE) {
+      /* If fetch_cached_data returned `Uncompilable` we fallback to `input_to_output`
+        This happens if we have say, an unsafe YAML cache, but try to load it in safe mode */
+      if (bs_read_contents(current_fd, current_key.size, &contents, &errno_provenance) < 0) goto fail_errno;
+      input_data = rb_str_new(contents, current_key.size);
+      bs_input_to_output(handler, args, input_data, &output_data, &exception_tag);
+      if (exception_tag != 0) goto raise;
+      goto succeed;
+    } else if (res == CACHE_MISS || res == CACHE_STALE) valid_cache = 0;
     else if (res == ERROR_WITH_ERRNO) goto fail_errno;
     else if (!NIL_P(output_data)) goto succeed; /* fast-path, goal */
   }
@@ -772,9 +785,13 @@ bs_fetch(char * path, VALUE path_v, char * cache_path, VALUE handler, VALUE args
   exception_tag = bs_storage_to_output(handler, args, storage_data, &output_data);
   if (exception_tag != 0) goto raise;
 
-  /* If output_data is nil, delete the cache entry and generate the output
-   * using input_to_output */
-  if (NIL_P(output_data)) {
+  if (output_data == uncompilable) {
+    /* If storage_to_output returned `Uncompilable` we fallback to `input_to_output` */
+    bs_input_to_output(handler, args, input_data, &output_data, &exception_tag);
+    if (exception_tag != 0) goto raise;
+  } else if (NIL_P(output_data)) {
+    /* If output_data is nil, delete the cache entry and generate the output
+     * using input_to_output */
     if (unlink(cache_path) < 0) {
       errno_provenance = "bs_fetch:unlink";
       goto fail_errno;
@@ -919,12 +936,27 @@ struct i2s_data {
 };
 
 static VALUE
-prot_storage_to_output(VALUE arg)
+rescue_uncompilable(VALUE arg, VALUE e)
+{
+  return uncompilable;
+}
+
+static VALUE
+try_storage_to_output(VALUE arg)
 {
   struct s2o_data * data = (struct s2o_data *)arg;
   return rb_funcall(data->handler, rb_intern("storage_to_output"), 2, data->storage_data, data->args);
 }
 
+static VALUE
+prot_storage_to_output(VALUE arg)
+{
+  return rb_rescue2(
+    try_storage_to_output, arg,
+    rescue_uncompilable, Qnil,
+    rb_eBootsnap_CompileCache_Uncompilable, 0);
+}
+
 static int
 bs_storage_to_output(VALUE handler, VALUE args, VALUE storage_data, VALUE * output_data)
 {
@@ -963,19 +995,13 @@ try_input_to_storage(VALUE arg)
   return rb_funcall(data->handler, rb_intern("input_to_storage"), 2, data->input_data, data->pathval);
 }
 
-static VALUE
-rescue_input_to_storage(VALUE arg, VALUE e)
-{
-  return uncompilable;
-}
-
 static VALUE
 prot_input_to_storage(VALUE arg)
 {
   struct i2s_data * data = (struct i2s_data *)arg;
   return rb_rescue2(
       try_input_to_storage, (VALUE)data,
-      rescue_input_to_storage, Qnil,
+      rescue_uncompilable, Qnil,
       rb_eBootsnap_CompileCache_Uncompilable, 0);
 }
 

lib/bootsnap/cli.rb

@@ -138,7 +138,7 @@ def precompile_yaml_files(load_paths, exclude: self.exclude)
 
     def precompile_yaml(*yaml_files)
       Array(yaml_files).each do |yaml_file|
-        if CompileCache::YAML.precompile(yaml_file, cache_dir: cache_dir)
+        if CompileCache::YAML.precompile(yaml_file)
           STDERR.puts(yaml_file) if verbose
         end
       end
@@ -161,7 +161,7 @@ def precompile_json_files(load_paths, exclude: self.exclude)
 
     def precompile_json(*json_files)
       Array(json_files).each do |json_file|
-        if CompileCache::JSON.precompile(json_file, cache_dir: cache_dir)
+        if CompileCache::JSON.precompile(json_file)
           STDERR.puts(json_file) if verbose
         end
       end
@@ -183,7 +183,7 @@ def precompile_ruby_files(load_paths, exclude: self.exclude)
 
     def precompile_ruby(*ruby_files)
       Array(ruby_files).each do |ruby_file|
-        if CompileCache::ISeq.precompile(ruby_file, cache_dir: cache_dir)
+        if CompileCache::ISeq.precompile(ruby_file)
           STDERR.puts(ruby_file) if verbose
         end
       end

lib/bootsnap/compile_cache/iseq.rb

@@ -7,7 +7,11 @@ module Bootsnap
   module CompileCache
     module ISeq
       class << self
-        attr_accessor(:cache_dir)
+        attr_reader(:cache_dir)
+
+        def cache_dir=(cache_dir)
+          @cache_dir = cache_dir.end_with?("/") ? "#{cache_dir}iseq" : "#{cache_dir}-iseq"
+        end
       end
 
       has_ruby_bug_18250 = begin # https://bugs.ruby-lang.org/issues/18250
@@ -61,7 +65,7 @@ def self.fetch(path, cache_dir: ISeq.cache_dir)
         )
       end
 
-      def self.precompile(path, cache_dir: ISeq.cache_dir)
+      def self.precompile(path)
         Bootsnap::CompileCache::Native.precompile(
           cache_dir,
           path.to_s,

lib/bootsnap/compile_cache/json.rb

@@ -6,7 +6,12 @@ module Bootsnap
   module CompileCache
     module JSON
       class << self
-        attr_accessor(:msgpack_factory, :cache_dir, :supported_options)
+        attr_accessor(:msgpack_factory, :supported_options)
+        attr_reader(:cache_dir)
+
+        def cache_dir=(cache_dir)
+          @cache_dir = cache_dir.end_with?("/") ? "#{cache_dir}json" : "#{cache_dir}-json"
+        end
 
         def input_to_storage(payload, _)
           obj = ::JSON.parse(payload)
@@ -24,7 +29,7 @@ def input_to_output(data, kwargs)
           ::JSON.parse(data, **(kwargs || {}))
         end
 
-        def precompile(path, cache_dir: self.cache_dir)
+        def precompile(path)
           Bootsnap::CompileCache::Native.precompile(
             cache_dir,
             path.to_s,