udp: fix handling empty datagrams

saghul · Jul 3, 2014 · 64faa00 · 64faa00
1 parent 44d041e
commit 64faa00
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 3 deletions.
diff --git a/src/udp.c b/src/udp.c
@@ -24,14 +24,19 @@ on_udp_read(uv_udp_t* handle, int nread, const uv_buf_t* buf, struct sockaddr* a
     /* Object could go out of scope in the callback, increase refcount to avoid it */
     Py_INCREF(self);
 
-    if (nread == 0) {
+    if (nread == 0 && addr == NULL) {
+        /* libuv got EAGAIN and is returning the buffer to us */
         goto done;
     }
 
-    if (nread > 0) {
+    if (nread >= 0) {
         ASSERT(addr);
         address_tuple = makesockaddr(addr);
-        data = PyBytes_FromStringAndSize(buf->base, nread);
+        if (nread == 0) {
+            data = PyBytes_FromString("");
+        } else {
+            data = PyBytes_FromStringAndSize(buf->base, nread);
+        }
         py_errorno = Py_None;
         Py_INCREF(Py_None);
     } else {

diff --git a/tests/test_udp.py b/tests/test_udp.py
@@ -9,6 +9,7 @@
 TEST_PORT2 = 12346
 MULTICAST_ADDRESS = "239.255.0.1"
 
+
 class UDPTest(TestCase):
 
     def setUp(self):
@@ -58,6 +59,35 @@ def test_udp_pingpong(self):
         self.assertEqual(self.on_close_called, 3)
 
 
+class UDPEmptyDatagramTest(TestCase):
+
+    def setUp(self):
+        super(UDPEmptyDatagramTest, self).setUp()
+        self.server = None
+        self.client = None
+        self.on_close_called = 0
+
+    def on_close(self, handle):
+        self.on_close_called += 1
+
+    def on_client_recv(self, handle, ip_port, flags, data, error):
+        self.assertEqual(flags, 0)
+        ip, port = ip_port
+        self.assertEqual(error, None)
+        self.assertEqual(data, b"")
+        self.client.close(self.on_close)
+        self.server.close(self.on_close)
+
+    def test_udp_empty_datagram(self):
+        self.server = pyuv.UDP(self.loop)
+        self.client = pyuv.UDP(self.loop)
+        self.client.bind(("0.0.0.0", TEST_PORT2))
+        self.client.start_recv(self.on_client_recv)
+        self.server.send(("127.0.0.1", TEST_PORT2), b"")
+        self.loop.run()
+        self.assertEqual(self.on_close_called, 2)
+
+
 class UDPTestNull(TestCase):
 
     def setUp(self):