Add resolve_timeout to TCPSocket [Feature #17134]

7 files changed, 66 insertions, 8 deletions

diff --git a/ext/socket/ipsocket.c b/ext/socket/ipsocket.c
index a2cb6e0e12..e2b7308b83 100644
--- a/ext/socket/ipsocket.c
+++ b/ext/socket/ipsocket.c
@@ -19,6 +19,7 @@ struct inetsock_arg
     } remote, local;
     int type;
     int fd;
+    VALUE resolv_timeout;
 };
 
 static VALUE
@@ -49,10 +50,20 @@ init_inetsock_internal(VALUE v)
     int fd, status = 0, local = 0;
     int family = AF_UNSPEC;
     const char *syscall = 0;
-
+    VALUE resolv_timeout = arg->resolv_timeout;
+
+#ifdef HAVE_GETADDRINFO_A
+    arg->remote.res = rsock_addrinfo_a(arg->remote.host, arg->remote.serv,
+				       family, SOCK_STREAM,
+				       (type == INET_SERVER) ? AI_PASSIVE : 0,
+				       resolv_timeout);
+#else
     arg->remote.res = rsock_addrinfo(arg->remote.host, arg->remote.serv,
 				     family, SOCK_STREAM,
 				     (type == INET_SERVER) ? AI_PASSIVE : 0);
+#endif
+
+
     /*
      * Maybe also accept a local address
      */
@@ -157,7 +168,8 @@ init_inetsock_internal(VALUE v)
 
 VALUE
 rsock_init_inetsock(VALUE sock, VALUE remote_host, VALUE remote_serv,
-	            VALUE local_host, VALUE local_serv, int type)
+	            VALUE local_host, VALUE local_serv, int type,
+		    VALUE resolv_timeout)
 {
     struct inetsock_arg arg;
     arg.sock = sock;
@@ -169,6 +181,7 @@ rsock_init_inetsock(VALUE sock, VALUE remote_host, VALUE remote_serv,
     arg.local.res = 0;
     arg.type = type;
     arg.fd = -1;
+    arg.resolv_timeout = resolv_timeout;
     return rb_ensure(init_inetsock_internal, (VALUE)&arg,
 		     inetsock_cleanup, (VALUE)&arg);
 }
diff --git a/ext/socket/raddrinfo.c b/ext/socket/raddrinfo.c
index 4dd2867781..211f05c7eb 100644
--- a/ext/socket/raddrinfo.c
+++ b/ext/socket/raddrinfo.c
@@ -662,6 +662,20 @@ rsock_addrinfo(VALUE host, VALUE port, int family, int socktype, int flags)
     return rsock_getaddrinfo(host, port, &hints, 1);
 }
 
+#ifdef HAVE_GETADDRINFO_A
+struct rb_addrinfo*
+rsock_addrinfo_a(VALUE host, VALUE port, int family, int socktype, int flags, VALUE timeout)
+{
+    struct addrinfo hints;
+
+    MEMZERO(&hints, struct addrinfo, 1);
+    hints.ai_family = family;
+    hints.ai_socktype = socktype;
+    hints.ai_flags = flags;
+    return rsock_getaddrinfo_a(host, port, &hints, 1, timeout);
+}
+#endif
+
 VALUE
 rsock_ipaddr(struct sockaddr *sockaddr, socklen_t sockaddrlen, int norevlookup)
 {
diff --git a/ext/socket/rubysocket.h b/ext/socket/rubysocket.h
index 30b7a7777a..91b446d3a1 100644
--- a/ext/socket/rubysocket.h
+++ b/ext/socket/rubysocket.h
@@ -321,6 +321,7 @@ int rsock_fd_family(int fd);
 struct rb_addrinfo *rsock_addrinfo(VALUE host, VALUE port, int family, int socktype, int flags);
 struct rb_addrinfo *rsock_getaddrinfo(VALUE host, VALUE port, struct addrinfo *hints, int socktype_hack);
 #ifdef HAVE_GETADDRINFO_A
+struct rb_addrinfo *rsock_addrinfo_a(VALUE host, VALUE port, int family, int socktype, int flags, VALUE timeout);
 struct rb_addrinfo *rsock_getaddrinfo_a(VALUE host, VALUE port, struct addrinfo *hints, int socktype_hack, VALUE timeout);
 #endif
 
@@ -349,7 +350,7 @@ int rsock_socket(int domain, int type, int proto);
 int rsock_detect_cloexec(int fd);
 VALUE rsock_init_sock(VALUE sock, int fd);
 VALUE rsock_sock_s_socketpair(int argc, VALUE *argv, VALUE klass);
-VALUE rsock_init_inetsock(VALUE sock, VALUE remote_host, VALUE remote_serv, VALUE local_host, VALUE local_serv, int type);
+VALUE rsock_init_inetsock(VALUE sock, VALUE remote_host, VALUE remote_serv, VALUE local_host, VALUE local_serv, int type, VALUE resolv_timeout);
 VALUE rsock_init_unixsock(VALUE sock, VALUE path, int server);
 
 struct rsock_send_arg {
diff --git a/ext/socket/sockssocket.c b/ext/socket/sockssocket.c
index 82789eeaab..78b0055ccc 100644
--- a/ext/socket/sockssocket.c
+++ b/ext/socket/sockssocket.c
@@ -34,7 +34,7 @@ socks_init(VALUE sock, VALUE host, VALUE port)
 	init = 1;
     }
 
-    return rsock_init_inetsock(sock, host, port, Qnil, Qnil, INET_SOCKS);
+    return rsock_init_inetsock(sock, host, port, Qnil, Qnil, INET_SOCKS, Qnil);
 }
 
 #ifdef SOCKS5
diff --git a/ext/socket/tcpserver.c b/ext/socket/tcpserver.c
index 1bbb31adcf..ad31e16306 100644
--- a/ext/socket/tcpserver.c
+++ b/ext/socket/tcpserver.c
@@ -36,7 +36,7 @@ tcp_svr_init(int argc, VALUE *argv, VALUE sock)
     VALUE hostname, port;
 
     rb_scan_args(argc, argv, "011", &hostname, &port);
-    return rsock_init_inetsock(sock, hostname, port, Qnil, Qnil, INET_SERVER);
+    return rsock_init_inetsock(sock, hostname, port, Qnil, Qnil, INET_SERVER, Qnil);
 }
 
 /*
diff --git a/ext/socket/tcpsocket.c b/ext/socket/tcpsocket.c
index f3fcee781d..6baf367709 100644
--- a/ext/socket/tcpsocket.c
+++ b/ext/socket/tcpsocket.c
@@ -23,12 +23,28 @@ tcp_init(int argc, VALUE *argv, VALUE sock)
 {
     VALUE remote_host, remote_serv;
     VALUE local_host, local_serv;
+    VALUE opt;
+    static ID keyword_ids[1];
+    VALUE kwargs[1];
+    VALUE resolv_timeout = Qnil;
 
-    rb_scan_args(argc, argv, "22", &remote_host, &remote_serv,
-			&local_host, &local_serv);
+    if (!keyword_ids[0]) {
+	CONST_ID(keyword_ids[0], "resolv_timeout");
+    }
+
+    rb_scan_args(argc, argv, "22:", &remote_host, &remote_serv,
+			&local_host, &local_serv, &opt);
+
+    if (!NIL_P(opt)) {
+	rb_get_kwargs(opt, keyword_ids, 0, 1, kwargs);
+	if (kwargs[0] != Qundef) {
+	    resolv_timeout = kwargs[0];
+	}
+    }
 
     return rsock_init_inetsock(sock, remote_host, remote_serv,
-			       local_host, local_serv, INET_CLIENT);
+			       local_host, local_serv, INET_CLIENT,
+			       resolv_timeout);
 }
 
 static VALUE
diff --git a/test/socket/test_tcp.rb b/test/socket/test_tcp.rb
index 11325fdedb..15c79b7519 100644
--- a/test/socket/test_tcp.rb
+++ b/test/socket/test_tcp.rb
@@ -55,6 +55,20 @@ class TestSocket_TCPSocket < Test::Unit::TestCase
     t.close if t && !t.closed?
   end
 
+  def test_initialize_resolv_timeout
+    TCPServer.open("localhost", 0) do |svr|
+      th = Thread.new {
+        c = svr.accept
+        c.close
+      }
+      addr = svr.addr
+      s = TCPSocket.new(addr[3], addr[1], resolv_timeout: 10)
+      th.join
+    ensure
+      s.close()
+    end
+  end
+
   def test_recvfrom
     TCPServer.open("localhost", 0) {|svr|
       th = Thread.new {