Fix use-after-free of the CDict/DDict referenced by a stream

ext/zstdruby/common.h

@@ -24,7 +24,11 @@ static int convert_compression_level(ZSTD_CCtx* ctx, VALUE compression_level_val
   return NUM2INT(compression_level_value);
 }
 
-static void set_compress_params(ZSTD_CCtx* const ctx, VALUE kwargs)
+/* Returns the Zstd::CDict given as `dict:`, or Qnil. ZSTD_CCtx_refCDict only
+   borrows the pointer, so a caller that keeps the ZSTD_CCtx alive beyond this
+   call has to keep the returned object reachable for just as long. A String
+   dictionary needs no such handling: ZSTD_CCtx_loadDictionary copies it. */
+static VALUE set_compress_params(ZSTD_CCtx* const ctx, VALUE kwargs)
 {
   ID kwargs_keys[2];
   kwargs_keys[0] = rb_intern("level");
@@ -46,6 +50,7 @@ static void set_compress_params(ZSTD_CCtx* const ctx, VALUE kwargs)
         ZSTD_freeCCtx(ctx);
         rb_raise(rb_eRuntimeError, "%s", "ZSTD_CCtx_refCDict failed");
       }
+      return kwargs_values[1];
     } else if (TYPE(kwargs_values[1]) == T_STRING) {
       char* dict_buffer = RSTRING_PTR(kwargs_values[1]);
       size_t dict_size = RSTRING_LEN(kwargs_values[1]);
@@ -59,6 +64,7 @@ static void set_compress_params(ZSTD_CCtx* const ctx, VALUE kwargs)
       rb_raise(rb_eArgError, "`dict:` must be a Zstd::CDict or a String");
     }
   }
+  return Qnil;
 }
 
 struct stream_compress_params {
@@ -122,7 +128,9 @@ static size_t zstd_compress(ZSTD_CCtx* const ctx, char* output_data, size_t outp
 #endif
 }
 
-static void set_decompress_params(ZSTD_DCtx* const dctx, VALUE kwargs)
+/* Returns the Zstd::DDict given as `dict:`, or Qnil. See set_compress_params:
+   ZSTD_DCtx_refDDict borrows, ZSTD_DCtx_loadDictionary copies. */
+static VALUE set_decompress_params(ZSTD_DCtx* const dctx, VALUE kwargs)
 {
   ID kwargs_keys[1];
   kwargs_keys[0] = rb_intern("dict");
@@ -137,6 +145,7 @@ static void set_decompress_params(ZSTD_DCtx* const dctx, VALUE kwargs)
         ZSTD_freeDCtx(dctx);
         rb_raise(rb_eRuntimeError, "%s", "ZSTD_DCtx_refDDict failed");
       }
+      return kwargs_values[0];
     } else if (TYPE(kwargs_values[0]) == T_STRING) {
       char* dict_buffer = RSTRING_PTR(kwargs_values[0]);
       size_t dict_size = RSTRING_LEN(kwargs_values[0]);
@@ -150,6 +159,7 @@ static void set_decompress_params(ZSTD_DCtx* const dctx, VALUE kwargs)
       rb_raise(rb_eArgError, "`dict:` must be a Zstd::DDict or a String");
     }
   }
+  return Qnil;
 }
 
 struct stream_decompress_params {

ext/zstdruby/streaming_compress.c

@@ -5,6 +5,7 @@ struct streaming_compress_t {
   VALUE buf;
   size_t buf_size;
   VALUE pending;   /* accumulate compressed bytes produced by write() */
+  VALUE dict;      /* Zstd::CDict the ctx borrows a pointer into, or Qnil */
 };
 
 static void
@@ -14,9 +15,11 @@ streaming_compress_mark(void *p)
 #ifdef HAVE_RB_GC_MARK_MOVABLE
   rb_gc_mark_movable(sc->buf);
   rb_gc_mark_movable(sc->pending);
+  rb_gc_mark_movable(sc->dict);
 #else
   rb_gc_mark(sc->buf);
   rb_gc_mark(sc->pending);
+  rb_gc_mark(sc->dict);
 #endif
 }
 
@@ -44,6 +47,7 @@ streaming_compress_compact(void *p)
   struct streaming_compress_t *sc = p;
   sc->buf = rb_gc_location(sc->buf);
   sc->pending = rb_gc_location(sc->pending);
+  sc->dict = rb_gc_location(sc->dict);
 }
 #endif
 
@@ -69,6 +73,7 @@ rb_streaming_compress_allocate(VALUE klass)
   RB_OBJ_WRITE(obj, &sc->buf, Qnil);
   sc->buf_size = 0;
   RB_OBJ_WRITE(obj, &sc->pending, Qnil);
+  RB_OBJ_WRITE(obj, &sc->dict, Qnil);
   return obj;
 }
 
@@ -86,9 +91,10 @@ rb_streaming_compress_initialize(int argc, VALUE *argv, VALUE obj)
   if (ctx == NULL) {
     rb_raise(rb_eRuntimeError, "%s", "ZSTD_createCCtx error");
   }
-  set_compress_params(ctx, kwargs);
+  VALUE dict = set_compress_params(ctx, kwargs);
 
   sc->ctx = ctx;
+  RB_OBJ_WRITE(obj, &sc->dict, dict);
   RB_OBJ_WRITE(obj, &sc->buf, rb_str_new(NULL, buffOutSize));
   sc->buf_size = buffOutSize;
   RB_OBJ_WRITE(obj, &sc->pending, rb_str_new(0, 0));

ext/zstdruby/streaming_decompress.c

@@ -4,6 +4,7 @@ struct streaming_decompress_t {
   ZSTD_DCtx* dctx;
   VALUE buf;
   size_t buf_size;
+  VALUE dict;      /* Zstd::DDict the dctx borrows a pointer into, or Qnil */
 };
 
 static void
@@ -12,8 +13,10 @@ streaming_decompress_mark(void *p)
   struct streaming_decompress_t *sd = p;
 #ifdef HAVE_RB_GC_MARK_MOVABLE
   rb_gc_mark_movable(sd->buf);
+  rb_gc_mark_movable(sd->dict);
 #else
   rb_gc_mark(sd->buf);
+  rb_gc_mark(sd->dict);
 #endif
 }
 
@@ -40,6 +43,7 @@ streaming_decompress_compact(void *p)
 {
   struct streaming_decompress_t *sd = p;
   sd->buf = rb_gc_location(sd->buf);
+  sd->dict = rb_gc_location(sd->dict);
 }
 #endif
 
@@ -64,6 +68,7 @@ rb_streaming_decompress_allocate(VALUE klass)
   sd->dctx = NULL;
   RB_OBJ_WRITE(obj, &sd->buf, Qnil);
   sd->buf_size = 0;
+  RB_OBJ_WRITE(obj, &sd->dict, Qnil);
   return obj;
 }
 
@@ -81,9 +86,10 @@ rb_streaming_decompress_initialize(int argc, VALUE *argv, VALUE obj)
   if (dctx == NULL) {
     rb_raise(rb_eRuntimeError, "%s", "ZSTD_createDCtx error");
   }
-  set_decompress_params(dctx, kwargs);
+  VALUE dict = set_decompress_params(dctx, kwargs);
 
   sd->dctx = dctx;
+  RB_OBJ_WRITE(obj, &sd->dict, dict);
   RB_OBJ_WRITE(obj, &sd->buf, rb_str_new(NULL, buffOutSize));
   sd->buf_size = buffOutSize;
 

spec/zstd-ruby-streaming-compress_spec.rb

@@ -91,6 +91,28 @@
     end
   end
 
+  describe 'Zstd::CDict dictionary the caller does not keep' do
+    let(:dictionary) do
+      File.read("#{__dir__}/dictionary")
+    end
+    let(:user_json) do
+      File.read("#{__dir__}/user_springmt.json")
+    end
+    it 'stays alive as long as the stream that references it' do
+      # The CDict is never stored anywhere: the stream is its only reference.
+      stream = Zstd::StreamingCompress.new(dict: Zstd::CDict.new(dictionary, 5))
+      GC.start
+      GC.compact
+      1000.times { |i| "fill the slot the CDict would have freed #{i}" }
+      GC.start
+
+      stream << user_json
+      compressed = stream.finish
+
+      expect(Zstd.decompress(compressed, dict: dictionary)).to eq(user_json)
+    end
+  end
+
   describe 'nil dictionary' do
     let(:user_json) do
       File.read("#{__dir__}/user_springmt.json")

spec/zstd-ruby-streaming-decompress_spec.rb

@@ -138,6 +138,29 @@
     end
   end
 
+  describe 'Zstd::DDict dictionary the caller does not keep' do
+    let(:dictionary) do
+      File.read("#{__dir__}/dictionary")
+    end
+    let(:user_json) do
+      File.read("#{__dir__}/user_springmt.json")
+    end
+    it 'stays alive as long as the stream that references it' do
+      compressed_json = Zstd.compress(user_json, dict: dictionary)
+      # The DDict is never stored anywhere: the stream is its only reference.
+      stream = Zstd::StreamingDecompress.new(dict: Zstd::DDict.new(dictionary))
+      GC.start
+      GC.compact
+      1000.times { |i| "fill the slot the DDict would have freed #{i}" }
+      GC.start
+
+      result = +''
+      result << stream.decompress(compressed_json[0, 5])
+      result << stream.decompress(compressed_json[5..-1])
+      expect(result).to eq(user_json)
+    end
+  end
+
   describe 'nil dictionary streaming decompress + GC.compact' do
     let(:dictionary) do
       File.read("#{__dir__}/dictionary")