Skip to content

Commit 8f4de31

Browse files
committed
Add upper limit on initial buffer size in MessagePack::Unpacker
Currently, the initial buffer size is specified in rb_ary_new2() or rb_hash_new_capa(). If a huge size is specified, a large amount of memory is allocated and system memory might be depleted. We want to unpack the data received over the network. However the service may stop due to large amount of memory allocation with crafted data. So this patch add upper limit on initial buffer size. If the buffer runs out, Ruby API will be reallocated automatically. ## Test code ```ruby require "msgpack" puts "msgpack version: #{MessagePack::VERSION}" unpacker = MessagePack::Unpacker.new unpacker.feed_each("\xDF\x20\x00\x00\x00") {} puts "Memory Usage: #{`ps -o rss= -p #{Process.pid}`.strip} KB" ``` ## Before Before it apply this patch, it allocates 8 GB memory on my environment. ``` $ ruby -v test.rb ruby 3.3.5 (2024-09-03 revision ef084cc8f4) [x86_64-linux] msgpack version: 1.7.2 Memory Usage: 8403320 KB ``` ## After ``` ruby -v test.rb ruby 3.3.5 (2024-09-03 revision ef084cc8f4) [x86_64-linux] msgpack version: 1.7.2 Memory Usage: 14480 KB ```
1 parent 9330593 commit 8f4de31

File tree

1 file changed

+13
-6
lines changed

1 file changed

+13
-6
lines changed

ext/msgpack/unpacker.c

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,15 @@
2020
#include "rmem.h"
2121
#include "extension_value_class.h"
2222
#include <assert.h>
23+
#include <limits.h>
2324

2425
#if !defined(HAVE_RB_PROC_CALL_WITH_BLOCK)
2526
#define rb_proc_call_with_block(recv, argc, argv, block) rb_funcallv(recv, rb_intern("call"), argc, argv)
2627
#endif
2728

2829
static int RAW_TYPE_STRING = 256;
2930
static int RAW_TYPE_BINARY = 257;
31+
static int16_t INITIAL_BUFFER_CAPACITY_MAX = SHRT_MAX;
3032

3133
static msgpack_rmem_t s_stack_rmem;
3234

@@ -37,6 +39,11 @@ static inline VALUE rb_hash_new_capa(long capa)
3739
}
3840
#endif
3941

42+
static inline int16_t initial_buffer_size(long size)
43+
{
44+
return (size > INITIAL_BUFFER_CAPACITY_MAX) ? INITIAL_BUFFER_CAPACITY_MAX : size;
45+
}
46+
4047
void msgpack_unpacker_static_init(void)
4148
{
4249
assert(sizeof(msgpack_unpacker_stack_entry_t) * MSGPACK_UNPACKER_STACK_CAPACITY <= MSGPACK_RMEM_PAGE_SIZE);
@@ -375,14 +382,14 @@ static int read_primitive(msgpack_unpacker_t* uk)
375382
if(count == 0) {
376383
return object_complete(uk, rb_ary_new());
377384
}
378-
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(count));
385+
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(initial_buffer_size(count)));
379386

380387
SWITCH_RANGE(b, 0x80, 0x8f) // FixMap
381388
int count = b & 0x0f;
382389
if(count == 0) {
383390
return object_complete(uk, rb_hash_new());
384391
}
385-
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(count));
392+
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(initial_buffer_size(count)));
386393

387394
SWITCH_RANGE(b, 0xc0, 0xdf) // Variable
388395
switch(b) {
@@ -605,7 +612,7 @@ static int read_primitive(msgpack_unpacker_t* uk)
605612
if(count == 0) {
606613
return object_complete(uk, rb_ary_new());
607614
}
608-
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(count));
615+
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(initial_buffer_size(count)));
609616
}
610617

611618
case 0xdd: // array 32
@@ -615,7 +622,7 @@ static int read_primitive(msgpack_unpacker_t* uk)
615622
if(count == 0) {
616623
return object_complete(uk, rb_ary_new());
617624
}
618-
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(count));
625+
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(initial_buffer_size(count)));
619626
}
620627

621628
case 0xde: // map 16
@@ -625,7 +632,7 @@ static int read_primitive(msgpack_unpacker_t* uk)
625632
if(count == 0) {
626633
return object_complete(uk, rb_hash_new());
627634
}
628-
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(count));
635+
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(initial_buffer_size(count)));
629636
}
630637

631638
case 0xdf: // map 32
@@ -635,7 +642,7 @@ static int read_primitive(msgpack_unpacker_t* uk)
635642
if(count == 0) {
636643
return object_complete(uk, rb_hash_new());
637644
}
638-
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(count));
645+
return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(initial_buffer_size(count)));
639646
}
640647

641648
default:

0 commit comments

Comments
 (0)