From 44df509dc2dee1acdc8f62af3430612c9a56107d Mon Sep 17 00:00:00 2001
From: Peter Zhu <peter@peterzhu.ca>
Date: Thu, 21 Mar 2024 12:17:15 -0400
Subject: [PATCH] Fix memory leak when exception is raised during JSON
 generation

If an exception is raised the FBuffer is leaked.

For example, the following script leaks memory:

    o = Object.new
    def o.to_json(a) = raise

    10.times do
      100_000.times do
        begin
          JSON(o)
        rescue
        end
      end

      puts `ps -o rss= -p #{$$}`
    end

Before:

    31824
    35696
    40240
    44304
    47424
    50944
    54000
    58384
    62416
    65296

After:

    24416
    24640
    24640
    24736
    24736
    24736
    24736
    24736
    24736
    24736
---
 ext/json/ext/generator/generator.c | 40 ++++++++++++++++++++++++++----
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/ext/json/ext/generator/generator.c b/ext/json/ext/generator/generator.c
index a33df848..6d78284b 100644
--- a/ext/json/ext/generator/generator.c
+++ b/ext/json/ext/generator/generator.c
@@ -892,7 +892,6 @@ static void generate_json_object(FBuffer *buffer, VALUE Vstate, JSON_Generator_S
     struct hash_foreach_arg arg;
 
     if (max_nesting != 0 && depth > max_nesting) {
-        fbuffer_free(buffer);
         rb_raise(eNestingError, "nesting of %ld is too deep", --state->depth);
     }
     fbuffer_append_char(buffer, '{');
@@ -927,7 +926,6 @@ static void generate_json_array(FBuffer *buffer, VALUE Vstate, JSON_Generator_St
     long depth = ++state->depth;
     int i, j;
     if (max_nesting != 0 && depth > max_nesting) {
-        fbuffer_free(buffer);
         rb_raise(eNestingError, "nesting of %ld is too deep", --state->depth);
     }
     fbuffer_append_char(buffer, '[');
@@ -1020,10 +1018,8 @@ static void generate_json_float(FBuffer *buffer, VALUE Vstate, JSON_Generator_St
     VALUE tmp = rb_funcall(obj, i_to_s, 0);
     if (!allow_nan) {
         if (isinf(value)) {
-            fbuffer_free(buffer);
             rb_raise(eGeneratorError, "%"PRIsVALUE" not allowed in JSON", RB_OBJ_STRING(tmp));
         } else if (isnan(value)) {
-            fbuffer_free(buffer);
             rb_raise(eGeneratorError, "%"PRIsVALUE" not allowed in JSON", RB_OBJ_STRING(tmp));
         }
     }
@@ -1096,11 +1092,45 @@ static FBuffer *cState_prepare_buffer(VALUE self)
     return buffer;
 }
 
+struct generate_json_data {
+    FBuffer *buffer;
+    VALUE vstate;
+    JSON_Generator_State *state;
+    VALUE obj;
+};
+
+static VALUE generate_json_try(VALUE d)
+{
+    struct generate_json_data *data = (struct generate_json_data *)d;
+
+    generate_json(data->buffer, data->vstate, data->state, data->obj);
+
+    return Qnil;
+}
+
+static VALUE generate_json_rescue(VALUE d, VALUE exc)
+{
+    struct generate_json_data *data = (struct generate_json_data *)d;
+    fbuffer_free(data->buffer);
+
+    rb_exc_raise(exc);
+
+    return Qundef;
+}
+
 static VALUE cState_partial_generate(VALUE self, VALUE obj)
 {
     FBuffer *buffer = cState_prepare_buffer(self);
     GET_STATE(self);
-    generate_json(buffer, self, state, obj);
+
+    struct generate_json_data data = {
+        .buffer = buffer,
+        .vstate = self,
+        .state = state,
+        .obj = obj
+    };
+    rb_rescue(generate_json_try, (VALUE)&data, generate_json_rescue, (VALUE)&data);
+
     return fbuffer_to_s(buffer);
 }