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

This commit is contained in:
Alex Grönholm 2022-03-13 18:42:29 +02:00 committed by GitHub
parent 7e473e94a5
commit 9f04ee569c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 489 additions and 7 deletions

View file

@ -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,7 @@ _overlapped_Overlapped_WSARecvFrom_impl(OverlappedObject *self,
Py_END_ALLOW_THREADS
self->error = err = (ret < 0 ? WSAGetLastError() : ERROR_SUCCESS);
switch(err) {
switch (err) {
case ERROR_BROKEN_PIPE:
mark_as_completed(&self->overlapped);
return SetFromWindowsErr(err);
@ -1817,6 +1871,74 @@ _overlapped_Overlapped_WSARecvFrom_impl(OverlappedObject *self,
}
}
/*[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);
case ERROR_SUCCESS:
case ERROR_MORE_DATA:
case ERROR_IO_PENDING:
Py_RETURN_NONE;
default:
self->type = TYPE_NOT_STARTED;
return SetFromWindowsErr(err);
}
}
#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