Merge pull request #371 from Shopify/recursive-raise-leak

Prevent memory leak when a recursive unpacker raises an exception

Author: byroot

ChangeLog

@@ -1,3 +1,5 @@
+* Fixed a potental memory leak when recursive unpacker raise.
+
 2024-10-03 1.7.3
 
 * Limit initial containers pre-allocation to `SHRT_MAX` (32k) entries.

ext/msgpack/unpacker.c

@@ -26,6 +26,27 @@
 #define rb_proc_call_with_block(recv, argc, argv, block) rb_funcallv(recv, rb_intern("call"), argc, argv)
 #endif
 
+struct protected_proc_call_args {
+    VALUE proc;
+    int argc;
+    VALUE *argv;
+};
+
+static VALUE protected_proc_call_safe(VALUE _args) {
+    struct protected_proc_call_args *args = (struct protected_proc_call_args *)_args;
+
+    return rb_proc_call_with_block(args->proc, args->argc, args->argv, Qnil);
+}
+
+static VALUE protected_proc_call(VALUE proc, int argc, VALUE *argv, int *raised) {
+    struct protected_proc_call_args args = {
+      .proc = proc,
+      .argc = argc,
+      .argv = argv,
+    };
+    return rb_protect(protected_proc_call_safe, (VALUE)&args, raised);
+}
+
 static int RAW_TYPE_STRING = 256;
 static int RAW_TYPE_BINARY = 257;
 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)
 
 void _msgpack_unpacker_destroy(msgpack_unpacker_t* uk)
 {
-    _msgpack_unpacker_free_stack(uk->stack);
+    msgpack_unpacker_stack_t *stack;
+    while ((stack = uk->stack)) {
+        uk->stack = stack->parent;
+        _msgpack_unpacker_free_stack(stack);
+    }
+
     msgpack_buffer_destroy(UNPACKER_BUFFER_(uk));
 }
 
@@ -186,7 +212,12 @@ static inline int object_complete_ext(msgpack_unpacker_t* uk, int ext_type, VALU
     if(proc != Qnil) {
         VALUE obj;
         VALUE arg = (str == Qnil ? rb_str_buf_new(0) : str);
-        obj = rb_proc_call_with_block(proc, 1, &arg, Qnil);
+        int raised;
+        obj = protected_proc_call(proc, 1, &arg, &raised);
+        if (raised) {
+            uk->last_object = rb_errinfo();
+            return PRIMITIVE_RECURSIVE_RAISED;
+        }
         return object_complete(uk, obj);
     }
 
@@ -316,11 +347,16 @@ static inline int read_raw_body_begin(msgpack_unpacker_t* uk, int raw_type)
             child_stack->parent = uk->stack;
             uk->stack = child_stack;
 
-            obj = rb_proc_call_with_block(proc, 1, &uk->self, Qnil);
-
+            int raised;
+            obj = protected_proc_call(proc, 1, &uk->self, &raised);
             uk->stack = child_stack->parent;
             _msgpack_unpacker_free_stack(child_stack);
 
+            if (raised) {
+                uk->last_object = rb_errinfo();
+                return PRIMITIVE_RECURSIVE_RAISED;
+            }
+
             return object_complete(uk, obj);
         }
     }

ext/msgpack/unpacker.h

@@ -119,6 +119,7 @@ static inline void msgpack_unpacker_set_allow_unknown_ext(msgpack_unpacker_t* uk
 #define PRIMITIVE_STACK_TOO_DEEP -3
 #define PRIMITIVE_UNEXPECTED_TYPE -4
 #define PRIMITIVE_UNEXPECTED_EXT_TYPE -5
+#define PRIMITIVE_RECURSIVE_RAISED -6
 
 int msgpack_unpacker_read(msgpack_unpacker_t* uk, size_t target_stack_depth);
 

ext/msgpack/unpacker_class.c

@@ -65,7 +65,11 @@ static size_t Unpacker_memsize(const void *ptr)
         total_size += sizeof(msgpack_unpacker_ext_registry_t) / (uk->ext_registry->borrow_count + 1);
     }
 
-    total_size += (uk->stack->depth + 1) * sizeof(msgpack_unpacker_stack_t);
+    msgpack_unpacker_stack_t *stack = uk->stack;
+    while (stack) {
+        total_size += (stack->depth + 1) * sizeof(msgpack_unpacker_stack_t);
+        stack = stack->parent;
+    }
 
     return total_size + msgpack_buffer_memsize(&uk->buffer);
 }
@@ -156,20 +160,28 @@ static VALUE Unpacker_allow_unknown_ext_p(VALUE self)
     return uk->allow_unknown_ext ? Qtrue : Qfalse;
 }
 
-NORETURN(static void raise_unpacker_error(int r))
+NORETURN(static void raise_unpacker_error(msgpack_unpacker_t *uk, int r))
 {
+    uk->stack->depth = 0;
     switch(r) {
     case PRIMITIVE_EOF:
         rb_raise(rb_eEOFError, "end of buffer reached");
+        break;
     case PRIMITIVE_INVALID_BYTE:
         rb_raise(eMalformedFormatError, "invalid byte");
+        break;
     case PRIMITIVE_STACK_TOO_DEEP:
         rb_raise(eStackError, "stack level too deep");
+        break;
     case PRIMITIVE_UNEXPECTED_TYPE:
         rb_raise(eUnexpectedTypeError, "unexpected type");
+        break;
     case PRIMITIVE_UNEXPECTED_EXT_TYPE:
-    // rb_bug("unexpected extension type");
         rb_raise(eUnknownExtTypeError, "unexpected extension type");
+        break;
+    case PRIMITIVE_RECURSIVE_RAISED:
+        rb_exc_raise(msgpack_unpacker_get_last_object(uk));
+        break;
     default:
         rb_raise(eUnpackError, "logically unknown error %d", r);
     }
@@ -190,7 +202,7 @@ static VALUE Unpacker_read(VALUE self)
 
     int r = msgpack_unpacker_read(uk, 0);
     if(r < 0) {
-        raise_unpacker_error(r);
+        raise_unpacker_error(uk, r);
     }
 
     return msgpack_unpacker_get_last_object(uk);
@@ -202,7 +214,7 @@ static VALUE Unpacker_skip(VALUE self)
 
     int r = msgpack_unpacker_skip(uk, 0);
     if(r < 0) {
-        raise_unpacker_error(r);
+        raise_unpacker_error(uk, r);
     }
 
     return Qnil;
@@ -214,7 +226,7 @@ static VALUE Unpacker_skip_nil(VALUE self)
 
     int r = msgpack_unpacker_skip_nil(uk);
     if(r < 0) {
-        raise_unpacker_error(r);
+        raise_unpacker_error(uk, r);
     }
 
     if(r) {
@@ -230,7 +242,7 @@ static VALUE Unpacker_read_array_header(VALUE self)
     uint32_t size;
     int r = msgpack_unpacker_read_array_header(uk, &size);
     if(r < 0) {
-        raise_unpacker_error(r);
+        raise_unpacker_error(uk, r);
     }
 
     return ULONG2NUM(size); // long at least 32 bits
@@ -243,7 +255,7 @@ static VALUE Unpacker_read_map_header(VALUE self)
     uint32_t size;
     int r = msgpack_unpacker_read_map_header(uk, &size);
     if(r < 0) {
-        raise_unpacker_error((int)r);
+        raise_unpacker_error(uk, r);
     }
 
     return ULONG2NUM(size); // long at least 32 bits
@@ -270,7 +282,7 @@ static VALUE Unpacker_each_impl(VALUE self)
             if(r == PRIMITIVE_EOF) {
                 return Qnil;
             }
-            raise_unpacker_error(r);
+            raise_unpacker_error(uk, r);
         }
         VALUE v = msgpack_unpacker_get_last_object(uk);
 #ifdef JRUBY
@@ -369,7 +381,7 @@ static VALUE Unpacker_full_unpack(VALUE self)
 
     int r = msgpack_unpacker_read(uk, 0);
     if(r < 0) {
-        raise_unpacker_error(r);
+        raise_unpacker_error(uk, r);
     }
 
     /* raise if extra bytes follow */

spec/spec_helper.rb

@@ -22,8 +22,8 @@
   # This method was added in Ruby 3.0.0. Calling it this way asks the GC to
   # move objects around, helping to find object movement bugs.
   begin
-    GC.verify_compaction_references(double_heap: true, toward: :empty)
-  rescue NotImplementedError
+    GC.verify_compaction_references(expand_heap: true, toward: :empty)
+  rescue NotImplementedError, ArgumentError
     # Some platforms don't support compaction
   end
 end

spec/unpacker_spec.rb

@@ -901,4 +901,45 @@ def flatten(struct, results = [])
       GC.stress = stress
     end
   end
+
+  if RUBY_PLATFORM != "java"
+    it "doesn't leak when a recursive unpacker raises" do
+      hash_with_indifferent_access = Class.new(Hash)
+      msgpack = MessagePack::Factory.new
+      msgpack.register_type(
+        0x02,
+        hash_with_indifferent_access,
+        packer: ->(value, packer) do
+          packer.write(value.to_h)
+        end,
+        unpacker: ->(unpacker) { raise RuntimeError, "Ooops" },
+        recursive: true
+      )
+
+      packer = msgpack.packer
+      data = [[[[[[[hash_with_indifferent_access.new]]]]]]]
+      payload = msgpack.dump(data)
+
+      unpacker = msgpack.unpacker
+      2.times do
+        unpacker.buffer.clear
+        unpacker.feed(payload)
+        expect {
+          unpacker.full_unpack
+        }.to raise_error(RuntimeError, "Ooops")
+      end
+
+      memsize = ObjectSpace.memsize_of(unpacker)
+
+      10.times do
+        unpacker.buffer.clear
+        unpacker.feed(payload)
+        expect {
+          unpacker.full_unpack
+        }.to raise_error(RuntimeError, "Ooops")
+      end
+
+      expect(memsize).to eq ObjectSpace.memsize_of(unpacker)
+    end
+  end
 end