Skip to content

Commit 9bac145

Browse files
byrootbyroot
authored
Merge pull request msgpack#371 from Shopify/recursive-raise-leak
Prevent memory leak when a recursive unpacker raises an exception
2 parents 6bbaa97 + 8a79706 commit 9bac145

File tree

6 files changed

+108
-16
lines changed

6 files changed

+108
-16
lines changed

‎ChangeLog‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
* Fixed a potental memory leak when recursive unpacker raise.
2+
13
2024-10-03 1.7.3
24

35
* Limit initial containers pre-allocation to `SHRT_MAX` (32k) entries.

‎ext/msgpack/unpacker.c‎

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,27 @@
2626
#define rb_proc_call_with_block(recv, argc, argv, block) rb_funcallv(recv, rb_intern("call"), argc, argv)
2727
#endif
2828

29+
struct protected_proc_call_args {
30+
VALUE proc;
31+
int argc;
32+
VALUE *argv;
33+
};
34+
35+
static VALUE protected_proc_call_safe(VALUE _args) {
36+
struct protected_proc_call_args *args = (struct protected_proc_call_args *)_args;
37+
38+
return rb_proc_call_with_block(args->proc, args->argc, args->argv, Qnil);
39+
}
40+
41+
static VALUE protected_proc_call(VALUE proc, int argc, VALUE *argv, int *raised) {
42+
struct protected_proc_call_args args = {
43+
.proc = proc,
44+
.argc = argc,
45+
.argv = argv,
46+
};
47+
return rb_protect(protected_proc_call_safe, (VALUE)&args, raised);
48+
}
49+
2950
static int RAW_TYPE_STRING = 256;
3051
static int RAW_TYPE_BINARY = 257;
3152
static int16_t INITIAL_BUFFER_CAPACITY_MAX = SHRT_MAX;
@@ -87,7 +108,12 @@ static inline void _msgpack_unpacker_free_stack(msgpack_unpacker_stack_t* stack)
87108

88109
void _msgpack_unpacker_destroy(msgpack_unpacker_t* uk)
89110
{
90-
_msgpack_unpacker_free_stack(uk->stack);
111+
msgpack_unpacker_stack_t *stack;
112+
while ((stack = uk->stack)) {
113+
uk->stack = stack->parent;
114+
_msgpack_unpacker_free_stack(stack);
115+
}
116+
91117
msgpack_buffer_destroy(UNPACKER_BUFFER_(uk));
92118
}
93119

@@ -186,7 +212,12 @@ static inline int object_complete_ext(msgpack_unpacker_t* uk, int ext_type, VALU
186212
if(proc != Qnil) {
187213
VALUE obj;
188214
VALUE arg = (str == Qnil ? rb_str_buf_new(0) : str);
189-
obj = rb_proc_call_with_block(proc, 1, &arg, Qnil);
215+
int raised;
216+
obj = protected_proc_call(proc, 1, &arg, &raised);
217+
if (raised) {
218+
uk->last_object = rb_errinfo();
219+
return PRIMITIVE_RECURSIVE_RAISED;
220+
}
190221
return object_complete(uk, obj);
191222
}
192223

@@ -316,11 +347,16 @@ static inline int read_raw_body_begin(msgpack_unpacker_t* uk, int raw_type)
316347
child_stack->parent = uk->stack;
317348
uk->stack = child_stack;
318349

319-
obj = rb_proc_call_with_block(proc, 1, &uk->self, Qnil);
320-
350+
int raised;
351+
obj = protected_proc_call(proc, 1, &uk->self, &raised);
321352
uk->stack = child_stack->parent;
322353
_msgpack_unpacker_free_stack(child_stack);
323354

355+
if (raised) {
356+
uk->last_object = rb_errinfo();
357+
return PRIMITIVE_RECURSIVE_RAISED;
358+
}
359+
324360
return object_complete(uk, obj);
325361
}
326362
}

‎ext/msgpack/unpacker.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ static inline void msgpack_unpacker_set_allow_unknown_ext(msgpack_unpacker_t* uk
119119
#define PRIMITIVE_STACK_TOO_DEEP -3
120120
#define PRIMITIVE_UNEXPECTED_TYPE -4
121121
#define PRIMITIVE_UNEXPECTED_EXT_TYPE -5
122+
#define PRIMITIVE_RECURSIVE_RAISED -6
122123

123124
int msgpack_unpacker_read(msgpack_unpacker_t* uk, size_t target_stack_depth);
124125

‎ext/msgpack/unpacker_class.c‎

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,11 @@ static size_t Unpacker_memsize(const void *ptr)
6565
total_size += sizeof(msgpack_unpacker_ext_registry_t) / (uk->ext_registry->borrow_count + 1);
6666
}
6767

68-
total_size += (uk->stack->depth + 1) * sizeof(msgpack_unpacker_stack_t);
68+
msgpack_unpacker_stack_t *stack = uk->stack;
69+
while (stack) {
70+
total_size += (stack->depth + 1) * sizeof(msgpack_unpacker_stack_t);
71+
stack = stack->parent;
72+
}
6973

7074
return total_size + msgpack_buffer_memsize(&uk->buffer);
7175
}
@@ -156,20 +160,28 @@ static VALUE Unpacker_allow_unknown_ext_p(VALUE self)
156160
return uk->allow_unknown_ext ? Qtrue : Qfalse;
157161
}
158162

159-
NORETURN(static void raise_unpacker_error(int r))
163+
NORETURN(static void raise_unpacker_error(msgpack_unpacker_t *uk, int r))
160164
{
165+
uk->stack->depth = 0;
161166
switch(r) {
162167
case PRIMITIVE_EOF:
163168
rb_raise(rb_eEOFError, "end of buffer reached");
169+
break;
164170
case PRIMITIVE_INVALID_BYTE:
165171
rb_raise(eMalformedFormatError, "invalid byte");
172+
break;
166173
case PRIMITIVE_STACK_TOO_DEEP:
167174
rb_raise(eStackError, "stack level too deep");
175+
break;
168176
case PRIMITIVE_UNEXPECTED_TYPE:
169177
rb_raise(eUnexpectedTypeError, "unexpected type");
178+
break;
170179
case PRIMITIVE_UNEXPECTED_EXT_TYPE:
171-
// rb_bug("unexpected extension type");
172180
rb_raise(eUnknownExtTypeError, "unexpected extension type");
181+
break;
182+
case PRIMITIVE_RECURSIVE_RAISED:
183+
rb_exc_raise(msgpack_unpacker_get_last_object(uk));
184+
break;
173185
default:
174186
rb_raise(eUnpackError, "logically unknown error %d", r);
175187
}
@@ -190,7 +202,7 @@ static VALUE Unpacker_read(VALUE self)
190202

191203
int r = msgpack_unpacker_read(uk, 0);
192204
if(r < 0) {
193-
raise_unpacker_error(r);
205+
raise_unpacker_error(uk, r);
194206
}
195207

196208
return msgpack_unpacker_get_last_object(uk);
@@ -202,7 +214,7 @@ static VALUE Unpacker_skip(VALUE self)
202214

203215
int r = msgpack_unpacker_skip(uk, 0);
204216
if(r < 0) {
205-
raise_unpacker_error(r);
217+
raise_unpacker_error(uk, r);
206218
}
207219

208220
return Qnil;
@@ -214,7 +226,7 @@ static VALUE Unpacker_skip_nil(VALUE self)
214226

215227
int r = msgpack_unpacker_skip_nil(uk);
216228
if(r < 0) {
217-
raise_unpacker_error(r);
229+
raise_unpacker_error(uk, r);
218230
}
219231

220232
if(r) {
@@ -230,7 +242,7 @@ static VALUE Unpacker_read_array_header(VALUE self)
230242
uint32_t size;
231243
int r = msgpack_unpacker_read_array_header(uk, &size);
232244
if(r < 0) {
233-
raise_unpacker_error(r);
245+
raise_unpacker_error(uk, r);
234246
}
235247

236248
return ULONG2NUM(size); // long at least 32 bits
@@ -243,7 +255,7 @@ static VALUE Unpacker_read_map_header(VALUE self)
243255
uint32_t size;
244256
int r = msgpack_unpacker_read_map_header(uk, &size);
245257
if(r < 0) {
246-
raise_unpacker_error((int)r);
258+
raise_unpacker_error(uk, r);
247259
}
248260

249261
return ULONG2NUM(size); // long at least 32 bits
@@ -270,7 +282,7 @@ static VALUE Unpacker_each_impl(VALUE self)
270282
if(r == PRIMITIVE_EOF) {
271283
return Qnil;
272284
}
273-
raise_unpacker_error(r);
285+
raise_unpacker_error(uk, r);
274286
}
275287
VALUE v = msgpack_unpacker_get_last_object(uk);
276288
#ifdef JRUBY
@@ -369,7 +381,7 @@ static VALUE Unpacker_full_unpack(VALUE self)
369381

370382
int r = msgpack_unpacker_read(uk, 0);
371383
if(r < 0) {
372-
raise_unpacker_error(r);
384+
raise_unpacker_error(uk, r);
373385
}
374386

375387
/* raise if extra bytes follow */

‎spec/spec_helper.rb‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@
2222
# This method was added in Ruby 3.0.0. Calling it this way asks the GC to
2323
# move objects around, helping to find object movement bugs.
2424
begin
25-
GC.verify_compaction_references(double_heap: true, toward: :empty)
26-
rescue NotImplementedError
25+
GC.verify_compaction_references(expand_heap: true, toward: :empty)
26+
rescue NotImplementedError, ArgumentError
2727
# Some platforms don't support compaction
2828
end
2929
end

‎spec/unpacker_spec.rb‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -901,4 +901,45 @@ def flatten(struct, results = [])
901901
GC.stress = stress
902902
end
903903
end
904+
905+
if RUBY_PLATFORM != "java"
906+
it "doesn't leak when a recursive unpacker raises" do
907+
hash_with_indifferent_access = Class.new(Hash)
908+
msgpack = MessagePack::Factory.new
909+
msgpack.register_type(
910+
0x02,
911+
hash_with_indifferent_access,
912+
packer: ->(value, packer) do
913+
packer.write(value.to_h)
914+
end,
915+
unpacker: ->(unpacker) { raise RuntimeError, "Ooops" },
916+
recursive: true
917+
)
918+
919+
packer = msgpack.packer
920+
data = [[[[[[[hash_with_indifferent_access.new]]]]]]]
921+
payload = msgpack.dump(data)
922+
923+
unpacker = msgpack.unpacker
924+
2.times do
925+
unpacker.buffer.clear
926+
unpacker.feed(payload)
927+
expect {
928+
unpacker.full_unpack
929+
}.to raise_error(RuntimeError, "Ooops")
930+
end
931+
932+
memsize = ObjectSpace.memsize_of(unpacker)
933+
934+
10.times do
935+
unpacker.buffer.clear
936+
unpacker.feed(payload)
937+
expect {
938+
unpacker.full_unpack
939+
}.to raise_error(RuntimeError, "Ooops")
940+
end
941+
942+
expect(memsize).to eq ObjectSpace.memsize_of(unpacker)
943+
end
944+
end
904945
end

0 commit comments

Comments
 (0)