Ruby: make sure map parsing frames are GC-rooted.

3 files changed, 58 insertions, 8 deletions

diff --git a/ruby/ext/google/protobuf_c/encode_decode.c b/ruby/ext/google/protobuf_c/encode_decode.c
index dd21a046..08c72bcc 100644
--- a/ruby/ext/google/protobuf_c/encode_decode.c
+++ b/ruby/ext/google/protobuf_c/encode_decode.c
@@ -255,10 +255,54 @@ typedef struct {
 // value into the map.
 typedef struct {
   VALUE map;
+  const map_handlerdata_t* handlerdata;
   char key_storage[NATIVE_SLOT_MAX_SIZE];
   char value_storage[NATIVE_SLOT_MAX_SIZE];
 } map_parse_frame_t;
 
+static void MapParseFrame_mark(void* _self) {
+  map_parse_frame_t* frame = _self;
+
+  // This shouldn't strictly be necessary since this should be rooted by the
+  // message itself, but it can't hurt.
+  rb_gc_mark(frame->map);
+
+  native_slot_mark(frame->handlerdata->key_field_type, &frame->key_storage);
+  native_slot_mark(frame->handlerdata->value_field_type, &frame->value_storage);
+}
+
+void MapParseFrame_free(void* self) {
+  xfree(self);
+}
+
+rb_data_type_t MapParseFrame_type = {
+  "MapParseFrame",
+  { MapParseFrame_mark, MapParseFrame_free, NULL },
+};
+
+// Array of Ruby objects wrapping map_parse_frame_t.
+// We don't allow multiple concurrent decodes, so we assume that this global
+// variable is specific to the "current" decode.
+VALUE map_parse_frames;
+
+static map_parse_frame_t* map_push_frame(VALUE map,
+                                         const map_handlerdata_t* handlerdata) {
+  map_parse_frame_t* frame = ALLOC(map_parse_frame_t);
+  frame->handlerdata = handlerdata;
+  frame->map = map;
+  native_slot_init(handlerdata->key_field_type, &frame->key_storage);
+  native_slot_init(handlerdata->value_field_type, &frame->value_storage);
+
+  rb_ary_push(map_parse_frames,
+              TypedData_Wrap_Struct(rb_cObject, &MapParseFrame_type, frame));
+
+  return frame;
+}
+
+static void map_pop_frame() {
+  rb_ary_pop(map_parse_frames);
+}
+
 // Handler to begin a map entry: allocates a temporary frame. This is the
 // 'startsubmsg' handler on the msgdef that contains the map field.
 static void *startmapentry_handler(void *closure, const void *hd) {
@@ -266,13 +310,7 @@ static void *startmapentry_handler(void *closure, const void *hd) {
   const map_handlerdata_t* mapdata = hd;
   VALUE map_rb = DEREF(msg, mapdata->ofs, VALUE);
 
-  map_parse_frame_t* frame = ALLOC(map_parse_frame_t);
-  frame->map = map_rb;
-
-  native_slot_init(mapdata->key_field_type, &frame->key_storage);
-  native_slot_init(mapdata->value_field_type, &frame->value_storage);
-
-  return frame;
+  return map_push_frame(map_rb, mapdata);
 }
 
 // Handler to end a map entry: inserts the value defined during the message into
@@ -298,7 +336,7 @@ static bool endmap_handler(void *closure, const void *hd, upb_status* s) {
       &frame->value_storage);
 
   Map_index_set(frame->map, key, value);
-  xfree(frame);
+  map_pop_frame();
 
   return true;
 }
@@ -737,6 +775,10 @@ VALUE Message_decode(VALUE klass, VALUE data) {
   msg_rb = rb_class_new_instance(0, NULL, msgklass);
   TypedData_Get_Struct(msg_rb, MessageHeader, &Message_type, msg);
 
+  // We generally expect this to be clear already, but clear it in case parsing
+  // previously got interrupted somehow.
+  rb_ary_clear(map_parse_frames);
+
   {
     const upb_pbdecodermethod* method = msgdef_decodermethod(desc);
     const upb_handlers* h = upb_pbdecodermethod_desthandlers(method);
@@ -781,6 +823,10 @@ VALUE Message_decode_json(VALUE klass, VALUE data) {
   msg_rb = rb_class_new_instance(0, NULL, msgklass);
   TypedData_Get_Struct(msg_rb, MessageHeader, &Message_type, msg);
 
+  // We generally expect this to be clear already, but clear it in case parsing
+  // previously got interrupted somehow.
+  rb_ary_clear(map_parse_frames);
+
   {
     const upb_json_parsermethod* method = msgdef_jsonparsermethod(desc);
     stackenv se;
diff --git a/ruby/ext/google/protobuf_c/protobuf.c b/ruby/ext/google/protobuf_c/protobuf.c
index 7cde4aec..98963667 100644
--- a/ruby/ext/google/protobuf_c/protobuf.c
+++ b/ruby/ext/google/protobuf_c/protobuf.c
@@ -112,4 +112,6 @@ void Init_protobuf_c() {
 
   upb_def_to_ruby_obj_map = rb_hash_new();
   rb_gc_register_address(&upb_def_to_ruby_obj_map);
+  map_parse_frames = rb_ary_new();
+  rb_gc_register_address(&map_parse_frames);
 }
diff --git a/ruby/ext/google/protobuf_c/protobuf.h b/ruby/ext/google/protobuf_c/protobuf.h
index 08b8d516..d5ced567 100644
--- a/ruby/ext/google/protobuf_c/protobuf.h
+++ b/ruby/ext/google/protobuf_c/protobuf.h
@@ -166,6 +166,8 @@ extern VALUE cBuilder;
 extern VALUE cError;
 extern VALUE cParseError;
 
+extern VALUE map_parse_frames;
+
 // We forward-declare all of the Ruby method implementations here because we
 // sometimes call the methods directly across .c files, rather than going
 // through Ruby's method dispatching (e.g. during message parse). It's cleaner