From 9d370e08511c7cf60413e2537adde69303767a28 Mon Sep 17 00:00:00 2001
From: "Yukihiro \"Matz\" Matsumoto" <matz@ruby.or.jp>
Date: Tue, 8 Aug 2023 12:00:49 +0900
Subject: [PATCH] mruby-socket/socket.c: ensure `struct addrinfo` to be freed

Instead of referring the last addrinfo via Addrinfo's class variable, we
use `mrb_ensure()` from `mruby-error` gem.
---
 mrbgems/mruby-socket/mrbgem.rake  |  1 +
 mrbgems/mruby-socket/src/socket.c | 60 +++++++++++++++----------------
 2 files changed, 31 insertions(+), 30 deletions(-)

diff --git a/mrbgems/mruby-socket/mrbgem.rake b/mrbgems/mruby-socket/mrbgem.rake
index f32f26fcd..3c8278c71 100644
--- a/mrbgems/mruby-socket/mrbgem.rake
+++ b/mrbgems/mruby-socket/mrbgem.rake
@@ -13,5 +13,6 @@ MRuby::Gem::Specification.new('mruby-socket') do |spec|
 
   spec.add_dependency('mruby-io', :core => 'mruby-io')
   spec.add_dependency('mruby-pack', :core => 'mruby-pack')
+  spec.add_dependency('mruby-error', :core => 'mruby-error')
   # spec.add_dependency('mruby-mtest')
 end
diff --git a/mrbgems/mruby-socket/src/socket.c b/mrbgems/mruby-socket/src/socket.c
index c3747fa6a..81ebf8adf 100644
--- a/mrbgems/mruby-socket/src/socket.c
+++ b/mrbgems/mruby-socket/src/socket.c
@@ -102,18 +102,40 @@ static int inet_pton(int af, const char *src, void *dst)
 
 #endif
 
+static mrb_value
+gen_addrinfo(mrb_state *mrb, mrb_value addrinfo)
+{
+  mrb_value ary = mrb_ary_new(mrb);
+  int arena_idx = mrb_gc_arena_save(mrb);  /* ary must be on arena! */
+  struct addrinfo *res0 = (struct addrinfo*)mrb_cptr(addrinfo);
+  mrb_value klass = mrb_obj_value(mrb_class_get_id(mrb, MRB_SYM(Addrinfo)));
+
+  for (struct addrinfo *res = res0; res != NULL; res = res->ai_next) {
+    mrb_value sa = mrb_str_new(mrb, (char*)res->ai_addr, res->ai_addrlen);
+    mrb_value ai = mrb_funcall_id(mrb, klass, MRB_SYM(new), 4, sa, mrb_fixnum_value(res->ai_family), mrb_fixnum_value(res->ai_socktype), mrb_fixnum_value(res->ai_protocol));
+    mrb_ary_push(mrb, ary, ai);
+    mrb_gc_arena_restore(mrb, arena_idx);
+  }
+
+  return ary;
+}
+
+static mrb_value
+free_addrinfo(mrb_state *mrb, mrb_value addrinfo)
+{
+  freeaddrinfo((struct addrinfo*)mrb_cptr(addrinfo));
+  return mrb_nil_value();
+}
+
 static mrb_value
 mrb_addrinfo_getaddrinfo(mrb_state *mrb, mrb_value klass)
 {
-  struct addrinfo hints = {0}, *res0, *res;
-  mrb_value ai, ary, family, lastai, nodename, protocol, sa, service, socktype;
+  struct addrinfo hints = {0}, *res0;
+  mrb_value family, nodename, protocol, service, socktype;
   mrb_int flags;
-  int arena_idx, error;
+  int error;
   const char *hostname = NULL, *servname = NULL;
 
-  ary = mrb_ary_new(mrb);
-  arena_idx = mrb_gc_arena_save(mrb);  /* ary must be on arena! */
-
   family = socktype = protocol = mrb_nil_value();
   flags = 0;
   mrb_get_args(mrb, "oo|oooi", &nodename, &service, &family, &socktype, &protocol, &flags);
@@ -155,29 +177,13 @@ mrb_addrinfo_getaddrinfo(mrb_state *mrb, mrb_value klass)
     hints.ai_protocol = (int)mrb_integer(protocol);
   }
 
-  lastai = mrb_cv_get(mrb, klass, MRB_SYM(_lastai));
-  if (mrb_cptr_p(lastai)) {
-    freeaddrinfo((struct addrinfo*)mrb_cptr(lastai));
-    mrb_cv_set(mrb, klass, MRB_SYM(_lastai), mrb_nil_value());
-  }
-
   error = getaddrinfo(hostname, servname, &hints, &res0);
   if (error) {
     mrb_raisef(mrb, E_SOCKET_ERROR, "getaddrinfo: %s", gai_strerror(error));
   }
-  mrb_cv_set(mrb, klass, MRB_SYM(_lastai), mrb_cptr_value(mrb, res0));
 
-  for (res = res0; res != NULL; res = res->ai_next) {
-    sa = mrb_str_new(mrb, (char*)res->ai_addr, res->ai_addrlen);
-    ai = mrb_funcall_id(mrb, klass, MRB_SYM(new), 4, sa, mrb_fixnum_value(res->ai_family), mrb_fixnum_value(res->ai_socktype), mrb_fixnum_value(res->ai_protocol));
-    mrb_ary_push(mrb, ary, ai);
-    mrb_gc_arena_restore(mrb, arena_idx);
-  }
-
-  freeaddrinfo(res0);
-  mrb_cv_set(mrb, klass, MRB_SYM(_lastai), mrb_nil_value());
-
-  return ary;
+  mrb_value addrinfo = mrb_cptr_value(mrb, res0);
+  return mrb_ensure(mrb, gen_addrinfo, addrinfo, free_addrinfo, addrinfo);
 }
 
 static mrb_value
@@ -867,7 +873,6 @@ mrb_mruby_socket_gem_init(mrb_state* mrb)
 #endif
 
   ai = mrb_define_class(mrb, "Addrinfo", mrb->object_class);
-  mrb_mod_cv_set(mrb, ai, MRB_SYM(_lastai), mrb_nil_value());
   mrb_define_class_method(mrb, ai, "getaddrinfo", mrb_addrinfo_getaddrinfo, MRB_ARGS_REQ(2)|MRB_ARGS_OPT(4));
   mrb_define_method(mrb, ai, "getnameinfo", mrb_addrinfo_getnameinfo, MRB_ARGS_OPT(1));
 #ifndef _WIN32
@@ -954,11 +959,6 @@ mrb_mruby_socket_gem_init(mrb_state* mrb)
 void
 mrb_mruby_socket_gem_final(mrb_state* mrb)
 {
-  mrb_value ai;
-  ai = mrb_mod_cv_get(mrb, mrb_class_get_id(mrb, MRB_SYM(Addrinfo)), MRB_SYM(_lastai));
-  if (mrb_cptr_p(ai)) {
-    freeaddrinfo((struct addrinfo*)mrb_cptr(ai));
-  }
 #ifdef _WIN32
   WSACleanup();
 #endif