From 31f8a396c3c1a8025baf2bb78177ee70b3cf5f49 Mon Sep 17 00:00:00 2001 From: Bala Kumar Date: Mon, 5 Oct 2026 21:44:06 +0530 Subject: [PATCH] Free the compression context when Zstd.compress raises Zstd.compress frees its ZSTD_CCtx only after compression returns normally, so any exception raised in between abandons the context. That includes interrupts: zstd_compress releases the GVL, and when it reacquires it Ruby delivers pending interrupts (Thread#raise, Timeout, Thread#kill, Sidekiq shutdown) by raising out of the call. An interrupted level-19 compression of a 32 MiB input leaks roughly 75-85 MB per call. The same applies to argument errors raised after the context exists: an unknown keyword or a level that does not fit in an int leaked the context from both Zstd.compress and StreamingCompress#initialize. Run the body of Zstd.compress under rb_ensure so the context is freed on every path. set_compress_params and convert_compression_level used to free the context themselves before raising; with an ensure also owning it that would be a double free, so they now raise and leave it to the owner. StreamingCompress#initialize therefore assigns sc->ctx before calling set_compress_params, so the object's free callback owns it. This mirrors the decompression-side change in #148. --- ext/zstdruby/common.h | 13 ++---- ext/zstdruby/streaming_compress.c | 4 +- ext/zstdruby/zstdruby.c | 54 ++++++++++++++++------- spec/zstd-ruby-streaming-compress_spec.rb | 12 +++++ spec/zstd-ruby_spec.rb | 27 ++++++++++++ 5 files changed, 83 insertions(+), 27 deletions(-) diff --git a/ext/zstdruby/common.h b/ext/zstdruby/common.h index f051440..21f6fd1 100644 --- a/ext/zstdruby/common.h +++ b/ext/zstdruby/common.h @@ -10,15 +10,12 @@ extern VALUE rb_cCDict, rb_cDDict; -static int convert_compression_level(ZSTD_CCtx* ctx, VALUE compression_level_value) +static int convert_compression_level(VALUE compression_level_value) { if (NIL_P(compression_level_value)) { return ZSTD_CLEVEL_DEFAULT; } if (!RB_INTEGER_TYPE_P(compression_level_value)) { - if (ctx) { - ZSTD_freeCCtx(ctx); - } rb_raise(rb_eTypeError, "compression level must be an Integer"); } return NUM2INT(compression_level_value); @@ -27,7 +24,8 @@ static int convert_compression_level(ZSTD_CCtx* ctx, VALUE compression_level_val /* 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. */ + dictionary needs no such handling: ZSTD_CCtx_loadDictionary copies it. + Raises without freeing ctx: the caller owns it and has to release it. */ static VALUE set_compress_params(ZSTD_CCtx* const ctx, VALUE kwargs) { ID kwargs_keys[2]; @@ -38,7 +36,7 @@ static VALUE set_compress_params(ZSTD_CCtx* const ctx, VALUE kwargs) int compression_level = ZSTD_CLEVEL_DEFAULT; if (kwargs_values[0] != Qundef && kwargs_values[0] != Qnil) { - compression_level = convert_compression_level(ctx, kwargs_values[0]); + compression_level = convert_compression_level(kwargs_values[0]); } ZSTD_CCtx_setParameter(ctx, ZSTD_c_compressionLevel, compression_level); @@ -47,7 +45,6 @@ static VALUE set_compress_params(ZSTD_CCtx* const ctx, VALUE kwargs) ZSTD_CDict* cdict = DATA_PTR(kwargs_values[1]); size_t ref_dict_ret = ZSTD_CCtx_refCDict(ctx, cdict); if (ZSTD_isError(ref_dict_ret)) { - ZSTD_freeCCtx(ctx); rb_raise(rb_eRuntimeError, "%s", "ZSTD_CCtx_refCDict failed"); } return kwargs_values[1]; @@ -56,11 +53,9 @@ static VALUE set_compress_params(ZSTD_CCtx* const ctx, VALUE kwargs) size_t dict_size = RSTRING_LEN(kwargs_values[1]); size_t load_dict_ret = ZSTD_CCtx_loadDictionary(ctx, dict_buffer, dict_size); if (ZSTD_isError(load_dict_ret)) { - ZSTD_freeCCtx(ctx); rb_raise(rb_eRuntimeError, "%s", "ZSTD_CCtx_loadDictionary failed"); } } else { - ZSTD_freeCCtx(ctx); rb_raise(rb_eArgError, "`dict:` must be a Zstd::CDict or a String"); } } diff --git a/ext/zstdruby/streaming_compress.c b/ext/zstdruby/streaming_compress.c index d53d6d5..ac3f651 100644 --- a/ext/zstdruby/streaming_compress.c +++ b/ext/zstdruby/streaming_compress.c @@ -91,9 +91,9 @@ rb_streaming_compress_initialize(int argc, VALUE *argv, VALUE obj) if (ctx == NULL) { rb_raise(rb_eRuntimeError, "%s", "ZSTD_createCCtx error"); } - VALUE dict = set_compress_params(ctx, kwargs); - + /* Before set_compress_params, which can raise: the free callback owns it. */ sc->ctx = ctx; + VALUE dict = set_compress_params(ctx, kwargs); RB_OBJ_WRITE(obj, &sc->dict, dict); RB_OBJ_WRITE(obj, &sc->buf, rb_str_new(NULL, buffOutSize)); sc->buf_size = buffOutSize; diff --git a/ext/zstdruby/zstdruby.c b/ext/zstdruby/zstdruby.c index 95d8c41..c3aeeb9 100644 --- a/ext/zstdruby/zstdruby.c +++ b/ext/zstdruby/zstdruby.c @@ -8,30 +8,25 @@ static VALUE zstdVersion(VALUE self) return INT2NUM(version); } -static VALUE rb_compress(int argc, VALUE *argv, VALUE self) -{ +struct compress_args { + ZSTD_CCtx* ctx; VALUE input_value; VALUE kwargs; - rb_scan_args(argc, argv, "10:", &input_value, &kwargs); - - StringValue(input_value); - - ZSTD_CCtx* const ctx = ZSTD_createCCtx(); - if (ctx == NULL) { - rb_raise(rb_eRuntimeError, "%s", "ZSTD_createCCtx error"); - } +}; - set_compress_params(ctx, kwargs); +static VALUE compress_body(VALUE arg) +{ + struct compress_args* args = (struct compress_args*)arg; + set_compress_params(args->ctx, args->kwargs); - char* input_data = RSTRING_PTR(input_value); - size_t input_size = RSTRING_LEN(input_value); + char* input_data = RSTRING_PTR(args->input_value); + size_t input_size = RSTRING_LEN(args->input_value); size_t max_compressed_size = ZSTD_compressBound(input_size); VALUE output = rb_str_new(NULL, max_compressed_size); char* output_data = RSTRING_PTR(output); - size_t const ret = zstd_compress(ctx, output_data, max_compressed_size, input_data, input_size, false); - ZSTD_freeCCtx(ctx); + size_t const ret = zstd_compress(args->ctx, output_data, max_compressed_size, input_data, input_size, false); if (ZSTD_isError(ret)) { rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(ret)); } @@ -40,6 +35,33 @@ static VALUE rb_compress(int argc, VALUE *argv, VALUE self) return output; } +static VALUE compress_ensure(VALUE arg) +{ + ZSTD_freeCCtx(((struct compress_args*)arg)->ctx); + return Qnil; +} + +static VALUE rb_compress(int argc, VALUE *argv, VALUE self) +{ + VALUE input_value; + VALUE kwargs; + rb_scan_args(argc, argv, "10:", &input_value, &kwargs); + + StringValue(input_value); + + ZSTD_CCtx* const ctx = ZSTD_createCCtx(); + if (ctx == NULL) { + rb_raise(rb_eRuntimeError, "%s", "ZSTD_createCCtx error"); + } + + /* The body can raise -- a bad keyword, or an interrupt delivered when + zstd_compress reacquires the GVL -- so the context is freed under ensure. */ + struct compress_args args = { ctx, input_value, kwargs }; + VALUE output = rb_ensure(compress_body, (VALUE)&args, compress_ensure, (VALUE)&args); + RB_GC_GUARD(input_value); + 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(); @@ -196,7 +218,7 @@ static VALUE rb_cdict_initialize(int argc, VALUE *argv, VALUE self) VALUE dict; VALUE compression_level_value; rb_scan_args(argc, argv, "11", &dict, &compression_level_value); - int compression_level = convert_compression_level(NULL, compression_level_value); + int compression_level = convert_compression_level(compression_level_value); StringValue(dict); char* dict_buffer = RSTRING_PTR(dict); diff --git a/spec/zstd-ruby-streaming-compress_spec.rb b/spec/zstd-ruby-streaming-compress_spec.rb index 77f4601..f05a8a7 100644 --- a/spec/zstd-ruby-streaming-compress_spec.rb +++ b/spec/zstd-ruby-streaming-compress_spec.rb @@ -129,6 +129,18 @@ end end + # initialize raises after creating the ZSTD_CCtx; the object's free callback + # has to own it by then. `rake spec:valgrind` catches a leak or a double free. + describe 'initialize raising after creating the compression context' do + it 'raises ArgumentError for an unknown keyword' do + expect { Zstd::StreamingCompress.new(unknown: 1) }.to raise_error(ArgumentError) + end + + it 'raises ArgumentError for a dict: that is neither a CDict nor a String' do + expect { Zstd::StreamingCompress.new(dict: 123) }.to raise_error(ArgumentError) + end + end + if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('3.0.0') describe 'Ractor' do it 'should be supported' do diff --git a/spec/zstd-ruby_spec.rb b/spec/zstd-ruby_spec.rb index 88a5846..3b92703 100644 --- a/spec/zstd-ruby_spec.rb +++ b/spec/zstd-ruby_spec.rb @@ -55,6 +55,33 @@ def to_str decompressed = Zstd.decompress(compressed) expect(decompressed).to eq('abc') end + + # Each of these raises after the ZSTD_CCtx has been created. The assertions + # pin the behaviour; `rake spec:valgrind` is what catches a leaked context. + context 'when it raises after creating the compression context' do + it 'raises ArgumentError for an unknown keyword' do + expect { Zstd.compress('abc', unknown: 1) }.to raise_error(ArgumentError) + end + + it 'raises RangeError for a level that does not fit in an int' do + expect { Zstd.compress('abc', level: 2**40) }.to raise_error(RangeError) + end + + it 'raises ArgumentError for a dict: that is neither a CDict nor a String' do + expect { Zstd.compress('abc', dict: 123) }.to raise_error(ArgumentError) + end + + it 'can be interrupted by Thread#raise while compressing' do + interrupt = Class.new(StandardError) + input = Random.new(42).bytes(16 * 1024 * 1024) + thread = Thread.new { Zstd.compress(input, level: 19) } + thread.report_on_exception = false + sleep 0.1 + thread.raise(interrupt) + + expect { thread.join }.to raise_error(interrupt) + end + end end describe 'decompress' do