fix(receive): a receive 0 should immediately return, not block

The fix is slightly more complex because reading from a closed
socket should still return a closed-error.

fixes: #427
fixes: https://github.com/ledgetech/lua-resty-http/pull/313

5 files changed, 53 insertions, 10 deletions

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 229fcbf..d516568 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -49,4 +49,5 @@ jobs:
           lua mimetest.lua
           lua urltest.lua
           lua test_socket_error.lua
+          lua test_receive_zero.lua
           kill %1
diff --git a/src/buffer.c b/src/buffer.c
index 3d48a09..b7e97f7 100644
--- a/src/buffer.c
+++ b/src/buffer.c
@@ -11,7 +11,7 @@
 static int recvraw(p_buffer buf, size_t wanted, luaL_Buffer *b);
 static int recvline(p_buffer buf, luaL_Buffer *b, size_t budget);
 static int recvall(p_buffer buf, luaL_Buffer *b, size_t budget);
-static int buffer_get(p_buffer buf, const char **data, size_t *count);
+static int buffer_get(p_buffer buf, const char **data, size_t *count, size_t wanted);
 static void buffer_skip(p_buffer buf, size_t count);
 static int sendraw(p_buffer buf, const char *data, size_t count, size_t *sent);
 
@@ -230,15 +230,14 @@ static int sendraw(p_buffer buf, const char *data, size_t count, size_t *sent) {
 static int recvraw(p_buffer buf, size_t wanted, luaL_Buffer *b) {
     int err = IO_DONE;
     size_t total = 0;
-    while (err == IO_DONE) {
+    do {
         size_t count; const char *data;
-        err = buffer_get(buf, &data, &count);
+        err = buffer_get(buf, &data, &count, wanted - total);
         count = MIN(count, wanted - total);
         luaL_addlstring(b, data, count);
         buffer_skip(buf, count);
         total += count;
-        if (total >= wanted) break;
-    }
+    } while (total < wanted && err == IO_DONE);
     return err;
 }
 
@@ -253,7 +252,7 @@ static int recvall(p_buffer buf, luaL_Buffer *b, size_t budget) {
     size_t total = 0;
     while (err == IO_DONE) {
         const char *data; size_t count;
-        err = buffer_get(buf, &data, &count);
+        err = buffer_get(buf, &data, &count, BUF_SIZE);
         if (budget && count > budget - total) {   /* strictly more than fits */
             count = budget - total;
             luaL_addlstring(b, data, count);
@@ -286,7 +285,7 @@ static int recvline(p_buffer buf, luaL_Buffer *b, size_t budget) {
     size_t total = 0;
     while (err == IO_DONE) {
         size_t count, pos; const char *data;
-        err = buffer_get(buf, &data, &count);
+        err = buffer_get(buf, &data, &count, BUF_SIZE);
         pos = 0;
         while (pos < count && data[pos] != '\n') {
             /* we ignore all \r's -- they are consumed but never counted */
@@ -325,15 +324,19 @@ static void buffer_skip(p_buffer buf, size_t count) {
 
 /*-------------------------------------------------------------------------*\
 * Return any data available in buffer, or get more data from transport layer
-* if buffer is empty
+* if buffer is empty. 'wanted' is how many more bytes the caller is still
+* after; when it is zero, the transport layer is still consulted (so an
+* already-closed connection is still reported), but no more than zero bytes
+* are requested from it, so a healthy connection with no data pending can
+* never block.
 \*-------------------------------------------------------------------------*/
-static int buffer_get(p_buffer buf, const char **data, size_t *count) {
+static int buffer_get(p_buffer buf, const char **data, size_t *count, size_t wanted) {
     int err = IO_DONE;
     p_io io = buf->io;
     p_timeout tm = buf->tm;
     if (buffer_isempty(buf)) {
         size_t got;
-        err = io->recv(io->ctx, buf->data, BUF_SIZE, &got, tm);
+        err = io->recv(io->ctx, buf->data, wanted == 0 ? 0 : BUF_SIZE, &got, tm);
         buf->first = 0;
         buf->last = got;
     }
diff --git a/src/usocket.c b/src/usocket.c
index 7965db6..e00c43d 100644
--- a/src/usocket.c
+++ b/src/usocket.c
@@ -258,6 +258,7 @@ int socket_recv(p_socket ps, char *data, size_t count, size_t *got, p_timeout tm
     int err;
     *got = 0;
     if (*ps == SOCKET_INVALID) return IO_CLOSED;
+    if (count == 0) return IO_DONE;
     for ( ;; ) {
         long taken = (long) recv(*ps, data, count, 0);
         if (taken > 0) {
@@ -343,6 +344,7 @@ int socket_read(p_socket ps, char *data, size_t count, size_t *got, p_timeout tm
     int err;
     *got = 0;
     if (*ps == SOCKET_INVALID) return IO_CLOSED;
+    if (count == 0) return IO_DONE;
     for ( ;; ) {
         long taken = (long) read(*ps, data, count);
         if (taken > 0) {
diff --git a/src/wsocket.c b/src/wsocket.c
index d3af9d4..86c6994 100644
--- a/src/wsocket.c
+++ b/src/wsocket.c
@@ -247,6 +247,7 @@ int socket_recv(p_socket ps, char *data, size_t count, size_t *got,
     int err, prev = IO_DONE;
     *got = 0;
     if (*ps == SOCKET_INVALID) return IO_CLOSED;
+    if (count == 0) return IO_DONE;
     for ( ;; ) {
         int taken = recv(*ps, data, (int) count, 0);
         if (taken > 0) {
diff --git a/test/test_receive_zero.lua b/test/test_receive_zero.lua
new file mode 100644
index 0000000..2259910
--- /dev/null
+++ b/test/test_receive_zero.lua
@@ -0,0 +1,36 @@
+-- a TCP receive(0) must never block: requesting zero bytes is trivially
+-- satisfied without touching the transport layer, POSIX recv(fd, buf, 0, 0)
+-- returns immediately regardless of whether data is available.
+local socket = require "socket"
+
+local host, port = "127.0.0.1", "5464"
+
+local server = assert(socket.bind(host, port))
+local client = assert(socket.connect(host, port))
+local peer = assert(server:accept())
+
+client:settimeout(2)
+
+-- no data has been sent by the peer: the read buffer is empty, so if
+-- receive(0) touches the network it will block until the timeout expires
+local t0 = socket.gettime()
+local data, err = client:receive(0)
+local elapsed = socket.gettime() - t0
+
+assert(data == "", "receive(0) on empty buffer returned " .. tostring(data))
+assert(err == nil, "receive(0) on empty buffer returned error " .. tostring(err))
+assert(elapsed < 1, "receive(0) on empty buffer blocked for " .. elapsed .. "s")
+
+-- receive(0) must also not consume any bytes when data *is* available
+assert(peer:send("hello"))
+data, err = client:receive(0)
+assert(data == "", "receive(0) with data pending returned " .. tostring(data))
+assert(err == nil, "receive(0) with data pending returned error " .. tostring(err))
+data = assert(client:receive(5))
+assert(data == "hello", "receive(0) consumed bytes meant for receive(5)")
+
+client:close()
+peer:close()
+server:close()
+
+print("done!")