bpo-31819: Add AbstractEventLoop.sock_recv_into() (#4051)

* bpo-31819: Add AbstractEventLoop.sock_recv_into()

* Add NEWS

* Add doc

1 files changed, 156 insertions, 58 deletions

diff --git a/Modules/overlapped.c b/Modules/overlapped.c
index 6099d46..9355ce6 100644
--- a/Modules/overlapped.c
+++ b/Modules/overlapped.c
@@ -37,8 +37,8 @@
 
 #define T_HANDLE T_POINTER
 
-enum {TYPE_NONE, TYPE_NOT_STARTED, TYPE_READ, TYPE_WRITE, TYPE_ACCEPT,
-      TYPE_CONNECT, TYPE_DISCONNECT, TYPE_CONNECT_NAMED_PIPE,
+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};
 
 typedef struct {
@@ -51,10 +51,10 @@ typedef struct {
     /* Type of operation */
     DWORD type;
     union {
-        /* Buffer used for reading: TYPE_READ and TYPE_ACCEPT */
-        PyObject *read_buffer;
-        /* Buffer used for writing: TYPE_WRITE */
-        Py_buffer write_buffer;
+        /* Buffer allocated by us: TYPE_READ and TYPE_ACCEPT */
+        PyObject *allocated_buffer;
+        /* Buffer passed by the user: TYPE_WRITE and TYPE_READINTO */
+        Py_buffer user_buffer;
     };
 } OverlappedObject;
 
@@ -550,9 +550,9 @@ Overlapped_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     self->handle = NULL;
     self->error = 0;
     self->type = TYPE_NONE;
-    self->read_buffer = NULL;
+    self->allocated_buffer = NULL;
     memset(&self->overlapped, 0, sizeof(OVERLAPPED));
-    memset(&self->write_buffer, 0, sizeof(Py_buffer));
+    memset(&self->user_buffer, 0, sizeof(Py_buffer));
     if (event)
         self->overlapped.hEvent = event;
     return (PyObject *)self;
@@ -597,11 +597,12 @@ Overlapped_dealloc(OverlappedObject *self)
     switch (self->type) {
     case TYPE_READ:
     case TYPE_ACCEPT:
-        Py_CLEAR(self->read_buffer);
+        Py_CLEAR(self->allocated_buffer);
         break;
     case TYPE_WRITE:
-        if (self->write_buffer.obj)
-            PyBuffer_Release(&self->write_buffer);
+    case TYPE_READINTO:
+        if (self->user_buffer.obj)
+            PyBuffer_Release(&self->user_buffer);
         break;
     }
     PyObject_Del(self);
@@ -676,7 +677,7 @@ Overlapped_getresult(OverlappedObject *self, PyObject *args)
         case ERROR_MORE_DATA:
             break;
         case ERROR_BROKEN_PIPE:
-            if ((self->type == TYPE_READ || self->type == TYPE_ACCEPT) && self->read_buffer != NULL)
+            if (self->type == TYPE_READ || self->type == TYPE_READINTO)
                 break;
             /* fall through */
         default:
@@ -685,17 +686,45 @@ Overlapped_getresult(OverlappedObject *self, PyObject *args)
 
     switch (self->type) {
         case TYPE_READ:
-            assert(PyBytes_CheckExact(self->read_buffer));
-            if (transferred != PyBytes_GET_SIZE(self->read_buffer) &&
-                _PyBytes_Resize(&self->read_buffer, transferred))
+            assert(PyBytes_CheckExact(self->allocated_buffer));
+            if (transferred != PyBytes_GET_SIZE(self->allocated_buffer) &&
+                _PyBytes_Resize(&self->allocated_buffer, transferred))
                 return NULL;
-            Py_INCREF(self->read_buffer);
-            return self->read_buffer;
+            Py_INCREF(self->allocated_buffer);
+            return self->allocated_buffer;
         default:
             return PyLong_FromUnsignedLong((unsigned long) transferred);
     }
 }
 
+static PyObject *
+do_ReadFile(OverlappedObject *self, HANDLE handle,
+            char *bufstart, DWORD buflen)
+{
+    DWORD nread;
+    int ret;
+    DWORD err;
+
+    Py_BEGIN_ALLOW_THREADS
+    ret = ReadFile(handle, bufstart, buflen, &nread,
+                   &self->overlapped);
+    Py_END_ALLOW_THREADS
+
+    self->error = err = ret ? ERROR_SUCCESS : GetLastError();
+    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);
+    }
+}
+
 PyDoc_STRVAR(
     Overlapped_ReadFile_doc,
     "ReadFile(handle, size) -> Overlapped[message]\n\n"
@@ -706,10 +735,7 @@ Overlapped_ReadFile(OverlappedObject *self, PyObject *args)
 {
     HANDLE handle;
     DWORD size;
-    DWORD nread;
     PyObject *buf;
-    BOOL ret;
-    DWORD err;
 
     if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD, &handle, &size))
         return NULL;
@@ -728,14 +754,66 @@ Overlapped_ReadFile(OverlappedObject *self, PyObject *args)
 
     self->type = TYPE_READ;
     self->handle = handle;
-    self->read_buffer = buf;
+    self->allocated_buffer = buf;
+
+    return do_ReadFile(self, handle, PyBytes_AS_STRING(buf), size);
+}
+
+PyDoc_STRVAR(
+    Overlapped_ReadFileInto_doc,
+    "ReadFileInto(handle, buf) -> Overlapped[bytes_transferred]\n\n"
+    "Start overlapped receive");
+
+static PyObject *
+Overlapped_ReadFileInto(OverlappedObject *self, PyObject *args)
+{
+    HANDLE handle;
+    PyObject *bufobj;
+
+    if (!PyArg_ParseTuple(args, F_HANDLE "O", &handle, &bufobj))
+        return NULL;
+
+    if (self->type != TYPE_NONE) {
+        PyErr_SetString(PyExc_ValueError, "operation already attempted");
+        return NULL;
+    }
+
+    if (!PyArg_Parse(bufobj, "y*", &self->user_buffer))
+        return NULL;
+
+#if SIZEOF_SIZE_T > SIZEOF_LONG
+    if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) {
+        PyBuffer_Release(&self->user_buffer);
+        PyErr_SetString(PyExc_ValueError, "buffer too large");
+        return NULL;
+    }
+#endif
+
+    self->type = TYPE_READINTO;
+    self->handle = handle;
+
+    return do_ReadFile(self, handle, self->user_buffer.buf,
+                       (DWORD)self->user_buffer.len);
+}
+
+static PyObject *
+do_WSARecv(OverlappedObject *self, HANDLE handle,
+           char *bufstart, DWORD buflen, DWORD flags)
+{
+    DWORD nread;
+    WSABUF wsabuf;
+    int ret;
+    DWORD err;
+    
+    wsabuf.buf = bufstart;
+    wsabuf.len = buflen;
 
     Py_BEGIN_ALLOW_THREADS
-    ret = ReadFile(handle, PyBytes_AS_STRING(buf), size, &nread,
-                   &self->overlapped);
+    ret = WSARecv((SOCKET)handle, &wsabuf, 1, &nread, &flags,
+                  &self->overlapped, NULL);
     Py_END_ALLOW_THREADS
 
-    self->error = err = ret ? ERROR_SUCCESS : GetLastError();
+    self->error = err = (ret < 0 ? WSAGetLastError() : ERROR_SUCCESS);
     switch (err) {
         case ERROR_BROKEN_PIPE:
             mark_as_completed(&self->overlapped);
@@ -761,11 +839,7 @@ Overlapped_WSARecv(OverlappedObject *self, PyObject *args)
     HANDLE handle;
     DWORD size;
     DWORD flags = 0;
-    DWORD nread;
     PyObject *buf;
-    WSABUF wsabuf;
-    int ret;
-    DWORD err;
 
     if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD "|" F_DWORD,
                           &handle, &size, &flags))
@@ -785,28 +859,48 @@ Overlapped_WSARecv(OverlappedObject *self, PyObject *args)
 
     self->type = TYPE_READ;
     self->handle = handle;
-    self->read_buffer = buf;
-    wsabuf.len = size;
-    wsabuf.buf = PyBytes_AS_STRING(buf);
+    self->allocated_buffer = buf;
 
-    Py_BEGIN_ALLOW_THREADS
-    ret = WSARecv((SOCKET)handle, &wsabuf, 1, &nread, &flags,
-                  &self->overlapped, NULL);
-    Py_END_ALLOW_THREADS
+    return do_WSARecv(self, handle, PyBytes_AS_STRING(buf), size, flags);
+}
 
-    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);
+PyDoc_STRVAR(
+    Overlapped_WSARecvInto_doc,
+    "WSARecvInto(handle, buf, flags) -> Overlapped[bytes_transferred]\n\n"
+    "Start overlapped receive");
+
+static PyObject *
+Overlapped_WSARecvInto(OverlappedObject *self, PyObject *args)
+{
+    HANDLE handle;
+    PyObject *bufobj;
+    DWORD flags;
+
+    if (!PyArg_ParseTuple(args, F_HANDLE "O" F_DWORD,
+                          &handle, &bufobj, &flags))
+        return NULL;
+
+    if (self->type != TYPE_NONE) {
+        PyErr_SetString(PyExc_ValueError, "operation already attempted");
+        return NULL;
+    }
+
+    if (!PyArg_Parse(bufobj, "y*", &self->user_buffer))
+        return NULL;
+
+#if SIZEOF_SIZE_T > SIZEOF_LONG
+    if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) {
+        PyBuffer_Release(&self->user_buffer);
+        PyErr_SetString(PyExc_ValueError, "buffer too large");
+        return NULL;
     }
+#endif
+
+    self->type = TYPE_READINTO;
+    self->handle = handle;
+
+    return do_WSARecv(self, handle, self->user_buffer.buf,
+                      (DWORD)self->user_buffer.len, flags);
 }
 
 PyDoc_STRVAR(
@@ -831,12 +925,12 @@ Overlapped_WriteFile(OverlappedObject *self, PyObject *args)
         return NULL;
     }
 
-    if (!PyArg_Parse(bufobj, "y*", &self->write_buffer))
+    if (!PyArg_Parse(bufobj, "y*", &self->user_buffer))
         return NULL;
 
 #if SIZEOF_SIZE_T > SIZEOF_LONG
-    if (self->write_buffer.len > (Py_ssize_t)ULONG_MAX) {
-        PyBuffer_Release(&self->write_buffer);
+    if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) {
+        PyBuffer_Release(&self->user_buffer);
         PyErr_SetString(PyExc_ValueError, "buffer too large");
         return NULL;
     }
@@ -846,8 +940,8 @@ Overlapped_WriteFile(OverlappedObject *self, PyObject *args)
     self->handle = handle;
 
     Py_BEGIN_ALLOW_THREADS
-    ret = WriteFile(handle, self->write_buffer.buf,
-                    (DWORD)self->write_buffer.len,
+    ret = WriteFile(handle, self->user_buffer.buf,
+                    (DWORD)self->user_buffer.len,
                     &written, &self->overlapped);
     Py_END_ALLOW_THREADS
 
@@ -887,12 +981,12 @@ Overlapped_WSASend(OverlappedObject *self, PyObject *args)
         return NULL;
     }
 
-    if (!PyArg_Parse(bufobj, "y*", &self->write_buffer))
+    if (!PyArg_Parse(bufobj, "y*", &self->user_buffer))
         return NULL;
 
 #if SIZEOF_SIZE_T > SIZEOF_LONG
-    if (self->write_buffer.len > (Py_ssize_t)ULONG_MAX) {
-        PyBuffer_Release(&self->write_buffer);
+    if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) {
+        PyBuffer_Release(&self->user_buffer);
         PyErr_SetString(PyExc_ValueError, "buffer too large");
         return NULL;
     }
@@ -900,8 +994,8 @@ Overlapped_WSASend(OverlappedObject *self, PyObject *args)
 
     self->type = TYPE_WRITE;
     self->handle = handle;
-    wsabuf.len = (DWORD)self->write_buffer.len;
-    wsabuf.buf = self->write_buffer.buf;
+    wsabuf.len = (DWORD)self->user_buffer.len;
+    wsabuf.buf = self->user_buffer.buf;
 
     Py_BEGIN_ALLOW_THREADS
     ret = WSASend((SOCKET)handle, &wsabuf, 1, &written, flags,
@@ -951,7 +1045,7 @@ Overlapped_AcceptEx(OverlappedObject *self, PyObject *args)
 
     self->type = TYPE_ACCEPT;
     self->handle = (HANDLE)ListenSocket;
-    self->read_buffer = buf;
+    self->allocated_buffer = buf;
 
     Py_BEGIN_ALLOW_THREADS
     ret = Py_AcceptEx(ListenSocket, AcceptSocket, PyBytes_AS_STRING(buf),
@@ -1193,8 +1287,12 @@ static PyMethodDef Overlapped_methods[] = {
      METH_NOARGS, Overlapped_cancel_doc},
     {"ReadFile", (PyCFunction) Overlapped_ReadFile,
      METH_VARARGS, Overlapped_ReadFile_doc},
+    {"ReadFileInto", (PyCFunction) Overlapped_ReadFileInto,
+     METH_VARARGS, Overlapped_ReadFileInto_doc},
     {"WSARecv", (PyCFunction) Overlapped_WSARecv,
      METH_VARARGS, Overlapped_WSARecv_doc},
+    {"WSARecvInto", (PyCFunction) Overlapped_WSARecvInto,
+     METH_VARARGS, Overlapped_WSARecvInto_doc},
     {"WriteFile", (PyCFunction) Overlapped_WriteFile,
      METH_VARARGS, Overlapped_WriteFile_doc},
     {"WSASend", (PyCFunction) Overlapped_WSASend,