Skip to content

Commit

Permalink
Add upper limit on initial buffer size in MessagePack::Unpacker
Browse files Browse the repository at this point in the history
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
```
  • Loading branch information
Watson1978 committed Oct 3, 2024
1 parent 9330593 commit 8f4de31
Showing 1 changed file with 13 additions and 6 deletions.
19 changes: 13 additions & 6 deletions ext/msgpack/unpacker.c
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,15 @@
#include "rmem.h"
#include "extension_value_class.h"
#include <assert.h>
#include <limits.h>

#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)
#endif

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;

Expand All @@ -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);
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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:
Expand Down

0 comments on commit 8f4de31

Please sign in to comment.