bpo-46805: Add low level UDP socket functions to asyncio (GH-31455)

1 files changed, 128 insertions, 4 deletions

diff --git a/Modules/overlapped.c b/Modules/overlapped.c
index 2ba48c8..ab9a2f0 100644
--- a/Modules/overlapped.c
+++ b/Modules/overlapped.c
@@ -64,7 +64,7 @@ class _overlapped.Overlapped "OverlappedObject *" "&OverlappedType"
 enum {TYPE_NONE, TYPE_NOT_STARTED, TYPE_READ, TYPE_READINTO, TYPE_WRITE,
       TYPE_ACCEPT, TYPE_CONNECT, TYPE_DISCONNECT, TYPE_CONNECT_NAMED_PIPE,
       TYPE_WAIT_NAMED_PIPE_AND_CONNECT, TYPE_TRANSMIT_FILE, TYPE_READ_FROM,
-      TYPE_WRITE_TO};
+      TYPE_WRITE_TO, TYPE_READ_FROM_INTO};
 
 typedef struct {
     PyObject_HEAD
@@ -91,6 +91,17 @@ typedef struct {
             struct sockaddr_in6 address;
             int address_length;
         } read_from;
+
+        /* Data used for reading from a connectionless socket:
+           TYPE_READ_FROM_INTO */
+        struct {
+            // A (number of bytes read, (host, port)) tuple
+            PyObject* result;
+            /* Buffer passed by the user */
+            Py_buffer *user_buffer;
+            struct sockaddr_in6 address;
+            int address_length;
+        } read_from_into;
     };
 } OverlappedObject;
 
@@ -662,6 +673,13 @@ Overlapped_clear(OverlappedObject *self)
             }
             break;
         }
+        case TYPE_READ_FROM_INTO: {
+            if (self->read_from_into.result) {
+                // We've received a message, free the result tuple.
+                Py_CLEAR(self->read_from_into.result);
+            }
+            break;
+        }
         case TYPE_WRITE:
         case TYPE_WRITE_TO:
         case TYPE_READINTO: {
@@ -866,6 +884,11 @@ _overlapped_Overlapped_getresult_impl(OverlappedObject *self, BOOL wait)
             {
                 break;
             }
+            else if (self->type == TYPE_READ_FROM_INTO &&
+                     self->read_from_into.result != NULL)
+            {
+                break;
+            }
             /* fall through */
         default:
             return SetFromWindowsErr(err);
@@ -914,6 +937,30 @@ _overlapped_Overlapped_getresult_impl(OverlappedObject *self, BOOL wait)
 
             Py_INCREF(self->read_from.result);
             return self->read_from.result;
+        case TYPE_READ_FROM_INTO:
+            // unparse the address
+            addr = unparse_address((SOCKADDR*)&self->read_from_into.address,
+                self->read_from_into.address_length);
+
+            if (addr == NULL) {
+                return NULL;
+            }
+
+            // The result is a two item tuple: (number of bytes read, address)
+            self->read_from_into.result = PyTuple_New(2);
+            if (self->read_from_into.result == NULL) {
+                Py_CLEAR(addr);
+                return NULL;
+            }
+
+            // first item: number of bytes read
+            PyTuple_SET_ITEM(self->read_from_into.result, 0,
+                PyLong_FromUnsignedLong((unsigned long)transferred));
+            // second item: address
+            PyTuple_SET_ITEM(self->read_from_into.result, 1, addr);
+
+            Py_INCREF(self->read_from_into.result);
+            return self->read_from_into.result;
         default:
             return PyLong_FromUnsignedLong((unsigned long) transferred);
     }
@@ -1053,6 +1100,7 @@ do_WSARecv(OverlappedObject *self, HANDLE handle,
     }
 }
 
+
 /*[clinic input]
 _overlapped.Overlapped.WSARecv
 
@@ -1617,6 +1665,13 @@ Overlapped_traverse(OverlappedObject *self, visitproc visit, void *arg)
     case TYPE_READ_FROM:
         Py_VISIT(self->read_from.result);
         Py_VISIT(self->read_from.allocated_buffer);
+        break;
+    case TYPE_READ_FROM_INTO:
+        Py_VISIT(self->read_from_into.result);
+        if (self->read_from_into.user_buffer->obj) {
+            Py_VISIT(&self->read_from_into.user_buffer->obj);
+        }
+        break;
     }
     return 0;
 }
@@ -1766,8 +1821,8 @@ _overlapped_Overlapped_WSARecvFrom_impl(OverlappedObject *self,
                                         DWORD flags)
 /*[clinic end generated code: output=13832a2025b86860 input=1b2663fa130e0286]*/
 {
-    DWORD nread;
     PyObject *buf;
+    DWORD nread;
     WSABUF wsabuf;
     int ret;
     DWORD err;
@@ -1785,8 +1840,8 @@ _overlapped_Overlapped_WSARecvFrom_impl(OverlappedObject *self,
         return NULL;
     }
 
-    wsabuf.len = size;
     wsabuf.buf = PyBytes_AS_STRING(buf);
+    wsabuf.len = size;
 
     self->type = TYPE_READ_FROM;
     self->handle = handle;
@@ -1802,8 +1857,74 @@ _overlapped_Overlapped_WSARecvFrom_impl(OverlappedObject *self,
     Py_END_ALLOW_THREADS
 
     self->error = err = (ret < 0 ? WSAGetLastError() : ERROR_SUCCESS);
+    switch (err) {
+    case ERROR_BROKEN_PIPE:
+        mark_as_completed(&self->overlapped);
+        return SetFromWindowsErr(err);
+    case ERROR_SUCCESS:
+    case ERROR_MORE_DATA:
+    case ERROR_IO_PENDING:
+        Py_RETURN_NONE;
+    default:
+        self->type = TYPE_NOT_STARTED;
+        return SetFromWindowsErr(err);
+    }
+}
 
-    switch(err) {
+
+/*[clinic input]
+_overlapped.Overlapped.WSARecvFromInto
+
+    handle: HANDLE
+    buf as bufobj: Py_buffer
+    size: DWORD
+    flags: DWORD = 0
+    /
+
+Start overlapped receive.
+[clinic start generated code]*/
+
+static PyObject *
+_overlapped_Overlapped_WSARecvFromInto_impl(OverlappedObject *self,
+                                            HANDLE handle, Py_buffer *bufobj,
+                                            DWORD size, DWORD flags)
+/*[clinic end generated code: output=30c7ea171a691757 input=4be4b08d03531e76]*/
+{
+    DWORD nread;
+    WSABUF wsabuf;
+    int ret;
+    DWORD err;
+
+    if (self->type != TYPE_NONE) {
+        PyErr_SetString(PyExc_ValueError, "operation already attempted");
+        return NULL;
+    }
+
+#if SIZEOF_SIZE_T > SIZEOF_LONG
+    if (bufobj->len > (Py_ssize_t)ULONG_MAX) {
+        PyErr_SetString(PyExc_ValueError, "buffer too large");
+        return NULL;
+    }
+#endif
+
+    wsabuf.buf = bufobj->buf;
+    wsabuf.len = size;
+
+    self->type = TYPE_READ_FROM_INTO;
+    self->handle = handle;
+    self->read_from_into.user_buffer = bufobj;
+    memset(&self->read_from_into.address, 0, sizeof(self->read_from_into.address));
+    self->read_from_into.address_length = sizeof(self->read_from_into.address);
+
+    Py_BEGIN_ALLOW_THREADS
+    ret = WSARecvFrom((SOCKET)handle, &wsabuf, 1, &nread, &flags,
+                      (SOCKADDR*)&self->read_from_into.address,
+                      &self->read_from_into.address_length,
+                      &self->overlapped, NULL);
+    Py_END_ALLOW_THREADS
+
+    self->error = err = (ret < 0 ? WSAGetLastError() : ERROR_SUCCESS);
+    switch (err) {
     case ERROR_BROKEN_PIPE:
         mark_as_completed(&self->overlapped);
         return SetFromWindowsErr(err);
@@ -1817,6 +1938,7 @@ _overlapped_Overlapped_WSARecvFrom_impl(OverlappedObject *self,
     }
 }
 
+
 #include "clinic/overlapped.c.h"
 
 static PyMethodDef Overlapped_methods[] = {
@@ -1826,6 +1948,8 @@ static PyMethodDef Overlapped_methods[] = {
     _OVERLAPPED_OVERLAPPED_READFILEINTO_METHODDEF
     _OVERLAPPED_OVERLAPPED_WSARECV_METHODDEF
     _OVERLAPPED_OVERLAPPED_WSARECVINTO_METHODDEF
+    _OVERLAPPED_OVERLAPPED_WSARECVFROM_METHODDEF
+    _OVERLAPPED_OVERLAPPED_WSARECVFROMINTO_METHODDEF
     _OVERLAPPED_OVERLAPPED_WRITEFILE_METHODDEF
     _OVERLAPPED_OVERLAPPED_WSASEND_METHODDEF
     _OVERLAPPED_OVERLAPPED_ACCEPTEX_METHODDEF