Skip to content

Commit 10dfa82

Browse files
committed
Make more robust for memory
Signed-off-by: Hiroshi Hatake <[email protected]>
1 parent 93df119 commit 10dfa82

File tree

2 files changed

+13
-9
lines changed

2 files changed

+13
-9
lines changed

ext/zstdruby/streaming_compress.c

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,6 @@ rb_streaming_compress_write(int argc, VALUE *argv, VALUE obj)
157157
const char* input_data = RSTRING_PTR(str);
158158
size_t input_size = RSTRING_LEN(str);
159159
ZSTD_inBuffer input = { input_data, input_size, 0 };
160-
VALUE result = rb_str_new(0, 0);
161160

162161
while (input.pos < input.size) {
163162
const char* output_data = RSTRING_PTR(sc->buf);
@@ -166,12 +165,11 @@ rb_streaming_compress_write(int argc, VALUE *argv, VALUE obj)
166165
if (ZSTD_isError(ret)) {
167166
rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(ret));
168167
}
169-
/* collect produced bytes */
168+
/* Directly append to the pending buffer */
170169
if (output.pos > 0) {
171-
rb_str_cat(result, output.dst, output.pos);
170+
rb_str_cat(sc->pending, output.dst, output.pos);
172171
}
173172
}
174-
rb_str_cat(sc->pending, RSTRING_PTR(result), RSTRING_LEN(result));
175173
total += RSTRING_LEN(str);
176174
}
177175

ext/zstdruby/streaming_decompress.c

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -101,16 +101,22 @@ rb_streaming_decompress_decompress(VALUE obj, VALUE src)
101101
struct streaming_decompress_t* sd;
102102
TypedData_Get_Struct(obj, struct streaming_decompress_t, &streaming_decompress_type, sd);
103103
VALUE result = rb_str_new(0, 0);
104-
size_t ret;
105-
do {
104+
105+
while (input.pos < input.size) {
106106
const char* output_data = RSTRING_PTR(sd->buf);
107107
ZSTD_outBuffer output = { (void*)output_data, sd->buf_size, 0 };
108-
ret = zstd_stream_decompress(sd->dctx, &output, &input, false);
108+
size_t const ret = zstd_stream_decompress(sd->dctx, &output, &input, false);
109+
109110
if (ZSTD_isError(ret)) {
110111
rb_raise(rb_eRuntimeError, "decompress error error code: %s", ZSTD_getErrorName(ret));
111112
}
112-
rb_str_cat(result, output.dst, output.pos);
113-
} while (input.pos < input.size && ret > 0);
113+
if (output.pos > 0) {
114+
rb_str_cat(result, output.dst, output.pos);
115+
}
116+
if (ret == 0 && output.pos == 0) {
117+
break;
118+
}
119+
}
114120
return result;
115121
}
116122

0 commit comments

Comments
 (0)