1 files changed, 56 insertions, 39 deletions

diff --git a/ext/socket/init.c b/ext/socket/init.c
index af46b8edaa..8eb8c8e901 100644
--- a/ext/socket/init.c
+++ b/ext/socket/init.c
@@ -166,7 +166,7 @@ recvfrom_locktmp(VALUE v)
 }
 
 VALUE
-rsock_s_recvfrom(VALUE sock, int argc, VALUE *argv, enum sock_recv_type from)
+rsock_s_recvfrom(VALUE socket, int argc, VALUE *argv, enum sock_recv_type from)
 {
     rb_io_t *fptr;
     VALUE str;
@@ -177,27 +177,35 @@ rsock_s_recvfrom(VALUE sock, int argc, VALUE *argv, enum sock_recv_type from)
 
     rb_scan_args(argc, argv, "12", &len, &flg, &str);
 
-    if (flg == Qnil) arg.flags = 0;
-    else             arg.flags = NUM2INT(flg);
+    if (flg == Qnil)
+        arg.flags = 0;
+    else
+        arg.flags = NUM2INT(flg);
+
     buflen = NUM2INT(len);
     str = rsock_strbuf(str, buflen);
 
-    GetOpenFile(sock, fptr);
+    RB_IO_POINTER(socket, fptr);
+
     if (rb_io_read_pending(fptr)) {
-	rb_raise(rb_eIOError, "recv for buffered IO");
+        rb_raise(rb_eIOError, "recv for buffered IO");
     }
+
     arg.fd = fptr->fd;
     arg.alen = (socklen_t)sizeof(arg.buf);
     arg.str = str;
     arg.length = buflen;
 
-    while (rb_io_check_closed(fptr),
-	   rsock_maybe_wait_fd(arg.fd),
-	   (slen = (long)rb_str_locktmp_ensure(str, recvfrom_locktmp,
-	                                       (VALUE)&arg)) < 0) {
-        if (!rb_io_wait_readable(fptr->fd)) {
+    while (true) {
+        rb_io_check_closed(fptr);
+        rsock_maybe_wait_fd(arg.fd);
+
+        slen = (long)rb_str_locktmp_ensure(str, recvfrom_locktmp, (VALUE)&arg);
+
+        if (slen >= 0) break;
+
+        if (!rb_io_maybe_wait_readable(errno, socket, Qnil))
             rb_sys_fail("recvfrom(2)");
-        }
     }
 
     /* Resize the string to the amount of data received */
@@ -221,7 +229,7 @@ rsock_s_recvfrom(VALUE sock, int argc, VALUE *argv, enum sock_recv_type from)
         return rb_assoc_new(str, rsock_unixaddr(&arg.buf.un, arg.alen));
 #endif
       case RECV_SOCKET:
-	return rb_assoc_new(str, rsock_io_socket_addrinfo(sock, &arg.buf.addr, arg.alen));
+	return rb_assoc_new(str, rsock_io_socket_addrinfo(socket, &arg.buf.addr, arg.alen));
       default:
 	rb_bug("rsock_s_recvfrom called with bad value");
     }
@@ -682,38 +690,47 @@ accept_blocking(void *data)
 }
 
 VALUE
-rsock_s_accept(VALUE klass, int fd, struct sockaddr *sockaddr, socklen_t *len)
+rsock_s_accept(VALUE klass, VALUE io, struct sockaddr *sockaddr, socklen_t *len)
 {
-    int fd2;
+    rb_io_t *fptr = NULL;
+    RB_IO_POINTER(io, fptr);
+
+    struct accept_arg accept_arg = {
+      .fd = fptr->fd,
+      .sockaddr = sockaddr,
+      .len = len
+    };
+
     int retry = 0;
-    struct accept_arg arg;
 
-    arg.fd = fd;
-    arg.sockaddr = sockaddr;
-    arg.len = len;
   retry:
-    rsock_maybe_wait_fd(fd);
-    fd2 = (int)BLOCKING_REGION_FD(accept_blocking, &arg);
-    if (fd2 < 0) {
-	int e = errno;
-	switch (e) {
-	  case EMFILE:
-	  case ENFILE:
-	  case ENOMEM:
-	    if (retry) break;
-	    rb_gc();
-	    retry = 1;
-	    goto retry;
-	  default:
-	    if (!rb_io_wait_readable(fd)) break;
-	    retry = 0;
-	    goto retry;
-	}
-	rb_syserr_fail(e, "accept(2)");
+    rsock_maybe_wait_fd(accept_arg.fd);
+    int peer = (int)BLOCKING_REGION_FD(accept_blocking, &accept_arg);
+    if (peer < 0) {
+        int error = errno;
+
+        switch (error) {
+          case EMFILE:
+          case ENFILE:
+          case ENOMEM:
+            if (retry) break;
+            rb_gc();
+            retry = 1;
+            goto retry;
+          default:
+            if (!rb_io_maybe_wait_readable(error, io, Qnil)) break;
+            retry = 0;
+            goto retry;
+        }
+
+        rb_syserr_fail(error, "accept(2)");
     }
-    rb_update_max_fd(fd2);
-    if (!klass) return INT2NUM(fd2);
-    return rsock_init_sock(rb_obj_alloc(klass), fd2);
+
+    rb_update_max_fd(peer);
+
+    if (!klass) return INT2NUM(peer);
+
+    return rsock_init_sock(rb_obj_alloc(klass), peer);
 }
 
 int