diff --git a/ext/zstdruby/zstdruby.c b/ext/zstdruby/zstdruby.c index 1649fe5..aa8e28e 100644 --- a/ext/zstdruby/zstdruby.c +++ b/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; } diff --git a/spec/zstd-ruby_spec.rb b/spec/zstd-ruby_spec.rb index ddfae47..401f552 100644 --- a/spec/zstd-ruby_spec.rb +++ b/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')