Skip to content

Commit 8b40eb1

Browse files
committed
Potential memory issue found by Grok Fast 1
1 parent b09d02b commit 8b40eb1

File tree

2 files changed

+11
-11
lines changed

2 files changed

+11
-11
lines changed

ext/zstdruby/context.c

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ static VALUE rb_cZstdCContext;
2121
static VALUE rb_cZstdDContext;
2222

2323
// Forward declaration of decompress_buffered from zstdruby.c
24-
extern VALUE decompress_buffered(ZSTD_DCtx* dctx, const char* input_data, size_t input_size);
24+
extern VALUE decompress_buffered(ZSTD_DCtx* dctx, const char* input_data, size_t input_size, bool free_ctx);
2525

2626
// CContext (compression-only) implementation
2727
static void zstd_ccontext_mark(void *ptr)
@@ -275,10 +275,10 @@ static VALUE zstd_dcontext_decompress(VALUE self, VALUE input_value)
275275
rb_raise(rb_eRuntimeError, "Not compressed by zstd: %s", ZSTD_getErrorName(uncompressed_size));
276276
}
277277

278-
if (uncompressed_size == ZSTD_CONTENTSIZE_UNKNOWN) {
279-
ctx->needs_reset = 1;
280-
return decompress_buffered(ctx->dctx, input_data, input_size);
281-
}
278+
if (uncompressed_size == ZSTD_CONTENTSIZE_UNKNOWN) {
279+
ctx->needs_reset = 1;
280+
return decompress_buffered(ctx->dctx, input_data, input_size, false);
281+
}
282282

283283
VALUE output = rb_str_new(NULL, uncompressed_size);
284284
char* output_data = RSTRING_PTR(output);

ext/zstdruby/zstdruby.c

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ static VALUE rb_compress(int argc, VALUE *argv, VALUE self)
3939
return output;
4040
}
4141

42-
VALUE decompress_buffered(ZSTD_DCtx* dctx, const char* input_data, size_t input_size)
42+
VALUE decompress_buffered(ZSTD_DCtx* dctx, const char* input_data, size_t input_size, bool free_ctx)
4343
{
4444
ZSTD_inBuffer input = { input_data, input_size, 0 };
4545
VALUE result = rb_str_new(0, 0);
@@ -52,12 +52,12 @@ VALUE decompress_buffered(ZSTD_DCtx* dctx, const char* input_data, size_t input_
5252

5353
size_t ret = zstd_stream_decompress(dctx, &output, &input, false);
5454
if (ZSTD_isError(ret)) {
55-
ZSTD_freeDCtx(dctx);
55+
if (free_ctx) ZSTD_freeDCtx(dctx);
5656
rb_raise(rb_eRuntimeError, "%s: %s", "ZSTD_decompressStream failed", ZSTD_getErrorName(ret));
5757
}
5858
rb_str_cat(result, output.dst, output.pos);
5959
}
60-
ZSTD_freeDCtx(dctx);
60+
if (free_ctx) ZSTD_freeDCtx(dctx);
6161
return result;
6262
}
6363

@@ -81,9 +81,9 @@ static VALUE rb_decompress(int argc, VALUE *argv, VALUE self)
8181
}
8282
// ZSTD_decompressStream may be called multiple times when ZSTD_CONTENTSIZE_UNKNOWN, causing slowness.
8383
// Therefore, we will not standardize on ZSTD_decompressStream
84-
if (uncompressed_size == ZSTD_CONTENTSIZE_UNKNOWN) {
85-
return decompress_buffered(dctx, input_data, input_size);
86-
}
84+
if (uncompressed_size == ZSTD_CONTENTSIZE_UNKNOWN) {
85+
return decompress_buffered(dctx, input_data, input_size, true);
86+
}
8787

8888
VALUE output = rb_str_new(NULL, uncompressed_size);
8989
char* output_data = RSTRING_PTR(output);

0 commit comments

Comments
 (0)