Fix ZSTD_DCtx leak when a frame fails to decode

ext/zstdruby/common.h

@@ -122,6 +122,7 @@ static size_t zstd_compress(ZSTD_CCtx* const ctx, char* output_data, size_t outp
 #endif
 }
 
+/* Raises without freeing dctx: the caller owns it and has to release it. */
 static void set_decompress_params(ZSTD_DCtx* const dctx, VALUE kwargs)
 {
   ID kwargs_keys[1];
@@ -134,19 +135,16 @@ static void set_decompress_params(ZSTD_DCtx* const dctx, VALUE kwargs)
       ZSTD_DDict* ddict = DATA_PTR(kwargs_values[0]);
       size_t ref_dict_ret = ZSTD_DCtx_refDDict(dctx, ddict);
       if (ZSTD_isError(ref_dict_ret)) {
-        ZSTD_freeDCtx(dctx);
         rb_raise(rb_eRuntimeError, "%s", "ZSTD_DCtx_refDDict failed");
       }
     } 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]);
       size_t load_dict_ret = ZSTD_DCtx_loadDictionary(dctx, dict_buffer, dict_size);
       if (ZSTD_isError(load_dict_ret)) {
-        ZSTD_freeDCtx(dctx);
         rb_raise(rb_eRuntimeError, "%s", "ZSTD_CCtx_loadDictionary failed");
       }
     } else {
-      ZSTD_freeDCtx(dctx);
       rb_raise(rb_eArgError, "`dict:` must be a Zstd::DDict or a String");
     }
   }

ext/zstdruby/streaming_decompress.c

@@ -81,9 +81,10 @@ rb_streaming_decompress_initialize(int argc, VALUE *argv, VALUE obj)
   if (dctx == NULL) {
     rb_raise(rb_eRuntimeError, "%s", "ZSTD_createDCtx error");
   }
+  /* Before set_decompress_params, which can raise: the free callback owns it. */
+  sd->dctx = dctx;
   set_decompress_params(dctx, kwargs);
 
-  sd->dctx = dctx;
   RB_OBJ_WRITE(obj, &sd->buf, rb_str_new(NULL, buffOutSize));
   sd->buf_size = buffOutSize;
 

ext/zstdruby/zstdruby.c

@@ -40,32 +40,53 @@ static VALUE rb_compress(int argc, VALUE *argv, VALUE self)
   return output;
 }
 
-static VALUE decode_one_frame(ZSTD_DCtx* dctx, const unsigned char* src, size_t size, VALUE kwargs, size_t* consumed) {
-  VALUE out = rb_str_buf_new(0);
-  size_t cap = ZSTD_DStreamOutSize();
-  char *buf = ALLOC_N(char, cap);
-  ZSTD_inBuffer in = (ZSTD_inBuffer){ src, size, 0 };
+struct decode_frame {
+  ZSTD_DCtx* dctx;
+  char* buf;
+  size_t cap;
+  ZSTD_inBuffer in;
+  VALUE out;
+};
 
-  ZSTD_DCtx_reset(dctx, ZSTD_reset_session_only);
-  set_decompress_params(dctx, kwargs);
+static VALUE decode_frame_body(VALUE arg) {
+  struct decode_frame* st = (struct decode_frame*)arg;
 
   for (;;) {
-    ZSTD_outBuffer o = (ZSTD_outBuffer){ buf, cap, 0 };
-    size_t ret = ZSTD_decompressStream(dctx, &o, &in);
+    ZSTD_outBuffer o = (ZSTD_outBuffer){ st->buf, st->cap, 0 };
+    size_t ret = ZSTD_decompressStream(st->dctx, &o, &st->in);
     if (ZSTD_isError(ret)) {
-      xfree(buf);
       rb_raise(rb_eRuntimeError, "ZSTD_decompressStream failed: %s", ZSTD_getErrorName(ret));
     }
     if (o.pos) {
-      rb_str_cat(out, buf, o.pos);
+      rb_str_cat(st->out, st->buf, o.pos);
     }
     if (ret == 0) {
       break;
     }
   }
-  xfree(buf);
+  return st->out;
+}
+
+static VALUE decode_frame_ensure(VALUE arg) {
+  struct decode_frame* st = (struct decode_frame*)arg;
+  xfree(st->buf);
+  return Qnil;
+}
+
+static VALUE decode_one_frame(ZSTD_DCtx* dctx, const unsigned char* src, size_t size, VALUE kwargs, size_t* consumed) {
+  ZSTD_DCtx_reset(dctx, ZSTD_reset_session_only);
+  set_decompress_params(dctx, kwargs);
+
+  struct decode_frame st;
+  st.dctx = dctx;
+  st.out = rb_str_buf_new(0);
+  st.cap = ZSTD_DStreamOutSize();
+  st.buf = ALLOC_N(char, st.cap);
+  st.in = (ZSTD_inBuffer){ src, size, 0 };
+
+  VALUE out = rb_ensure(decode_frame_body, (VALUE)&st, decode_frame_ensure, (VALUE)&st);
   if (consumed) {
-    *consumed = in.pos;
+    *consumed = st.in.pos;
   }
   return out;
 }
@@ -74,21 +95,24 @@ static VALUE decompress_buffered(ZSTD_DCtx* dctx, const char* data, size_t len)
   return decode_one_frame(dctx, (const unsigned char*)data, len, Qnil, NULL);
 }
 
-static VALUE rb_decompress(int argc, VALUE *argv, VALUE self)
-{
-  VALUE input_value, kwargs;
-  rb_scan_args(argc, argv, "10:", &input_value, &kwargs);
-  StringValue(input_value);
+struct decompress_scan {
+  const unsigned char* in;
+  size_t in_size;
+  VALUE kwargs;
+  ZSTD_DCtx* dctx;
+};
 
-  size_t in_size = RSTRING_LEN(input_value);
-  const unsigned char *in = (const unsigned char *)RSTRING_PTR(input_value);
+static VALUE decompress_scan_body(VALUE arg)
+{
+  struct decompress_scan* st = (struct decompress_scan*)arg;
+  const unsigned char *in = st->in;
+  size_t in_size = st->in_size;
 
   size_t off = 0;
   const uint32_t ZSTD_MAGIC = 0xFD2FB528U;
   const uint32_t SKIP_LO    = 0x184D2A50U; /* ...5F */
 
   VALUE result = Qnil;
-  ZSTD_DCtx *dctx = NULL;
 
   while (off + 4 <= in_size) {
     uint32_t magic = (uint32_t)in[off]
@@ -109,15 +133,15 @@ static VALUE rb_decompress(int argc, VALUE *argv, VALUE self)
     }
 
     if (magic == ZSTD_MAGIC) {
-      if (dctx == NULL) {
-        dctx = ZSTD_createDCtx();
-        if (!dctx) {
+      if (st->dctx == NULL) {
+        st->dctx = ZSTD_createDCtx();
+        if (!st->dctx) {
           rb_raise(rb_eRuntimeError, "ZSTD_createDCtx failed");
         }
       }
 
       size_t consumed = 0;
-      VALUE out = decode_one_frame(dctx, in + off, in_size - off, kwargs, &consumed);
+      VALUE out = decode_one_frame(st->dctx, in + off, in_size - off, st->kwargs, &consumed);
       if (result == Qnil) {
         /* First frame becomes the accumulator, avoiding a copy of its
            (potentially large) output in the common single-frame case. */
@@ -137,9 +161,31 @@ static VALUE rb_decompress(int argc, VALUE *argv, VALUE self)
     off += 1;
   }
 
-  if (dctx != NULL) {
-    ZSTD_freeDCtx(dctx);
+  return result;
+}
+
+static VALUE decompress_scan_ensure(VALUE arg)
+{
+  struct decompress_scan* st = (struct decompress_scan*)arg;
+  if (st->dctx != NULL) {
+    ZSTD_freeDCtx(st->dctx);
   }
+  return Qnil;
+}
+
+static VALUE rb_decompress(int argc, VALUE *argv, VALUE self)
+{
+  VALUE input_value, kwargs;
+  rb_scan_args(argc, argv, "10:", &input_value, &kwargs);
+  StringValue(input_value);
+
+  struct decompress_scan st;
+  st.in = (const unsigned char *)RSTRING_PTR(input_value);
+  st.in_size = RSTRING_LEN(input_value);
+  st.kwargs = kwargs;
+  st.dctx = NULL;
+
+  VALUE result = rb_ensure(decompress_scan_body, (VALUE)&st, decompress_scan_ensure, (VALUE)&st);
 
   RB_GC_GUARD(input_value);
   if (result == Qnil) {

spec/zstd-ruby_spec.rb

@@ -116,6 +116,21 @@ def to_str
       expect { Zstd.decompress(Object.new) }.to raise_error(TypeError)
     end
 
+    # These two walk the paths that used to leak the ZSTD_DCtx. Nothing here
+    # asserts the leak itself -- Valgrind or ASan on these examples reports it.
+    it 'should raise when a frame body fails to decode' do
+      # A valid frame header, so libzstd allocates its buffers from it, followed
+      # by a truncated body.
+      good = Zstd.compress(user_json * 50)
+      broken = good.byteslice(0, good.bytesize / 2) + ("\x00" * 32)
+
+      expect { Zstd.decompress(broken) }.to raise_error(RuntimeError)
+    end
+
+    it 'should raise when the dict argument is rejected' do
+      expect { Zstd.decompress(Zstd.compress('abc'), dict: 123) }.to raise_error(ArgumentError)
+    end
+
     class DummyForDecompress
       def to_str
         Zstd.compress('abc')