Fix ZSTD_DCtx leak when a frame fails to decode

ext/zstdruby/zstdruby.c

@@ -40,31 +40,55 @@ 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) {
-  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 out;
+  return st->out;
+}
+
+static VALUE decode_frame_ensure(VALUE arg) {
+  struct decode_frame* st = (struct decode_frame*)arg;
+  xfree(st->buf);
+  ZSTD_freeDCtx(st->dctx);
+  return Qnil;
+}
+
+/* Takes ownership of dctx: it is freed before this returns, on every path. */
+static VALUE decode_one_frame(ZSTD_DCtx* dctx, const unsigned char* src, size_t size, VALUE kwargs) {
+  ZSTD_DCtx_reset(dctx, ZSTD_reset_session_only);
+  /* set_decompress_params frees dctx itself before raising, so it has to run
+     before the ensure below takes ownership of it. */
+  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 };
+
+  return rb_ensure(decode_frame_body, (VALUE)&st, decode_frame_ensure, (VALUE)&st);
 }
 
 static VALUE decompress_buffered(ZSTD_DCtx* dctx, const char* data, size_t len) {
@@ -110,7 +134,6 @@ static VALUE rb_decompress(int argc, VALUE *argv, VALUE self)
 
       VALUE out = decode_one_frame(dctx, in + off, in_size - off, kwargs);
 
-      ZSTD_freeDCtx(dctx);
       RB_GC_GUARD(input_value);
       return out;
     }

spec/zstd-ruby_spec.rb

@@ -103,6 +103,18 @@ def to_str
       expect { Zstd.decompress(Object.new) }.to raise_error(TypeError)
     end
 
+    # Walks the failure path that used to leak the ZSTD_DCtx and the scratch
+    # buffer. Nothing here asserts the leak itself -- `rake spec:valgrind`
+    # 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(File.read("#{__dir__}/user_springmt.json") * 50)
+      broken = good.byteslice(0, good.bytesize / 2) + ("\x00" * 32)
+
+      expect { Zstd.decompress(broken) }.to raise_error(RuntimeError)
+    end
+
     class DummyForDecompress
       def to_str
         Zstd.compress('abc')