socket: fix BasicSocket#*_nonblock buffering bugs from r58400

IO#read_nonblock and IO#write_nonblock take into account
buffered data, so the Linux-only BasicSocket#read_nonblock
and BasicSocket#write_nonblock methods must, too.

This bug was only introduced in r58400
("socket: avoid fcntl for read/write_nonblock on Linux")
and does not affect any stable release.

* ext/socket/basicsocket.c (rsock_init_basicsocket):
* ext/socket/init.c (rsock_s_recvfrom_nonblock):
* ext/socket/init.c (rsock_init_socket_init):
* ext/socket/lib/socket.rb (def read_nonblock):
* ext/socket/lib/socket.rb (def write_nonblock):
* ext/socket/rubysocket.h (static inline void rsock_maybe_wait_fd):
* test/socket/test_basicsocket.rb (def test_read_write_nonblock):
  [Feature #13362]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@60496 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

5 files changed, 139 insertions, 7 deletions

diff --git a/ext/socket/basicsocket.c b/ext/socket/basicsocket.c
index 2937e31960..2641b4410b 100644
--- a/ext/socket/basicsocket.c
+++ b/ext/socket/basicsocket.c
@@ -725,6 +725,13 @@ rsock_init_basicsocket(void)
     rb_define_private_method(rb_cBasicSocket,
 			     "__recv_nonblock", bsock_recv_nonblock, 4);
 
+#if MSG_DONTWAIT_RELIABLE
+    rb_define_private_method(rb_cBasicSocket,
+			     "__read_nonblock", rsock_read_nonblock, 3);
+    rb_define_private_method(rb_cBasicSocket,
+			     "__write_nonblock", rsock_write_nonblock, 2);
+#endif
+
     /* in ancdata.c */
     rb_define_private_method(rb_cBasicSocket, "__sendmsg",
 			     rsock_bsock_sendmsg, 4);
diff --git a/ext/socket/init.c b/ext/socket/init.c
index 1ecd4fe352..189977dcba 100644
--- a/ext/socket/init.c
+++ b/ext/socket/init.c
@@ -280,6 +280,108 @@ rsock_s_recvfrom_nonblock(VALUE sock, VALUE len, VALUE flg, VALUE str,
     return rb_assoc_new(str, addr);
 }
 
+#if MSG_DONTWAIT_RELIABLE
+static VALUE sym_wait_writable;
+
+/* copied from io.c :< */
+static long
+read_buffered_data(char *ptr, long len, rb_io_t *fptr)
+{
+    int n = fptr->rbuf.len;
+
+    if (n <= 0) return 0;
+    if (n > len) n = (int)len;
+    MEMMOVE(ptr, fptr->rbuf.ptr+fptr->rbuf.off, char, n);
+    fptr->rbuf.off += n;
+    fptr->rbuf.len -= n;
+    return n;
+}
+
+/* :nodoc: */
+VALUE
+rsock_read_nonblock(VALUE sock, VALUE length, VALUE buf, VALUE ex)
+{
+    rb_io_t *fptr;
+    long n;
+    long len = NUM2LONG(length);
+    VALUE str = rsock_strbuf(buf, len);
+    char *ptr;
+
+    OBJ_TAINT(str);
+    GetOpenFile(sock, fptr);
+
+    if (len == 0) {
+	return str;
+    }
+
+    ptr = RSTRING_PTR(str);
+    n = read_buffered_data(ptr, len, fptr);
+    if (n <= 0) {
+	n = (long)recv(fptr->fd, ptr, len, MSG_DONTWAIT);
+	if (n < 0) {
+	    int e = errno;
+	    if ((e == EWOULDBLOCK || e == EAGAIN)) {
+		if (ex == Qfalse) return sym_wait_readable;
+		rb_readwrite_syserr_fail(RB_IO_WAIT_READABLE,
+					 e, "read would block");
+	    }
+	    rb_syserr_fail_path(e, fptr->pathv);
+	}
+    }
+    if (len != n) {
+	rb_str_modify(str);
+	rb_str_set_len(str, n);
+	if (str != buf) {
+	    rb_str_resize(str, n);
+	}
+    }
+    if (n == 0) {
+	if (ex == Qfalse) return Qnil;
+	rb_eof_error();
+    }
+
+    return str;
+}
+
+/* :nodoc: */
+VALUE
+rsock_write_nonblock(VALUE sock, VALUE str, VALUE ex)
+{
+    rb_io_t *fptr;
+    long n;
+
+    if (!RB_TYPE_P(str, T_STRING))
+	str = rb_obj_as_string(str);
+
+    sock = rb_io_get_write_io(sock);
+    GetOpenFile(sock, fptr);
+    rb_io_check_writable(fptr);
+
+    /*
+     * As with IO#write_nonblock, we may block if somebody is relying on
+     * buffered I/O; but nobody actually hits this because pipes and sockets
+     * are not userspace-buffered in Ruby by default.
+     */
+    if (fptr->wbuf.len > 0) {
+	rb_io_flush(sock);
+    }
+
+    n = (long)send(fptr->fd, RSTRING_PTR(str), RSTRING_LEN(str), MSG_DONTWAIT);
+    if (n < 0) {
+	int e = errno;
+
+	if (e == EWOULDBLOCK || e == EAGAIN) {
+	    if (ex == Qfalse) return sym_wait_writable;
+	    rb_readwrite_syserr_fail(RB_IO_WAIT_WRITABLE, e,
+				    "write would block");
+	}
+	rb_syserr_fail_path(e, fptr->pathv);
+    }
+
+    return LONG2FIX(n);
+}
+#endif /* MSG_DONTWAIT_RELIABLE */
+
 /* returns true if SOCK_CLOEXEC is supported */
 int rsock_detect_cloexec(int fd)
 {
@@ -680,4 +782,8 @@ rsock_init_socket_init(void)
 
 #undef rb_intern
     sym_wait_readable = ID2SYM(rb_intern("wait_readable"));
+
+#if MSG_DONTWAIT_RELIABLE
+    sym_wait_writable = ID2SYM(rb_intern("wait_writable"));
+#endif
 }
diff --git a/ext/socket/lib/socket.rb b/ext/socket/lib/socket.rb
index 53d961e680..18f79b3f08 100644
--- a/ext/socket/lib/socket.rb
+++ b/ext/socket/lib/socket.rb
@@ -449,16 +449,11 @@ class BasicSocket < IO
   # Do other platforms support MSG_DONTWAIT reliably?
   if RUBY_PLATFORM =~ /linux/ && Socket.const_defined?(:MSG_DONTWAIT)
     def read_nonblock(len, str = nil, exception: true) # :nodoc:
-      case rv = __recv_nonblock(len, 0, str, exception)
-      when '' # recv_nonblock returns empty string on EOF
-        exception ? raise(EOFError, 'end of file reached') : nil
-      else
-        rv
-      end
+      __read_nonblock(len, str, exception)
     end
 
     def write_nonblock(buf, exception: true) # :nodoc:
-      __sendmsg_nonblock(buf, 0, nil, nil, exception)
+      __write_nonblock(buf, exception)
     end
   end
 end
diff --git a/ext/socket/rubysocket.h b/ext/socket/rubysocket.h
index 352da8c56e..922df9b87b 100644
--- a/ext/socket/rubysocket.h
+++ b/ext/socket/rubysocket.h
@@ -430,6 +430,9 @@ static inline void rsock_maybe_wait_fd(int fd) { }
 #  define MSG_DONTWAIT_RELIABLE 0
 #endif
 
+VALUE rsock_read_nonblock(VALUE sock, VALUE length, VALUE buf, VALUE ex);
+VALUE rsock_write_nonblock(VALUE sock, VALUE buf, VALUE ex);
+
 #if !defined HAVE_INET_NTOP && ! defined _WIN32
 const char *inet_ntop(int, const void *, char *, size_t);
 #elif defined __MINGW32__
diff --git a/test/socket/test_basicsocket.rb b/test/socket/test_basicsocket.rb
index 0671e97c75..d388b4f0dd 100644
--- a/test/socket/test_basicsocket.rb
+++ b/test/socket/test_basicsocket.rb
@@ -205,4 +205,25 @@ class TestSocket_BasicSocket < Test::Unit::TestCase
       assert_not_predicate(ssock, :nonblock?) unless set_nb
     end
   end
+
+  def test_read_nonblock_mix_buffered
+    socks do |sserv, ssock, csock|
+      ssock.write("hello\nworld\n")
+      assert_equal "hello\n", csock.gets
+      IO.select([csock], nil, nil, 10) or
+        flunk 'socket did not become readable'
+      assert_equal "world\n", csock.read_nonblock(8)
+    end
+  end
+
+  def test_write_nonblock_buffered
+    socks do |sserv, ssock, csock|
+      ssock.sync = false
+      ssock.write("h")
+      assert_equal :wait_readable, csock.read_nonblock(1, exception: false)
+      assert_equal 4, ssock.write_nonblock("ello")
+      ssock.close
+      assert_equal "hello", csock.read(5)
+    end
+  end
 end if defined?(BasicSocket)