From 8f4de31adf89e752f4ae6f118f3e8ba98b2eee45 Mon Sep 17 00:00:00 2001 From: Shizuo Fujita Date: Thu, 3 Oct 2024 12:02:58 +0900 Subject: [PATCH] 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 ``` --- ext/msgpack/unpacker.c | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/ext/msgpack/unpacker.c b/ext/msgpack/unpacker.c index 4b05443d..4fd5942d 100644 --- a/ext/msgpack/unpacker.c +++ b/ext/msgpack/unpacker.c @@ -20,6 +20,7 @@ #include "rmem.h" #include "extension_value_class.h" #include +#include #if !defined(HAVE_RB_PROC_CALL_WITH_BLOCK) #define rb_proc_call_with_block(recv, argc, argv, block) rb_funcallv(recv, rb_intern("call"), argc, argv) @@ -27,6 +28,7 @@ static int RAW_TYPE_STRING = 256; static int RAW_TYPE_BINARY = 257; +static int16_t INITIAL_BUFFER_CAPACITY_MAX = SHRT_MAX; static msgpack_rmem_t s_stack_rmem; @@ -37,6 +39,11 @@ static inline VALUE rb_hash_new_capa(long capa) } #endif +static inline int16_t initial_buffer_size(long size) +{ + return (size > INITIAL_BUFFER_CAPACITY_MAX) ? INITIAL_BUFFER_CAPACITY_MAX : size; +} + void msgpack_unpacker_static_init(void) { 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) if(count == 0) { return object_complete(uk, rb_ary_new()); } - return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(count)); + return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(initial_buffer_size(count))); SWITCH_RANGE(b, 0x80, 0x8f) // FixMap int count = b & 0x0f; if(count == 0) { return object_complete(uk, rb_hash_new()); } - return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(count)); + return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(initial_buffer_size(count))); SWITCH_RANGE(b, 0xc0, 0xdf) // Variable switch(b) { @@ -605,7 +612,7 @@ static int read_primitive(msgpack_unpacker_t* uk) if(count == 0) { return object_complete(uk, rb_ary_new()); } - return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(count)); + return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(initial_buffer_size(count))); } case 0xdd: // array 32 @@ -615,7 +622,7 @@ static int read_primitive(msgpack_unpacker_t* uk) if(count == 0) { return object_complete(uk, rb_ary_new()); } - return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(count)); + return _msgpack_unpacker_stack_push(uk, STACK_TYPE_ARRAY, count, rb_ary_new2(initial_buffer_size(count))); } case 0xde: // map 16 @@ -625,7 +632,7 @@ static int read_primitive(msgpack_unpacker_t* uk) if(count == 0) { return object_complete(uk, rb_hash_new()); } - return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(count)); + return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(initial_buffer_size(count))); } case 0xdf: // map 32 @@ -635,7 +642,7 @@ static int read_primitive(msgpack_unpacker_t* uk) if(count == 0) { return object_complete(uk, rb_hash_new()); } - return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(count)); + return _msgpack_unpacker_stack_push(uk, STACK_TYPE_MAP_KEY, count*2, rb_hash_new_capa(initial_buffer_size(count))); } default: