Implements handling of partial frames by the streaming decompressor

Fixes SpringMT#87
gucki · May 16, 2024 · e7f8b54 · e7f8b54
1 parent 183e9cb
commit e7f8b54
Showing 1 changed file with 22 additions and 0 deletions.
diff --git a/ext/zstdruby/streaming_decompress.c b/ext/zstdruby/streaming_decompress.c
@@ -113,6 +113,27 @@ rb_streaming_decompress_decompress(VALUE obj, VALUE src)
   return result;
 }
 
+static VALUE
+rb_streaming_decompress_decompress2(VALUE obj, VALUE src)
+{
+  StringValue(src);
+  const char* input_data = RSTRING_PTR(src);
+  size_t input_size = RSTRING_LEN(src);
+  ZSTD_inBuffer input = { input_data, input_size, 0 };
+
+  struct streaming_decompress_t* sd;
+  TypedData_Get_Struct(obj, struct streaming_decompress_t, &streaming_decompress_type, sd);
+  const char* output_data = RSTRING_PTR(sd->buf);
+  VALUE result = rb_str_new(0, 0);
+  ZSTD_outBuffer output = { (void*)output_data, sd->buf_size, 0 };
+  size_t const ret = zstd_stream_decompress(sd->dctx, &output, &input, false);
+  if (ZSTD_isError(ret)) {
+    rb_raise(rb_eRuntimeError, "decompress error error code: %s", ZSTD_getErrorName(ret));
+  }
+  rb_str_cat(result, output.dst, output.pos);
+  return rb_ary_new_from_args(3, UINT2NUM(ret), result, ULONG2NUM(input.pos));
+}
+
 extern VALUE rb_mZstd, cStreamingDecompress;
 void
 zstd_ruby_streaming_decompress_init(void)
@@ -121,4 +142,5 @@ zstd_ruby_streaming_decompress_init(void)
   rb_define_alloc_func(cStreamingDecompress, rb_streaming_decompress_allocate);
   rb_define_method(cStreamingDecompress, "initialize", rb_streaming_decompress_initialize, -1);
   rb_define_method(cStreamingDecompress, "decompress", rb_streaming_decompress_decompress, 1);
+  rb_define_method(cStreamingDecompress, "decompress2", rb_streaming_decompress_decompress2, 1);
 }