diff --git a/ext/zstdruby/common.h b/ext/zstdruby/common.h index e0b0f4f..469c9fc 100644 --- a/ext/zstdruby/common.h +++ b/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,7 +135,6 @@ 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) { @@ -142,11 +142,9 @@ static void set_decompress_params(ZSTD_DCtx* const dctx, VALUE kwargs) 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"); } } diff --git a/ext/zstdruby/streaming_decompress.c b/ext/zstdruby/streaming_decompress.c index c2d2b8f..820bc8f 100644 --- a/ext/zstdruby/streaming_decompress.c +++ b/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; diff --git a/ext/zstdruby/zstdruby.c b/ext/zstdruby/zstdruby.c index 3e8454b..2849f01 100644 --- a/ext/zstdruby/zstdruby.c +++ b/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) { diff --git a/spec/zstd-ruby_spec.rb b/spec/zstd-ruby_spec.rb index 6ef02d9..f5073a4 100644 --- a/spec/zstd-ruby_spec.rb +++ b/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')