Add a hook before fork() for getaddrinfo_a()

We need stop worker threads in getaddrinfo_a() before fork().
This change adds a hook before fork() that cancel all outstanding requests
and wait for all ongoing requests. Then, it waits for all worker
threads to be finished.

Fixes [Bug #17220]

5 files changed, 140 insertions, 0 deletions

diff --git a/configure.ac b/configure.ac
index 6c6d73f11f..84680cb29e 100644
--- a/configure.ac
+++ b/configure.ac
@@ -1152,6 +1152,7 @@ AC_CHECK_LIB(crypt, crypt)      # glibc (GNU/Linux, GNU/Hurd, GNU/kFreeBSD)
 AC_CHECK_LIB(dl, dlopen)	# Dynamic linking for SunOS/Solaris and SYSV
 AC_CHECK_LIB(dld, shl_load)	# Dynamic linking for HP-UX
 AC_CHECK_LIB(socket, shutdown)  # SunOS/Solaris
+AC_CHECK_LIB(anl, getaddrinfo_a)
 
 dnl Checks for header files.
 AC_HEADER_DIRENT
@@ -1943,6 +1944,7 @@ AC_CHECK_FUNCS(fsync)
 AC_CHECK_FUNCS(ftruncate)
 AC_CHECK_FUNCS(ftruncate64)		# used for Win32 platform
 AC_CHECK_FUNCS(getattrlist)
+AC_CHECK_FUNCS(getaddrinfo_a)
 AC_CHECK_FUNCS(getcwd)
 AC_CHECK_FUNCS(getgidx)
 AC_CHECK_FUNCS(getgrnam)
diff --git a/ext/socket/raddrinfo.c b/ext/socket/raddrinfo.c
index 211f05c7eb..07659a760d 100644
--- a/ext/socket/raddrinfo.c
+++ b/ext/socket/raddrinfo.c
@@ -340,6 +340,116 @@ rb_getaddrinfo(const char *node, const char *service,
 }
 
 #ifdef HAVE_GETADDRINFO_A
+struct gaicbs {
+    struct gaicbs *next;
+    struct gaicb *gaicb;
+};
+
+/* linked list to retain all outstanding and ongoing requests */
+static struct gaicbs *requests = NULL;
+
+static void
+gaicbs_add(struct gaicb *req)
+{
+    struct gaicbs *request;
+
+    if (!req) return;
+    request = (struct gaicbs *)xmalloc(sizeof(struct gaicbs));
+    request->gaicb = req;
+    request->next = requests;
+
+    requests = request;
+}
+
+static void
+gaicbs_remove(struct gaicb *req)
+{
+    struct gaicbs *request = requests;
+    struct gaicbs *prev = NULL;
+
+    if (!req) return;
+
+    while (request) {
+        if (request->gaicb == req) break;
+        prev = request;
+        request = request->next;
+    }
+
+    if (request) {
+        if (prev) {
+            prev->next = request->next;
+        } else {
+            requests = request->next;
+        }
+        xfree(request);
+    }
+}
+
+static void
+gaicbs_cancel_all(void)
+{
+    struct gaicbs *request = requests;
+    struct gaicbs *tmp, *prev = NULL;
+    int ret;
+
+    while (request) {
+        ret = gai_cancel(request->gaicb);
+        if (ret == EAI_NOTCANCELED) {
+            // continue to next request
+            prev = request;
+            request = request->next;
+        } else {
+            // remove the request from the list
+            if (prev) {
+                prev->next = request->next;
+            } else {
+                requests = request->next;
+            }
+            tmp = request;
+            request = request->next;
+            xfree(tmp);
+        }
+    }
+}
+
+static void
+gaicbs_wait_all(void)
+{
+    struct gaicbs *request = requests;
+    int size = 0;
+
+    // count gaicbs
+    while (request) {
+        size++;
+        request = request->next;
+    }
+
+    // create list to wait
+    const struct gaicb *reqs[size];
+    request = requests;
+    for (int i=0; request; i++) {
+        reqs[i] = request->gaicb;
+        request = request->next;
+    }
+
+    // wait requests
+    gai_suspend(reqs, size, NULL); // ignore result intentionally
+}
+
+/*  A mitigation for [Bug #17220].
+    It cancels all outstanding requests and waits for ongoing requests.
+    Then, it waits internal worker threads in getaddrinfo_a(3) to be finished. */
+void
+rb_getaddrinfo_a_before_exec(void)
+{
+    gaicbs_cancel_all();
+    gaicbs_wait_all();
+
+    /* wait worker threads in getaddrinfo_a(3) to be finished.
+       they will finish after 1 second sleep. */
+    sleep(1);
+}
+
 int
 rb_getaddrinfo_a(const char *node, const char *service,
                const struct addrinfo *hints,
@@ -364,11 +474,13 @@ rb_getaddrinfo_a(const char *node, const char *service,
 	reqs[0] = &req;
 	ret = getaddrinfo_a(GAI_NOWAIT, reqs, 1, NULL);
 	if (ret) return ret;
+        gaicbs_add(&req);
 
 	arg.req = &req;
 	arg.timeout = timeout;
 
 	ret = (int)(VALUE)rb_thread_call_without_gvl(nogvl_gai_suspend, &arg, RUBY_UBF_IO, 0);
+        gaicbs_remove(&req);
         if (ret && ret != EAI_ALLDONE) {
             /* EAI_ALLDONE indicates that the request already completed and gai_suspend was redundant */
             /* on Ubuntu 18.04 (or other systems), gai_suspend(3) returns EAI_SYSTEM/ENOENT on timeout */
@@ -2758,4 +2870,8 @@ rsock_init_addrinfo(void)
 
     rb_define_method(rb_cAddrinfo, "marshal_dump", addrinfo_mdump, 0);
     rb_define_method(rb_cAddrinfo, "marshal_load", addrinfo_mload, 1);
+
+#ifdef HAVE_GETADDRINFO_A
+    rb_socket_before_exec_func = rb_getaddrinfo_a_before_exec;
+#endif
 }
diff --git a/include/ruby/internal/intern/process.h b/include/ruby/internal/intern/process.h
index 2b1005a205..5cc0ca242e 100644
--- a/include/ruby/internal/intern/process.h
+++ b/include/ruby/internal/intern/process.h
@@ -28,6 +28,8 @@
 RBIMPL_SYMBOL_EXPORT_BEGIN()
 
 /* process.c */
+RUBY_EXTERN void (* rb_socket_before_exec_func)();
+
 void rb_last_status_set(int status, rb_pid_t pid);
 VALUE rb_last_status_get(void);
 int rb_proc_exec(const char*);
diff --git a/process.c b/process.c
index 612d319947..3a0616d6e3 100644
--- a/process.c
+++ b/process.c
@@ -331,6 +331,9 @@ static ID id_CLOCK_BASED_CLOCK_PROCESS_CPUTIME_ID;
 static ID id_MACH_ABSOLUTE_TIME_BASED_CLOCK_MONOTONIC;
 #endif
 static ID id_hertz;
+#ifdef HAVE_GETADDRINFO_A
+void (* rb_socket_before_exec_func)() = NULL;
+#endif
 
 /* execv and execl are async-signal-safe since SUSv4 (POSIX.1-2008, XPG7) */
 #if defined(__sun) && !defined(_XPG7) /* Solaris 10, 9, ... */
@@ -1545,6 +1548,13 @@ before_exec_async_signal_safe(void)
 static void
 before_exec_non_async_signal_safe(void)
 {
+#ifdef HAVE_GETADDRINFO_A
+    if (rb_socket_before_exec_func) {
+        /* A mitigation for [Bug #17220]. See ext/socket/raddrinfo.c */
+        rb_socket_before_exec_func();
+    }
+#endif
+
     /*
      * On Mac OS X 10.5.x (Leopard) or earlier, exec() may return ENOTSUP
      * if the process have multiple threads. Therefore we have to kill
diff --git a/test/socket/test_socket.rb b/test/socket/test_socket.rb
index f1ec927c4c..04d677c747 100644
--- a/test/socket/test_socket.rb
+++ b/test/socket/test_socket.rb
@@ -102,6 +102,16 @@ class TestSocket < Test::Unit::TestCase
     assert_nothing_raised('[ruby-core:29427]'){ TCPServer.open('localhost', 0) {} }
   end
 
+  def test_getaddrinfo_after_fork
+    skip "fork not supported" unless Process.respond_to?(:fork)
+    assert_normal_exit(<<-"end;", '[ruby-core:100329] [Bug #17220]')
+      require "socket"
+      Socket.getaddrinfo("localhost", nil)
+      pid = fork { Socket.getaddrinfo("localhost", nil) }
+      assert_equal pid, Timeout.timeout(30) { Process.wait(pid) }
+    end;
+  end
+
 
   def test_getnameinfo
     assert_raise(SocketError) { Socket.getnameinfo(["AF_UNIX", 80, "0.0.0.0"]) }