vibe_zstd 1.1.0 → 1.3.0

data/ext/vibe_zstd/dctx.c

@@ -8,9 +8,24 @@ extern rb_data_type_t vibe_zstd_dctx_type;
 // Class-level default for initial capacity (0 = use ZSTD_DStreamOutSize)
 static size_t default_initial_capacity = 0;
 
+// Class-level default output-size limit (0 = unlimited)
+static size_t default_max_decompressed_size = 0;
+
+// VibeZstd::DecompressedSizeExceeded - raised when output exceeds the limit.
+// Defined in vibe_zstd_dctx_init_class, cached here for use on the error path.
+static VALUE rb_eDecompressedSizeExceeded;
+
 // Helper to set DCtx parameter from Ruby keyword argument
 static int
 vibe_zstd_dctx_init_param_iter(VALUE key, VALUE value, VALUE self) {
+    // Guard: only Symbol keys are valid.  A non-Symbol key (e.g. a String like
+    // "format" => 1) would make SYM2ID undefined behaviour, so reject it early.
+    if (!SYMBOL_P(key)) {
+        rb_raise(rb_eArgError,
+                 "DCtx.new option keys must be Symbols (got %"PRIsVALUE")",
+                 rb_inspect(key));
+    }
+
     // Build the setter method name: key + "="
     const char* key_str = rb_id2name(SYM2ID(key));
     char setter[256];
@@ -51,7 +66,8 @@ typedef struct {
 } dctx_param_entry;
 
 static dctx_param_entry dctx_param_table[] = {
-    {0, ZSTD_d_windowLogMax, "window_log_max"}
+    {0, ZSTD_d_windowLogMax, "window_log_max"},
+    {0, ZSTD_d_format, "format"}
 };
 
 #define DCTX_PARAM_TABLE_SIZE (sizeof(dctx_param_table) / sizeof(dctx_param_entry))
@@ -136,6 +152,7 @@ vibe_zstd_dctx_get_param_generic(VALUE self, ZSTD_dParameter param, const char*
 
 // Define all DCtx parameter accessors
 DEFINE_DCTX_PARAM_ACCESSORS(window_log_max, ZSTD_d_windowLogMax, "window_log_max")
+DEFINE_DCTX_PARAM_ACCESSORS(format, ZSTD_d_format, "format")
 
 // DCtx parameter_bounds - query parameter bounds (class method, kept for introspection)
 static VALUE
@@ -215,6 +232,54 @@ vibe_zstd_dctx_set_initial_capacity(VALUE self, VALUE value) {
     return value;
 }
 
+// DCtx default_max_decompressed_size getter (class method); 0 = unlimited
+static VALUE
+vibe_zstd_dctx_get_default_max_decompressed_size(VALUE self) {
+    return SIZET2NUM(default_max_decompressed_size);
+}
+
+// DCtx default_max_decompressed_size setter (class method)
+static VALUE
+vibe_zstd_dctx_set_default_max_decompressed_size(VALUE self, VALUE value) {
+    if (NIL_P(value)) {
+        default_max_decompressed_size = 0;  // unlimited
+    } else {
+        default_max_decompressed_size = NUM2SIZET(value);
+    }
+    return value;
+}
+
+// DCtx max_decompressed_size getter (instance method); reports the effective
+// limit, falling back to the class default. Returns 0 when unlimited.
+static VALUE
+vibe_zstd_dctx_get_max_decompressed_size(VALUE self) {
+    vibe_zstd_dctx* dctx;
+    TypedData_Get_Struct(self, vibe_zstd_dctx, &vibe_zstd_dctx_type, dctx);
+
+    if (dctx->max_decompressed_size == 0) {
+        return SIZET2NUM(default_max_decompressed_size);
+    }
+    return SIZET2NUM(dctx->max_decompressed_size);
+}
+
+// DCtx max_decompressed_size setter (instance method); nil = inherit class default
+static VALUE
+vibe_zstd_dctx_set_max_decompressed_size(VALUE self, VALUE value) {
+    vibe_zstd_dctx* dctx;
+    TypedData_Get_Struct(self, vibe_zstd_dctx, &vibe_zstd_dctx_type, dctx);
+
+    if (NIL_P(value)) {
+        dctx->max_decompressed_size = 0;  // inherit class default
+    } else {
+        size_t limit = NUM2SIZET(value);
+        if (limit == 0) {
+            rb_raise(rb_eArgError, "max_decompressed_size must be positive (or nil to inherit the class default)");
+        }
+        dctx->max_decompressed_size = limit;
+    }
+    return value;
+}
+
 // Decompress args for GVL release
 // This structure packages all arguments needed for decompression so we can
 // call ZSTD functions without holding Ruby's Global VM Lock (GVL).
@@ -253,7 +318,10 @@ typedef struct {
     size_t dst_capacity;
     size_t dst_size;
     size_t initial_capacity;
+    size_t max_size;   // 0 = unlimited; otherwise output must not exceed this
     int error;
+    int limit_exceeded;  // set if output would exceed max_size
+    int truncated;        // set if input was exhausted before the frame completed
     const char *error_name;
 } decompress_stream_nogvl_args;
 
@@ -264,9 +332,15 @@ static void*
 decompress_stream_without_gvl(void* arg) {
     decompress_stream_nogvl_args* args = arg;
     args->error = 0;
+    args->limit_exceeded = 0;
+    args->truncated = 0;
     args->error_name = NULL;
 
     args->dst_capacity = args->initial_capacity;
+    // Never allocate more than the configured limit up front.
+    if (args->max_size && args->dst_capacity > args->max_size) {
+        args->dst_capacity = args->max_size;
+    }
     args->dst = malloc(args->dst_capacity);
     if (!args->dst) {
         args->error = 1;
@@ -276,11 +350,21 @@ decompress_stream_without_gvl(void* arg) {
     args->dst_size = 0;
 
     ZSTD_inBuffer input = { args->src, args->src_size, 0 };
+    size_t last_ret = 1;  // sentinel: non-zero = frame not yet complete
 
     while (input.pos < input.size) {
         // Ensure we have room for output
         if (args->dst_size >= args->dst_capacity) {
             size_t new_capacity = args->dst_capacity * 2;
+            // Clamp growth to the configured limit. If we cannot grow past the
+            // current capacity, the output would exceed the limit.
+            if (args->max_size && new_capacity > args->max_size) {
+                new_capacity = args->max_size;
+            }
+            if (new_capacity <= args->dst_capacity) {
+                args->limit_exceeded = 1;
+                return NULL;
+            }
             char* new_buf = realloc(args->dst, new_capacity);
             if (!new_buf) {
                 args->error = 1;
@@ -305,14 +389,74 @@ decompress_stream_without_gvl(void* arg) {
         }
 
         args->dst_size += output.pos;
+        last_ret = ret;
 
         // ret == 0 means frame is complete
         if (ret == 0) break;
     }
 
+    // If we consumed all input but the last call still reported a non-zero hint
+    // (more input needed), the frame was cut short — flag it as truncated.
+    if (last_ret != 0) {
+        args->truncated = 1;
+    }
+
     return NULL;
 }
 
+// State for the rb_ensure-wrapped unknown-size decompression path.
+// Groups everything the body needs to run the no-GVL stream loop and everything
+// the cleanup needs to release on any exit (raise, async exception, success).
+typedef struct {
+    ZSTD_DCtx* dctx;
+    ZSTD_DDict* ddict;
+    decompress_stream_nogvl_args* args;
+    VALUE data;
+    size_t max_size;
+} dctx_stream_decompress_state;
+
+// Body: run the no-GVL stream loop (source string locked), check the outcome,
+// and build the result string. Raising here is safe: cleanup always runs.
+static VALUE
+vibe_zstd_dctx_stream_decompress_body(VALUE p) {
+    dctx_stream_decompress_state* state = (dctx_stream_decompress_state*)p;
+
+    // Lock the source string while the GVL is released: another Ruby thread
+    // holding the same string must not mutate or GC it mid-decompression.
+    vibe_zstd_nogvl_with_str_locked(decompress_stream_without_gvl, state->args, state->data);
+
+    if (state->args->limit_exceeded) {
+        rb_raise(rb_eDecompressedSizeExceeded,
+                 "Decompressed output exceeds limit of %zu bytes", state->max_size);
+    }
+
+    if (state->args->error) {
+        rb_raise(rb_eRuntimeError, "Decompression failed: %s", state->args->error_name);
+    }
+
+    if (state->args->truncated) {
+        rb_raise(rb_eRuntimeError, "Truncated frame: incomplete zstd data");
+    }
+
+    // Create Ruby string from the C buffer; cleanup frees the buffer
+    return rb_str_new(state->args->dst, state->args->dst_size);
+}
+
+// Cleanup: free the C output buffer and return the context to no-dictionary
+// mode so subsequent calls on this DCtx are not affected.
+static VALUE
+vibe_zstd_dctx_stream_decompress_cleanup(VALUE p) {
+    dctx_stream_decompress_state* state = (dctx_stream_decompress_state*)p;
+    if (state->args->dst) {
+        free(state->args->dst);
+        state->args->dst = NULL;
+    }
+    if (state->ddict) {
+        ZSTD_DCtx_refDDict(state->dctx, NULL);
+    }
+    return Qnil;
+}
+
 // DCtx frame_content_size - class method to get frame content size
 static VALUE
 vibe_zstd_dctx_frame_content_size(VALUE self, VALUE data) {
@@ -353,35 +497,52 @@ vibe_zstd_dctx_decompress(int argc, VALUE* argv, VALUE self) {
     size_t srcSize = RSTRING_LEN(data);
     size_t offset = 0;
 
-    // Skip any leading skippable frames
-    while (offset < srcSize && ZSTD_isSkippableFrame(src + offset, srcSize - offset)) {
-        size_t frameSize = ZSTD_findFrameCompressedSize(src + offset, srcSize - offset);
-        if (ZSTD_isError(frameSize)) {
-            rb_raise(rb_eRuntimeError, "Invalid skippable frame at offset %zu: %s", offset, ZSTD_getErrorName(frameSize));
+    // Magicless frames (format = ZSTD_f_zstd1_magicless) carry no magic number,
+    // so frame introspection (content size, dict ID, skippable detection) cannot
+    // be performed. Force the streaming decompress path, which honors the format
+    // parameter set on the context via ZSTD_decompressStream.
+    int dformat = 0;
+    (void)ZSTD_DCtx_getParameter(dctx->dctx, ZSTD_d_format, &dformat);
+    int magicless = (dformat == ZSTD_f_zstd1_magicless);
+
+    unsigned long long contentSize;
+    unsigned int frame_dict_id;
+
+    if (magicless) {
+        contentSize = ZSTD_CONTENTSIZE_UNKNOWN;  // route to streaming path
+        frame_dict_id = 0;                        // cannot read dict ID without magic
+    } else {
+        // Skip any leading skippable frames
+        while (offset < srcSize && ZSTD_isSkippableFrame(src + offset, srcSize - offset)) {
+            size_t frameSize = ZSTD_findFrameCompressedSize(src + offset, srcSize - offset);
+            if (ZSTD_isError(frameSize)) {
+                rb_raise(rb_eRuntimeError, "Invalid skippable frame at offset %zu: %s", offset, ZSTD_getErrorName(frameSize));
+            }
+            offset += frameSize;
         }
-        offset += frameSize;
-    }
 
-    // Now check the actual compressed frame
-    if (offset >= srcSize) {
-        rb_raise(rb_eRuntimeError, "No compressed frame found in %zu bytes (only skippable frames)", srcSize);
-    }
+        // Now check the actual compressed frame
+        if (offset >= srcSize) {
+            rb_raise(rb_eRuntimeError, "No compressed frame found in %zu bytes (only skippable frames)", srcSize);
+        }
 
-    src += offset;
-    srcSize -= offset;
+        src += offset;
+        srcSize -= offset;
 
-    unsigned long long contentSize = ZSTD_getFrameContentSize(src, srcSize);
-    if (contentSize == ZSTD_CONTENTSIZE_ERROR) {
-        rb_raise(rb_eRuntimeError, "Invalid compressed data: not a valid zstd frame (size: %zu bytes)", srcSize);
-    }
+        contentSize = ZSTD_getFrameContentSize(src, srcSize);
+        if (contentSize == ZSTD_CONTENTSIZE_ERROR) {
+            rb_raise(rb_eRuntimeError, "Invalid compressed data: not a valid zstd frame (size: %zu bytes)", srcSize);
+        }
 
-    // Check dictionary requirements from the frame
-    unsigned int frame_dict_id = ZSTD_getDictID_fromFrame(src, srcSize);
+        // Check dictionary requirements from the frame
+        frame_dict_id = ZSTD_getDictID_fromFrame(src, srcSize);
+    }
 
     // Extract keyword arguments
     ZSTD_DDict* ddict = NULL;
     unsigned int provided_dict_id = 0;
     size_t initial_capacity = 0;  // 0 = not specified in per-call options
+    size_t max_size = 0;          // 0 = not specified in per-call options
 
     if (!NIL_P(options)) {
         VALUE dict_val = rb_hash_aref(options, ID2SYM(rb_intern("dict")));
@@ -399,6 +560,27 @@ vibe_zstd_dctx_decompress(int argc, VALUE* argv, VALUE self) {
                 rb_raise(rb_eArgError, "initial_capacity must be positive");
             }
         }
+
+        // Per-call output-size limit; accepts :max_decompressed_size or :max_size.
+        VALUE max_size_val = rb_hash_aref(options, ID2SYM(rb_intern("max_decompressed_size")));
+        if (NIL_P(max_size_val)) {
+            max_size_val = rb_hash_aref(options, ID2SYM(rb_intern("max_size")));
+        }
+        if (!NIL_P(max_size_val)) {
+            max_size = NUM2SIZET(max_size_val);
+            if (max_size == 0) {
+                rb_raise(rb_eArgError, "max_decompressed_size must be positive");
+            }
+        }
+    }
+
+    // Resolve max_size fallback chain: per-call > instance > class default.
+    // A value of 0 at every level means unlimited.
+    if (max_size == 0) {
+        max_size = dctx->max_decompressed_size;  // instance
+        if (max_size == 0) {
+            max_size = default_max_decompressed_size;  // class
+        }
     }
 
     // Validate dictionary matches frame requirements
@@ -426,6 +608,17 @@ vibe_zstd_dctx_decompress(int argc, VALUE* argv, VALUE self) {
     // Releases GVL to allow other Ruby threads to run during decompression.
     // Uses C malloc/realloc (not Ruby allocators) since Ruby API calls are forbidden without GVL.
     if (contentSize == ZSTD_CONTENTSIZE_UNKNOWN) {
+        // Reference the dictionary on the context before streaming decompression.
+        // ZSTD_decompressStream uses whatever dict is referenced on the DCtx, so
+        // without this the dictionary would be ignored on the unknown-size path
+        // (every dict frame produced by CompressWriter has unknown content size).
+        if (ddict) {
+            size_t rd = ZSTD_DCtx_refDDict(dctx->dctx, ddict);
+            if (ZSTD_isError(rd)) {
+                rb_raise(rb_eRuntimeError, "Failed to reference dictionary: %s", ZSTD_getErrorName(rd));
+            }
+        }
+
         decompress_stream_nogvl_args stream_args = {
             .dctx = dctx->dctx,
             .src = src,
@@ -434,22 +627,34 @@ vibe_zstd_dctx_decompress(int argc, VALUE* argv, VALUE self) {
             .dst_capacity = 0,
             .dst_size = 0,
             .initial_capacity = initial_capacity,
+            .max_size = max_size,
             .error = 0,
+            .limit_exceeded = 0,
+            .truncated = 0,
             .error_name = NULL
         };
 
-        rb_thread_call_without_gvl(decompress_stream_without_gvl, &stream_args, NULL, NULL);
-
-        if (stream_args.error) {
-            if (stream_args.dst) free(stream_args.dst);
-            rb_raise(rb_eRuntimeError, "Decompression failed: %s", stream_args.error_name);
-        }
-
-        // Create Ruby string from the C buffer, then free the C buffer
-        VALUE result = rb_str_new(stream_args.dst, stream_args.dst_size);
-        free(stream_args.dst);
-        return result;
+        // Run the streaming decompression and build the result under rb_ensure:
+        // the cleanup frees the C buffer and un-references the dictionary on
+        // every exit path, including the raises below and async exceptions
+        // delivered when the GVL is reacquired.
+        dctx_stream_decompress_state state = {
+            .dctx = dctx->dctx,
+            .ddict = ddict,
+            .args = &stream_args,
+            .data = data,
+            .max_size = max_size
+        };
+        return rb_ensure(vibe_zstd_dctx_stream_decompress_body, (VALUE)&state,
+                         vibe_zstd_dctx_stream_decompress_cleanup, (VALUE)&state);
     }
+    // Reject a frame whose declared content size exceeds the limit before
+    // allocating the output buffer (the header is attacker-controlled).
+    if (max_size && contentSize > (unsigned long long)max_size) {
+        rb_raise(rb_eDecompressedSizeExceeded,
+                 "Declared content size %llu exceeds limit of %zu bytes", contentSize, max_size);
+    }
+
     VALUE result = rb_str_new(NULL, contentSize);
     decompress_args args = {
         .dctx = dctx->dctx,
@@ -460,7 +665,11 @@ vibe_zstd_dctx_decompress(int argc, VALUE* argv, VALUE self) {
         .dstCapacity = contentSize,
         .result = 0
     };
-    rb_thread_call_without_gvl(decompress_without_gvl, &args, NULL, NULL);
+    // Lock the source string while the GVL is released: another Ruby thread
+    // holding the same string must not mutate or GC it mid-decompression.
+    // The helper unlocks via rb_ensure so an async exception cannot leave
+    // the string permanently locked.
+    vibe_zstd_nogvl_with_str_locked(decompress_without_gvl, &args, data);
     if (ZSTD_isError(args.result)) {
         rb_raise(rb_eRuntimeError, "Decompression failed: %s", ZSTD_getErrorName(args.result));
     }
@@ -525,6 +734,13 @@ vibe_zstd_dctx_init_class(VALUE rb_cVibeZstdDCtx) {
     // Initialize parameter lookup table
     init_dctx_param_table();
 
+    // Define VibeZstd::Error (base) and VibeZstd::DecompressedSizeExceeded.
+    // Defined here (rather than only in Ruby) so the error is available even if
+    // the C extension is required without the Ruby wrapper. Ruby's
+    // `class Error < StandardError` simply reopens the same class.
+    VALUE rb_eVibeZstdError = rb_define_class_under(rb_mVibeZstd, "Error", rb_eStandardError);
+    rb_eDecompressedSizeExceeded = rb_define_class_under(rb_mVibeZstd, "DecompressedSizeExceeded", rb_eVibeZstdError);
+
     rb_define_alloc_func(rb_cVibeZstdDCtx, vibe_zstd_dctx_alloc);
     rb_define_method(rb_cVibeZstdDCtx, "initialize", vibe_zstd_dctx_initialize, -1);
     rb_define_method(rb_cVibeZstdDCtx, "decompress", vibe_zstd_dctx_decompress, -1);
@@ -543,8 +759,20 @@ vibe_zstd_dctx_init_class(VALUE rb_cVibeZstdDCtx) {
     rb_define_method(rb_cVibeZstdDCtx, "window_log_max", vibe_zstd_dctx_get_window_log_max, 0);
     rb_define_alias(rb_cVibeZstdDCtx, "max_window_log=", "window_log_max=");
     rb_define_alias(rb_cVibeZstdDCtx, "max_window_log", "window_log_max");
+    rb_define_method(rb_cVibeZstdDCtx, "format=", vibe_zstd_dctx_set_format, 1);
+    rb_define_method(rb_cVibeZstdDCtx, "format", vibe_zstd_dctx_get_format, 0);
 
     // Instance-level initial_capacity accessors
     rb_define_method(rb_cVibeZstdDCtx, "initial_capacity", vibe_zstd_dctx_get_initial_capacity, 0);
     rb_define_method(rb_cVibeZstdDCtx, "initial_capacity=", vibe_zstd_dctx_set_initial_capacity, 1);
+
+    // Class-level default_max_decompressed_size accessors (0 = unlimited)
+    rb_define_singleton_method(rb_cVibeZstdDCtx, "default_max_decompressed_size", vibe_zstd_dctx_get_default_max_decompressed_size, 0);
+    rb_define_singleton_method(rb_cVibeZstdDCtx, "default_max_decompressed_size=", vibe_zstd_dctx_set_default_max_decompressed_size, 1);
+
+    // Instance-level max_decompressed_size accessors (with shorter max_size alias)
+    rb_define_method(rb_cVibeZstdDCtx, "max_decompressed_size", vibe_zstd_dctx_get_max_decompressed_size, 0);
+    rb_define_method(rb_cVibeZstdDCtx, "max_decompressed_size=", vibe_zstd_dctx_set_max_decompressed_size, 1);
+    rb_define_alias(rb_cVibeZstdDCtx, "max_size", "max_decompressed_size");
+    rb_define_alias(rb_cVibeZstdDCtx, "max_size=", "max_decompressed_size=");
 }

data/ext/vibe_zstd/depend

@@ -0,0 +1,3 @@
+# vibe_zstd.c textually #includes the split implementation files, so the object
+# must be rebuilt when any of them (or the project headers) change.
+vibe_zstd.o: cctx.c dctx.c dict.c streaming.c frames.c vibe_zstd.h vibe_zstd_internal.h

data/ext/vibe_zstd/dict.c

@@ -125,13 +125,21 @@ dict_training_cleanup(VALUE arg) {
     return Qnil;
 }
 
-// Copy Ruby sample strings into contiguous C buffer for ZDICT functions
+// Copy Ruby sample strings into contiguous C buffer for ZDICT functions.
+// capacity is the total_samples_size measured during the validation pass.
+// If any sample has grown since validation (TOCTOU mutation), we raise rather
+// than overflow the buffer.
 static void
-copy_samples_to_buffer(dict_training_resources* resources, VALUE samples, long num_samples) {
+copy_samples_to_buffer(dict_training_resources* resources, VALUE samples, long num_samples,
+                       size_t capacity) {
     size_t offset = 0;
     for (long i = 0; i < num_samples; i++) {
         VALUE sample = rb_ary_entry(samples, i);
         size_t sample_len = RSTRING_LEN(sample);
+        // Guard against mutation between the validation pass and this copy.
+        if (sample_len > capacity - offset) {
+            rb_raise(rb_eRuntimeError, "sample mutated during dictionary training");
+        }
         resources->sample_sizes[i] = sample_len;
         memcpy(resources->samples_buffer + offset, RSTRING_PTR(sample), sample_len);
         offset += sample_len;
@@ -148,12 +156,14 @@ typedef struct {
     VALUE result;
     size_t max_dict_size;
     long num_samples;
-    VALUE samples;
+    size_t total_samples_size;  // measured during validation; passed to copy_samples_to_buffer
+    VALUE samples;              // private converted-samples array built during validation
 } dict_training_ctx;
 
 static VALUE train_dict_basic_body(VALUE arg) {
     dict_training_ctx* ctx = (dict_training_ctx*)arg;
-    copy_samples_to_buffer(ctx->resources, ctx->samples, ctx->num_samples);
+    copy_samples_to_buffer(ctx->resources, ctx->samples, ctx->num_samples,
+                           ctx->total_samples_size);
     size_t dict_size = ZDICT_trainFromBuffer(
         ctx->resources->dict_buffer, ctx->max_dict_size,
         ctx->resources->samples_buffer, ctx->resources->sample_sizes, (unsigned)ctx->num_samples
@@ -172,7 +182,8 @@ typedef struct {
 
 static VALUE train_dict_cover_body(VALUE arg) {
     train_dict_cover_ctx* ctx = (train_dict_cover_ctx*)arg;
-    copy_samples_to_buffer(ctx->base.resources, ctx->base.samples, ctx->base.num_samples);
+    copy_samples_to_buffer(ctx->base.resources, ctx->base.samples, ctx->base.num_samples,
+                           ctx->base.total_samples_size);
     size_t dict_size = ZDICT_trainFromBuffer_cover(
         ctx->base.resources->dict_buffer, ctx->base.max_dict_size,
         ctx->base.resources->samples_buffer, ctx->base.resources->sample_sizes, (unsigned)ctx->base.num_samples,
@@ -192,7 +203,8 @@ typedef struct {
 
 static VALUE train_dict_fast_cover_body(VALUE arg) {
     train_dict_fast_cover_ctx* ctx = (train_dict_fast_cover_ctx*)arg;
-    copy_samples_to_buffer(ctx->base.resources, ctx->base.samples, ctx->base.num_samples);
+    copy_samples_to_buffer(ctx->base.resources, ctx->base.samples, ctx->base.num_samples,
+                           ctx->base.total_samples_size);
     size_t dict_size = ZDICT_trainFromBuffer_fastCover(
         ctx->base.resources->dict_buffer, ctx->base.max_dict_size,
         ctx->base.resources->samples_buffer, ctx->base.resources->sample_sizes, (unsigned)ctx->base.num_samples,
@@ -213,7 +225,8 @@ typedef struct {
 
 static VALUE finalize_dict_body(VALUE arg) {
     finalize_dict_ctx* ctx = (finalize_dict_ctx*)arg;
-    copy_samples_to_buffer(ctx->base.resources, ctx->base.samples, ctx->base.num_samples);
+    copy_samples_to_buffer(ctx->base.resources, ctx->base.samples, ctx->base.num_samples,
+                           ctx->base.total_samples_size);
     size_t dict_size = ZDICT_finalizeDictionary(
         ctx->base.resources->dict_buffer, ctx->base.max_dict_size,
         RSTRING_PTR(ctx->content_val), RSTRING_LEN(ctx->content_val),
@@ -245,11 +258,17 @@ vibe_zstd_train_dict(int argc, VALUE* argv, VALUE self) {
         rb_raise(rb_eArgError, "samples array cannot be empty");
     }
 
-    // Validate all samples are strings and calculate sizes BEFORE allocating
+    // Validate all samples are strings and calculate sizes BEFORE allocating.
+    // Build a private converted-samples array so that StringValue conversions are
+    // retained: rb_ary_entry re-fetches the raw element, so storing the converted
+    // value here ensures copy_samples_to_buffer operates on real String objects
+    // rather than potentially-non-String objects that merely respond to to_str.
+    VALUE converted_samples = rb_ary_new_capa(num_samples);
     size_t total_samples_size = 0;
     for (long i = 0; i < num_samples; i++) {
         VALUE sample = rb_ary_entry(samples, i);
-        StringValue(sample);  // Validate type early - may raise TypeError
+        StringValue(sample);  // Validate type early - may raise TypeError; updates local
+        rb_ary_push(converted_samples, sample);
         total_samples_size += RSTRING_LEN(sample);
     }
 
@@ -274,7 +293,8 @@ vibe_zstd_train_dict(int argc, VALUE* argv, VALUE self) {
         .result = Qnil,
         .max_dict_size = max_dict_size,
         .num_samples = num_samples,
-        .samples = samples
+        .total_samples_size = total_samples_size,
+        .samples = converted_samples  // use private array, not caller's array
     };
 
     rb_ensure(train_dict_basic_body, (VALUE)&ctx, dict_training_cleanup, (VALUE)&resources);
@@ -298,11 +318,14 @@ vibe_zstd_train_dict_cover(int argc, VALUE* argv, VALUE self) {
         rb_raise(rb_eArgError, "samples array cannot be empty");
     }
 
-    // Validate all samples are strings and calculate sizes BEFORE allocating
+    // Validate all samples are strings and calculate sizes BEFORE allocating.
+    // Build a private converted-samples array (see vibe_zstd_train_dict for details).
+    VALUE converted_samples = rb_ary_new_capa(num_samples);
     size_t total_samples_size = 0;
     for (long i = 0; i < num_samples; i++) {
         VALUE sample = rb_ary_entry(samples, i);
-        StringValue(sample);  // Validate type early - may raise TypeError
+        StringValue(sample);  // Validate type early - may raise TypeError; updates local
+        rb_ary_push(converted_samples, sample);
         total_samples_size += RSTRING_LEN(sample);
     }
 
@@ -358,7 +381,8 @@ vibe_zstd_train_dict_cover(int argc, VALUE* argv, VALUE self) {
             .result = Qnil,
             .max_dict_size = max_dict_size,
             .num_samples = num_samples,
-            .samples = samples
+            .total_samples_size = total_samples_size,
+            .samples = converted_samples  // use private array, not caller's array
         },
         .params = params
     };
@@ -384,11 +408,14 @@ vibe_zstd_train_dict_fast_cover(int argc, VALUE* argv, VALUE self) {
         rb_raise(rb_eArgError, "samples array cannot be empty");
     }
 
-    // Validate all samples are strings and calculate sizes BEFORE allocating
+    // Validate all samples are strings and calculate sizes BEFORE allocating.
+    // Build a private converted-samples array (see vibe_zstd_train_dict for details).
+    VALUE converted_samples = rb_ary_new_capa(num_samples);
     size_t total_samples_size = 0;
     for (long i = 0; i < num_samples; i++) {
         VALUE sample = rb_ary_entry(samples, i);
-        StringValue(sample);  // Validate type early - may raise TypeError
+        StringValue(sample);  // Validate type early - may raise TypeError; updates local
+        rb_ary_push(converted_samples, sample);
         total_samples_size += RSTRING_LEN(sample);
     }
 
@@ -447,7 +474,8 @@ vibe_zstd_train_dict_fast_cover(int argc, VALUE* argv, VALUE self) {
             .result = Qnil,
             .max_dict_size = max_dict_size,
             .num_samples = num_samples,
-            .samples = samples
+            .total_samples_size = total_samples_size,
+            .samples = converted_samples  // use private array, not caller's array
         },
         .params = params
     };
@@ -514,11 +542,14 @@ vibe_zstd_finalize_dictionary(int argc, VALUE* argv, VALUE self) {
         rb_raise(rb_eArgError, "samples array cannot be empty");
     }
 
-    // Validate all samples are strings and calculate sizes BEFORE allocating
+    // Validate all samples are strings and calculate sizes BEFORE allocating.
+    // Build a private converted-samples array (see vibe_zstd_train_dict for details).
+    VALUE converted_samples = rb_ary_new_capa(num_samples);
     size_t total_samples_size = 0;
     for (long i = 0; i < num_samples; i++) {
         VALUE sample = rb_ary_entry(samples_val, i);
-        StringValue(sample);  // Validate type early - may raise TypeError
+        StringValue(sample);  // Validate type early - may raise TypeError; updates local
+        rb_ary_push(converted_samples, sample);
         total_samples_size += RSTRING_LEN(sample);
     }
 
@@ -546,7 +577,8 @@ vibe_zstd_finalize_dictionary(int argc, VALUE* argv, VALUE self) {
             .result = Qnil,
             .max_dict_size = max_size,
             .num_samples = num_samples,
-            .samples = samples_val
+            .total_samples_size = total_samples_size,
+            .samples = converted_samples  // use private array, not caller's array
         },
         .content_val = content_val,
         .params = params

data/ext/vibe_zstd/extconf.rb

@@ -14,10 +14,13 @@ $INCFLAGS << " -I#{LIBZSTD_DIR}/decompress"
 $INCFLAGS << " -I#{LIBZSTD_DIR}/dictBuilder"
 # standard:enable Style/GlobalVars
 
-# Add preprocessor definitions
-append_cflags("-DXXH_NAMESPACE=ZSTD_")
-append_cflags("-DZSTD_LEGACY_SUPPORT=0") # Disable legacy support to reduce size
-append_cflags("-DZSTD_MULTITHREAD") # Enable multithreading support
+# Add preprocessor definitions (use $defs so they appear in DEFS in the Makefile,
+# append_cflags only validates the flag but doesn't reliably propagate -D flags)
+# standard:disable Style/GlobalVars
+$defs << "-DXXH_NAMESPACE=ZSTD_"
+$defs << "-DZSTD_LEGACY_SUPPORT=0" # Disable legacy support to reduce size
+$defs << "-DZSTD_MULTITHREAD" # Enable multithreading support
+# standard:enable Style/GlobalVars
 
 # Link with pthread for multithreading
 have_library("pthread") || abort("pthread library is required for multithreading support")

data/ext/vibe_zstd/frames.c

@@ -81,12 +81,23 @@ vibe_zstd_read_skippable_frame(VALUE self, VALUE data) {
     uint32_t content_size;
     memcpy(&content_size, src + 4, 4);
 
-    VALUE result = rb_str_buf_new(content_size);
+    // The content size field is attacker-controlled and may claim up to ~4 GiB.
+    // A skippable frame's content cannot exceed the bytes actually provided
+    // (src_size minus the 8-byte header), so cap the allocation accordingly to
+    // prevent a tiny truncated input from forcing a huge allocation. A frame
+    // whose declared size exceeds what is present is malformed and
+    // ZSTD_readSkippableFrame reports the error below.
+    size_t capacity = content_size;
+    if (capacity > src_size - 8) {
+        capacity = src_size - 8;
+    }
+
+    VALUE result = rb_str_buf_new(capacity);
     unsigned magic_variant;
 
     size_t bytes_read = ZSTD_readSkippableFrame(
         RSTRING_PTR(result),
-        content_size,
+        capacity,
         &magic_variant,
         src,
         src_size