1 files changed, 57 insertions, 31 deletions

diff --git a/ext/io/nonblock/nonblock.c b/ext/io/nonblock/nonblock.c
index e5915fa128..d90538f735 100644
--- a/ext/io/nonblock/nonblock.c
+++ b/ext/io/nonblock/nonblock.c
@@ -17,16 +17,27 @@
 #endif
 #include <fcntl.h>
 
+#ifndef HAVE_RB_IO_DESCRIPTOR
+static int
+io_descriptor_fallback(VALUE io)
+{
+    rb_io_t *fptr;
+    GetOpenFile(io, fptr);
+    return fptr->fd;
+}
+#define rb_io_descriptor io_descriptor_fallback
+#endif
+
 #ifdef F_GETFL
 static int
-io_nonblock_mode(int fd)
+get_fcntl_flags(int fd)
 {
     int f = fcntl(fd, F_GETFL);
     if (f == -1) rb_sys_fail(0);
     return f;
 }
 #else
-#define io_nonblock_mode(fd) ((void)(fd), 0)
+#define get_fcntl_flags(fd) ((void)(fd), 0)
 #endif
 
 #ifdef F_GETFL
@@ -39,10 +50,8 @@ io_nonblock_mode(int fd)
 static VALUE
 rb_io_nonblock_p(VALUE io)
 {
-    rb_io_t *fptr;
-    GetOpenFile(io, fptr);
-    if (io_nonblock_mode(fptr->fd) & O_NONBLOCK)
-	return Qtrue;
+    if (get_fcntl_flags(rb_io_descriptor(io)) & O_NONBLOCK)
+        return Qtrue;
     return Qfalse;
 }
 #else
@@ -50,6 +59,15 @@ rb_io_nonblock_p(VALUE io)
 #endif
 
 #ifdef F_SETFL
+static void
+set_fcntl_flags(int fd, int f)
+{
+    if (fcntl(fd, F_SETFL, f) == -1)
+	rb_sys_fail(0);
+}
+
+#ifndef RUBY_IO_NONBLOCK_METHODS
+
 static int
 io_nonblock_set(int fd, int f, int nb)
 {
@@ -63,8 +81,7 @@ io_nonblock_set(int fd, int f, int nb)
 	    return 0;
 	f &= ~O_NONBLOCK;
     }
-    if (fcntl(fd, F_SETFL, f) == -1)
-	rb_sys_fail(0);
+    set_fcntl_flags(fd, f);
     return 1;
 }
 
@@ -116,23 +133,28 @@ io_nonblock_set(int fd, int f, int nb)
  *
  */
 static VALUE
-rb_io_nonblock_set(VALUE io, VALUE nb)
+rb_io_nonblock_set(VALUE self, VALUE value)
 {
-    rb_io_t *fptr;
-    GetOpenFile(io, fptr);
-    if (RTEST(nb))
-	rb_io_set_nonblock(fptr);
-    else
-	io_nonblock_set(fptr->fd, io_nonblock_mode(fptr->fd), RTEST(nb));
-    return io;
+    if (RTEST(value)) {
+        rb_io_t *fptr;
+        GetOpenFile(self, fptr);
+        rb_io_set_nonblock(fptr);
+    }
+    else {
+        int descriptor = rb_io_descriptor(self);
+        io_nonblock_set(descriptor, get_fcntl_flags(descriptor), RTEST(value));
+    }
+
+    return self;
 }
 
+#endif /* RUBY_IO_NONBLOCK_METHODS */
+
 static VALUE
 io_nonblock_restore(VALUE arg)
 {
     int *restore = (int *)arg;
-    if (fcntl(restore[0], F_SETFL, restore[1]) == -1)
-	rb_sys_fail(0);
+    set_fcntl_flags(restore[0], restore[1]);
     return Qnil;
 }
 
@@ -147,24 +169,25 @@ io_nonblock_restore(VALUE arg)
  * The original mode is restored after the block is executed.
  */
 static VALUE
-rb_io_nonblock_block(int argc, VALUE *argv, VALUE io)
+rb_io_nonblock_block(int argc, VALUE *argv, VALUE self)
 {
     int nb = 1;
-    rb_io_t *fptr;
-    int f, restore[2];
 
-    GetOpenFile(io, fptr);
+    int descriptor = rb_io_descriptor(self);
+
     if (argc > 0) {
-	VALUE v;
-	rb_scan_args(argc, argv, "01", &v);
-	nb = RTEST(v);
+        VALUE v;
+        rb_scan_args(argc, argv, "01", &v);
+        nb = RTEST(v);
     }
-    f = io_nonblock_mode(fptr->fd);
-    restore[0] = fptr->fd;
-    restore[1] = f;
-    if (!io_nonblock_set(fptr->fd, f, nb))
-	return rb_yield(io);
-    return rb_ensure(rb_yield, io, io_nonblock_restore, (VALUE)restore);
+
+    int current_flags = get_fcntl_flags(descriptor);
+    int restore[2] = {descriptor, current_flags};
+
+    if (!io_nonblock_set(descriptor, current_flags, nb))
+        return rb_yield(self);
+
+    return rb_ensure(rb_yield, self, io_nonblock_restore, (VALUE)restore);
 }
 #else
 #define rb_io_nonblock_set rb_f_notimplement
@@ -174,7 +197,10 @@ rb_io_nonblock_block(int argc, VALUE *argv, VALUE io)
 void
 Init_nonblock(void)
 {
+#ifndef RUBY_IO_NONBLOCK_METHODS
     rb_define_method(rb_cIO, "nonblock?", rb_io_nonblock_p, 0);
     rb_define_method(rb_cIO, "nonblock=", rb_io_nonblock_set, 1);
+#endif
+
     rb_define_method(rb_cIO, "nonblock", rb_io_nonblock_block, -1);
 }